refactor(research): extract shared UCT kernel - #190
Conversation
📝 WalkthroughWalkthroughA new ChangesUCT Kernel Extraction
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant FormulaMcts
participant hft_search_kernel
participant NodeTree
FormulaMcts->>hft_search_kernel: select_expandable(NodeTree, root, exploration)
hft_search_kernel->>NodeTree: validate tree and calculate UCT scores
hft_search_kernel-->>FormulaMcts: return expandable node
FormulaMcts->>hft_search_kernel: backpropagate(NodeTree, root, leaf, reward)
hft_search_kernel->>NodeTree: update lineage statistics
Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 22475efeb0
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| None => break, | ||
| } | ||
| } | ||
| let _ = backpropagate(&mut self.nodes, 0, node_id, evaluation.score); |
There was a problem hiding this comment.
When a restored, otherwise accepted checkpoint has a node at u64::MAX visits or a finite reward total near overflow, backpropagate returns StatsOverflow without mutating the tree. Discarding that error after removing the pending candidate lets AutoResearchKernel checkpoint and persist the evaluation as successful even though its reward was permanently omitted; before this extraction, the resulting invalid statistics caused checkpointing to fail. Preserve the failure so the run remains fail-closed instead of silently corrupting search state.
AGENTS.md reference: AGENTS.md:L59-L61
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
rust_hft/Cargo.toml (1)
85-85: 🧹 Nitpick | 🔵 TrivialWorkspace graph changed — run the required metadata check.
Adding a new member changes the workspace dependency graph. As per path instructions, "after workspace graph changes, run
cargo metadata --locked --no-deps" to confirm the graph resolves cleanly before merge.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust_hft/Cargo.toml` at line 85, After adding the research-core/search-kernel workspace member, run cargo metadata --locked --no-deps from the workspace root and confirm it completes successfully. Address any metadata or lockfile resolution errors before proceeding.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust_hft/alpha-harness/engine/src/engines/mcts.rs`:
- Around line 303-311: Update MCTS::observe so failures returned by
backpropagate are not silently discarded: capture the Result and emit an
appropriate error or warning with the candidate/node context when it fails,
while preserving the existing successful backpropagation path and
non-finite-score handling.
In `@rust_hft/research-core/search-kernel/src/lib.rs`:
- Around line 72-80: Update UctStats::validate to require best_reward to be
present exactly when visits is greater than zero, independently of total_reward;
retain the zero-visits requirement that total_reward is zero and best_reward is
absent. Also prevent direct Deserialize construction of UctStats, routing
deserialization through the validated from_parts path so all instances enforce
these invariants.
---
Nitpick comments:
In `@rust_hft/Cargo.toml`:
- Line 85: After adding the research-core/search-kernel workspace member, run
cargo metadata --locked --no-deps from the workspace root and confirm it
completes successfully. Address any metadata or lockfile resolution errors
before proceeding.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 15cec4e9-8f7d-4507-a348-b82c04b27f10
⛔ Files ignored due to path filters (1)
rust_hft/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (7)
rust_hft/ARCHITECTURE.mdrust_hft/Cargo.tomlrust_hft/alpha-harness/engine/Cargo.tomlrust_hft/alpha-harness/engine/src/engines/mcts.rsrust_hft/alpha-harness/engine/src/engines/mod.rsrust_hft/research-core/search-kernel/Cargo.tomlrust_hft/research-core/search-kernel/src/lib.rs
| fn observe(&mut self, proposal: &EngineProposal, evaluation: &CandidateEvaluation) { | ||
| let Some(mut node_id) = self.candidates.remove(&proposal.candidate_id) else { | ||
| let Some(node_id) = self.candidates.remove(&proposal.candidate_id) else { | ||
| return; | ||
| }; | ||
| if !evaluation.score.is_finite() { | ||
| return; | ||
| } | ||
| loop { | ||
| let node = &mut self.nodes[node_id]; | ||
| node.visits += 1; | ||
| node.total_reward += evaluation.score; | ||
| node.best_reward = Some( | ||
| node.best_reward | ||
| .map_or(evaluation.score, |best| best.max(evaluation.score)), | ||
| ); | ||
| match node.parent { | ||
| Some(parent) => node_id = parent, | ||
| None => break, | ||
| } | ||
| } | ||
| let _ = backpropagate(&mut self.nodes, 0, node_id, evaluation.score); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
observe() silently drops backpropagate failures.
let _ = backpropagate(...) discards any UctError (e.g. StatsOverflow, InvalidTopology). This is fail-closed for data integrity (no partial mutation), but a genuine observation is lost with zero observability — nothing signals that a proposed candidate's evaluation never got recorded.
🔍 Suggested minimal observability fix
- let _ = backpropagate(&mut self.nodes, 0, node_id, evaluation.score);
+ if let Err(error) = backpropagate(&mut self.nodes, 0, node_id, evaluation.score) {
+ tracing::warn!(%error, node_id, "MCTS observe: backpropagate rejected evaluation");
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| fn observe(&mut self, proposal: &EngineProposal, evaluation: &CandidateEvaluation) { | |
| let Some(mut node_id) = self.candidates.remove(&proposal.candidate_id) else { | |
| let Some(node_id) = self.candidates.remove(&proposal.candidate_id) else { | |
| return; | |
| }; | |
| if !evaluation.score.is_finite() { | |
| return; | |
| } | |
| loop { | |
| let node = &mut self.nodes[node_id]; | |
| node.visits += 1; | |
| node.total_reward += evaluation.score; | |
| node.best_reward = Some( | |
| node.best_reward | |
| .map_or(evaluation.score, |best| best.max(evaluation.score)), | |
| ); | |
| match node.parent { | |
| Some(parent) => node_id = parent, | |
| None => break, | |
| } | |
| } | |
| let _ = backpropagate(&mut self.nodes, 0, node_id, evaluation.score); | |
| } | |
| fn observe(&mut self, proposal: &EngineProposal, evaluation: &CandidateEvaluation) { | |
| let Some(node_id) = self.candidates.remove(&proposal.candidate_id) else { | |
| return; | |
| }; | |
| if !evaluation.score.is_finite() { | |
| return; | |
| } | |
| if let Err(error) = backpropagate(&mut self.nodes, 0, node_id, evaluation.score) { | |
| tracing::warn!(%error, node_id, "MCTS observe: backpropagate rejected evaluation"); | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust_hft/alpha-harness/engine/src/engines/mcts.rs` around lines 303 - 311,
Update MCTS::observe so failures returned by backpropagate are not silently
discarded: capture the Result and emit an appropriate error or warning with the
candidate/node context when it fails, while preserving the existing successful
backpropagation path and non-finite-score handling.
| fn validate(self) -> Result<(), UctError> { | ||
| if !self.total_reward.is_finite() | ||
| || self.best_reward.is_some_and(|reward| !reward.is_finite()) | ||
| || (self.visits == 0) != (self.total_reward == 0.0 && self.best_reward.is_none()) | ||
| { | ||
| return Err(UctError::InvalidStats); | ||
| } | ||
| Ok(()) | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
UctStats::validate() doesn't fully enforce the visits↔best_reward invariant.
The check (self.visits == 0) != (self.total_reward == 0.0 && self.best_reward.is_none()) only catches inconsistency when total_reward == 0.0. A node with visits > 0, a nonzero total_reward, and best_reward == None passes validation undetected — and that exact shape is trivially reachable via a deserialized checkpoint (best_reward: null with nonzero visits/total_reward). Since validate_tree (used by both checkpoint validation and backpropagate) calls node.stats() → UctStats::from_parts → this validate(), such a corrupted node is accepted rather than rejected, undermining the PR's fail-closed checkpoint acceptance criterion; the bogus best_reward then silently surfaces through trace()/MctsNodeSnapshot.
Separately, UctStats derives Deserialize directly (line 28) rather than only being constructible via from_parts; any future direct deserialization path would bypass this invariant entirely.
🛡️ Proposed fix for the invariant check
fn validate(self) -> Result<(), UctError> {
if !self.total_reward.is_finite()
|| self.best_reward.is_some_and(|reward| !reward.is_finite())
- || (self.visits == 0) != (self.total_reward == 0.0 && self.best_reward.is_none())
+ || (self.visits == 0) != self.best_reward.is_none()
+ || (self.visits == 0 && self.total_reward != 0.0)
{
return Err(UctError::InvalidStats);
}
Ok(())
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| fn validate(self) -> Result<(), UctError> { | |
| if !self.total_reward.is_finite() | |
| || self.best_reward.is_some_and(|reward| !reward.is_finite()) | |
| || (self.visits == 0) != (self.total_reward == 0.0 && self.best_reward.is_none()) | |
| { | |
| return Err(UctError::InvalidStats); | |
| } | |
| Ok(()) | |
| } | |
| fn validate(self) -> Result<(), UctError> { | |
| if !self.total_reward.is_finite() | |
| || self.best_reward.is_some_and(|reward| !reward.is_finite()) | |
| || (self.visits == 0) != self.best_reward.is_none() | |
| || (self.visits == 0 && self.total_reward != 0.0) | |
| { | |
| return Err(UctError::InvalidStats); | |
| } | |
| Ok(()) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust_hft/research-core/search-kernel/src/lib.rs` around lines 72 - 80, Update
UctStats::validate to require best_reward to be present exactly when visits is
greater than zero, independently of total_reward; retain the zero-visits
requirement that total_reward is zero and best_reward is absent. Also prevent
direct Deserialize construction of UctStats, routing deserialization through the
validated from_parts path so all instances enforce these invariants.
Change contract
Extract deterministic UCT tree mechanics from continuous Formula MCTS into a domain-neutral research-core kernel while preserving Formula candidate generation, seeded selection, checkpoint wire format, and evaluation behavior.
Out of scope
Polymarket candidate/evaluator logic, LLM proposal policy, collector or snapshot changes, cloud execution, result publication, and checkpoint deletion.
Dependency or merge order
Parent PRD: #184. This PR implements #185 and is the stack base for #186 (Polymarket adapter). Merge order is #185 -> #186 -> #187 -> #188 -> #189. This layer compiles and remains fail-closed on its own.
Focused validation
cargo test --locked -p hft-search-kernel(5 passed)cargo clippy --locked -p hft-search-kernel --all-targets -- -D warningscargo test --locked -p alpha-engine engines::mcts::tests --lib(8 passed)cargo check --locked -p alpha-engine --no-default-features --libgit diff --checkRollout / rollback impact
No deployment or schema cutover. Reverting this PR restores Formula-local UCT mechanics. The existing checkpoint JSON shape and version remain unchanged.
Closes #185
Summary by CodeRabbit
New Features
Documentation
Bug Fixes