From 684dfe56fb36ea655fe0329fa048d9fabbd992d2 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 17:35:56 -0500 Subject: [PATCH 1/9] feat(brain): add graph state sync endpoint and new mapping models - Introduced a new `/state` endpoint in `main.py` to return the synchronized graph state for use by the Rust Engine. - Includes metadata for asset end timestamps, implications, contradictions, and partition mappings. - Added `BrainStatePayload`, `ImplicationMapping`, `PartitionMapping`, and `ContradictionMapping` models in `models.py` to structure graph state data. - Integrated logic to fetch: - Conditional token mappings and asset end timestamps based on active markets. - Implications and contradictions using relationship graph data. - Partition mappings for mutually exclusive market groups, with deterministic condition IDs based on asset hashes. - Enhanced consistency and compatibility between the Brain Service and the Rust Engine by providing a complete state snapshot in a structured payload. This addition facilitates efficient polling and synchronization with the Rust Engine while maintaining flexibility for future extensions. --- services/brain/main.py | 97 ++++++++++++++++++++++++++++++++++++++++ services/brain/models.py | 30 +++++++++++++ 2 files changed, 127 insertions(+) diff --git a/services/brain/main.py b/services/brain/main.py index 482b4bf..c2b26f9 100644 --- a/services/brain/main.py +++ b/services/brain/main.py @@ -25,8 +25,12 @@ from models import ( AnalyzePairRequest, AnalyzePairResponse, + BrainStatePayload, GraphStats, HealthResponse, + ImplicationMapping, + PartitionMapping, + ContradictionMapping, Market, MarketRead, Relationship, @@ -478,3 +482,96 @@ async def get_relationships_for_market( async def graph_stats(): """Return high-level statistics about the in-memory relationship graph.""" return await _graph_manager.get_stats() + + +# ── Graph Sync Endpoint ──────────────────────────────────────────────────────── + +@app.get("/state", response_model=BrainStatePayload) +async def get_brain_state(session: AsyncSession = Depends(get_session)): + """Return the entire synchronized graph state for the Rust Engine to poll.""" + + # 1. Fetch metadata (tokens and timestamps) + result = await session.execute( + select(Market.condition_id, Market.end_date, Market.clob_token_ids) + .where(Market.active.is_(True), Market.closed.is_(False)) + ) + + asset_end_timestamps = {} + condition_to_yes_token = {} + + for condition_id, end_date, clob_token_ids in result.all(): + if not clob_token_ids or len(clob_token_ids) == 0: + continue + + # The YES token is almost always the first element in the array for binary markets + yes_token = clob_token_ids[0] + condition_to_yes_token[condition_id] = yes_token + + if end_date: + try: + # End dates from Polymarket are usually ISO8601 strings + dt = datetime.fromisoformat(end_date.replace("Z", "+00:00")) + asset_end_timestamps[yes_token] = int(dt.timestamp()) + except Exception: + pass + + # 2. Implications and Contradictions + implications_dicts = await _graph_manager.get_all_relationships() + implications = [] + contradictions = [] + + for rel in implications_dicts: + parent_token = condition_to_yes_token.get(rel["parent_condition_id"]) + child_token = condition_to_yes_token.get(rel["child_condition_id"]) + + if not parent_token or not child_token: + continue + + if rel["logic_type"] == "IMPLIES": + implications.append( + ImplicationMapping( + parent_asset_id=parent_token, + child_asset_id=child_token, + confidence=rel["confidence"], + ) + ) + elif rel["logic_type"] == "MUTUALLY_EXCLUSIVE": + contradictions.append( + ContradictionMapping( + asset_a=parent_token, + asset_b=child_token, + confidence=rel["confidence"], + ) + ) + + # 3. Partitions + partitions_lists = await _graph_manager.get_partition_groups() + partitions = [] + + import hashlib + for i, group in enumerate(partitions_lists): + if len(group) >= 2: + group_tokens = [] + for cid in group: + tok = condition_to_yes_token.get(cid) + if tok: + group_tokens.append(tok) + + if len(group_tokens) >= 2: + # Hash the sorted assets to get a deterministic B256-like condition ID + hash_id = hashlib.sha256("".join(sorted(group_tokens)).encode()).hexdigest() + partitions.append( + PartitionMapping( + condition_id=f"0x{hash_id}", + expected_outcomes_count=len(group_tokens), + assets=group_tokens, + confidence=1.0 # Inherited confidence + ) + ) + + return BrainStatePayload( + implications=implications, + partitions=partitions, + contradictions=contradictions, + asset_end_timestamps=asset_end_timestamps + ) diff --git a/services/brain/models.py b/services/brain/models.py index 6a05e17..f0c6939 100644 --- a/services/brain/models.py +++ b/services/brain/models.py @@ -138,3 +138,33 @@ class HealthResponse(BaseModel): status: str db_connected: bool llm_configured: bool + + +class ImplicationMapping(BaseModel): + """An implication edge between two markets.""" + parent_asset_id: str + child_asset_id: str + confidence: float + + +class PartitionMapping(BaseModel): + """A group of mutually exclusive markets representing options to a question.""" + condition_id: str + expected_outcomes_count: int + assets: list[str] + confidence: float + + +class ContradictionMapping(BaseModel): + """A pair of mutually exclusive markets.""" + asset_a: str + asset_b: str + confidence: float + + +class BrainStatePayload(BaseModel): + """The aggregate graph state polled by the Rust Engine.""" + implications: list[ImplicationMapping] + partitions: list[PartitionMapping] + contradictions: list[ContradictionMapping] + asset_end_timestamps: dict[str, int] From 6bff53d2ad8af6e414c03c146e7a5c3b4da45ca0 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 17:36:18 -0500 Subject: [PATCH 2/9] feat(engine): integrate contradiction evaluation into EngineActor - Added `ContradictionMapping` to `BrainStatePayload` for state synchronization. - Defined `ContradictionMapping` struct and implemented parsing logic (`parse_contradictions`). - Extended `EngineActor` with `contradiction_edges` for storing mutually exclusive asset relationships. - Implemented `evaluate_contradictions` in `EngineActor` to identify and execute profitable contradiction opportunities. - Logic includes validation of weights, gas fees, and profit margins. - Updated book handlers to trigger contradiction evaluations alongside implications and partitions. - Adjusted asset frequency tracking and expiry validation logic to account for contradiction mapping (`compute_tracked_assets`). - Added comprehensive unit test for contradiction evaluation flow (`test_engine_actor_triggers_contradiction_arb`). - Updated documentation in `strategies/mod.rs` to mark `ContradictionStrategy` as integrated. This implementation connects the Tier-3 contradiction strategy with the runtime, enabling evaluation and execution of mutually exclusive market inconsistencies. --- crates/engine/src/actor.rs | 228 +++++++++++++++++++++++++++- crates/engine/src/strategies/mod.rs | 4 +- 2 files changed, 227 insertions(+), 5 deletions(-) diff --git a/crates/engine/src/actor.rs b/crates/engine/src/actor.rs index 36727dc..ddb97ac 100644 --- a/crates/engine/src/actor.rs +++ b/crates/engine/src/actor.rs @@ -19,6 +19,7 @@ use std::str::FromStr; pub struct BrainStatePayload { pub implications: Vec, pub partitions: Vec, + pub contradictions: Vec, pub asset_end_timestamps: HashMap, } @@ -37,6 +38,13 @@ pub struct PartitionMapping { pub confidence: f64, } +#[derive(Debug, Deserialize)] +pub struct ContradictionMapping { + pub asset_a: String, + pub asset_b: String, + pub confidence: f64, +} + pub struct EngineActor { // State Caches pub orderbook_cache: HashMap, @@ -45,6 +53,9 @@ pub struct EngineActor { // Map of relation edges: Child -> List of Parents that imply it pub implication_reverse_edges: HashMap>, + // Map of undirected contradicting mutually-exclusive peers (A implies NOT B, etc.) + pub contradiction_edges: HashMap>, // Peer -> (Peer ID, Default Weight 1.0) + // Map of Market (Condition) -> Set of Asset IDs (Outcomes) pub market_outcomes: HashMap>, @@ -81,6 +92,7 @@ impl EngineActor { orderbook_cache: HashMap::new(), implication_edges: HashMap::new(), implication_reverse_edges: HashMap::new(), + contradiction_edges: HashMap::new(), market_outcomes: HashMap::new(), market_expected_outcome_counts: HashMap::new(), market_end_timestamps: HashMap::new(), @@ -176,6 +188,8 @@ impl EngineActor { .await; self.evaluate_partitions(&book, bot_tx, fee_rate, gas_per_leg_usdc) .await; + self.evaluate_contradictions(&book, bot_tx, fee_rate, gas_per_leg_usdc) + .await; } async fn evaluate_implications( @@ -590,6 +604,98 @@ impl EngineActor { } } + async fn evaluate_contradictions( + &mut self, + book: &NormalizedOrderbook, + bot_tx: &Sender, + fee_rate: Decimal, + gas_per_leg_usdc: Decimal, + ) { + let condition_id = book.asset_id; + + let peers = match self.contradiction_edges.get(&condition_id) { + Some(p) => p.clone(), + None => return, + }; + + for (peer_id, weight) in peers { + let peer_book = match self.orderbook_cache.get(&peer_id) { + Some(b) => b, + None => continue, + }; + + if book.bid_untradeable || peer_book.bid_untradeable { + continue; + } + + let ask_a = book.vwap_ask; + let ask_b = peer_book.vwap_ask; + + if ask_a <= Decimal::ZERO || ask_b <= Decimal::ZERO { + continue; + } + + // A Contradiction strategy also executes exactly 2 legs (Sell A, Sell B) + let total_est_gas_usdc = gas_per_leg_usdc * Decimal::from(2); + + // Weight logic: for binary markets, usually 1.0 vs 1.0. The tier 2 brain assigns 1.0 for mutually exclusive currently. + let weight_a = Decimal::ONE; + let weight_b = weight; + + if let Some(profit) = crate::strategies::contradiction::ContradictionStrategy::check( + condition_id, + ask_a, + weight_a, + peer_id, + ask_b, + weight_b, + fee_rate, + total_est_gas_usdc, + ) { + // Sizing follows partition strategy style: max(ask_a, ask_b) controls risk limit + let highest_price = ask_a.max(ask_b); + let consistent_size_shares = + arbos_core::constants::TARGET_LIQUIDITY / highest_price; + + let signal = ArbSignal { + legs: vec![ + TradeAction::Sell { + asset_id: condition_id, + size: consistent_size_shares, + }, + TradeAction::Sell { + asset_id: peer_id, + size: consistent_size_shares, + }, + ], + expected_profit_usdc: profit, + timestamp: book.timestamp, + }; + + info!( + asset_a = %condition_id, + asset_b = %peer_id, + profit = %profit, + "Routing Contradiction ArbSignal to Bot" + ); + if let Err(e) = bot_tx.send(signal).await { + error!( + "Failed to route Contradiction ArbSignal to Bot channel: {}", + e + ); + } + + // Eject executed legs + self.orderbook_cache.remove(&condition_id); + self.orderbook_cache.remove(&peer_id); + + // Break after executing one contradiction for this asset to avoid double-spend + // on the same loop tick if it contradicts multiple peers + break; + } + } + } + /// The Max_Duration filter. Markets resolving $>30$ days out trap capital if the Oracle takes weeks. fn violates_max_duration(&self, asset_id: U256) -> bool { if let Some(end_ts) = self.market_end_timestamps.get(&asset_id) { @@ -621,11 +727,16 @@ impl EngineActor { let (new_imp_edges, new_imp_rev) = Self::parse_implications(&payload.implications); let (new_market_outcomes, new_market_expected) = Self::parse_partitions(&payload.partitions); + let new_contradictions = Self::parse_contradictions(&payload.contradictions); let new_timestamps = Self::parse_timestamps(&payload.asset_end_timestamps); // 2. Compute Tracked Assets (with Expiry GC and MAX_TRACKED limits) - let new_tracked_assets = - Self::compute_tracked_assets(&new_imp_edges, &new_market_outcomes, &new_timestamps); + let new_tracked_assets = Self::compute_tracked_assets( + &new_imp_edges, + &new_contradictions, + &new_market_outcomes, + &new_timestamps, + ); // 3. Diff & Command Ingestor self.apply_asset_diffs(&new_tracked_assets, cmd_tx).await; @@ -634,6 +745,7 @@ impl EngineActor { self.tracked_assets = new_tracked_assets; self.implication_edges = new_imp_edges; self.implication_reverse_edges = new_imp_rev; + self.contradiction_edges = new_contradictions; self.market_outcomes = new_market_outcomes; self.market_expected_outcome_counts = new_market_expected; self.market_end_timestamps = new_timestamps; @@ -701,6 +813,35 @@ impl EngineActor { (edges, rev_edges) } + fn parse_contradictions( + contradictions: &[crate::actor::ContradictionMapping], + ) -> HashMap> { + let mut edges: HashMap> = HashMap::new(); + + for cont in contradictions { + if cont.confidence < MIN_CONFIDENCE_THRESHOLD { + continue; + } + match ( + U256::from_str_radix(&cont.asset_a, 10), + U256::from_str_radix(&cont.asset_b, 10), + ) { + (Ok(a_id), Ok(b_id)) => { + let weight = Decimal::ONE; // Using 1.0 for mutually exclusive edges by default + edges.entry(a_id).or_default().push((b_id, weight)); + edges.entry(b_id).or_default().push((a_id, weight)); + } + _ => { + warn!( + "Failed to parse Contradiction U256 IDs, silently dropping: A: {} | B: {}", + cont.asset_a, cont.asset_b + ); + } + } + } + edges + } + fn parse_partitions( partitions: &[crate::actor::PartitionMapping], ) -> ( @@ -757,6 +898,7 @@ impl EngineActor { fn compute_tracked_assets( edges: &HashMap>, + contradictions: &HashMap>, outcomes: &HashMap>, timestamps: &HashMap, ) -> HashSet { @@ -768,6 +910,12 @@ impl EngineActor { let mut frequency: HashMap = HashMap::new(); Self::tally_edge_frequencies(&mut frequency, edges, timestamps, current_ts); + Self::tally_contradiction_frequencies( + &mut frequency, + contradictions, + timestamps, + current_ts, + ); Self::tally_outcome_frequencies(&mut frequency, outcomes, timestamps, current_ts); Self::sort_and_limit_tracked_assets(frequency) @@ -799,6 +947,24 @@ impl EngineActor { } } + fn tally_contradiction_frequencies( + frequency: &mut HashMap, + edges: &HashMap>, + timestamps: &HashMap, + current_ts: i64, + ) { + for (&p, children) in edges { + if Self::is_asset_unexpired(p, timestamps, current_ts) { + *frequency.entry(p).or_insert(0) += 1; + } + for &(c, _) in children { + if Self::is_asset_unexpired(c, timestamps, current_ts) { + *frequency.entry(c).or_insert(0) += 1; + } + } + } + } + fn tally_outcome_frequencies( frequency: &mut HashMap, outcomes: &HashMap>, @@ -981,4 +1147,62 @@ mod tests { assert!(signal.expected_profit_usdc > dec!(0.0)); assert_eq!(signal.legs.len(), 2); } + + #[tokio::test] + async fn test_engine_actor_triggers_contradiction_arb() { + let mut actor = EngineActor::new(); + let (tx, mut rx) = mpsc::channel(10); + + let peer_a_id = U256::from(10_u64); + let peer_b_id = U256::from(20_u64); + + actor + .contradiction_edges + .insert(peer_a_id, vec![(peer_b_id, Decimal::ONE)]); + actor + .contradiction_edges + .insert(peer_b_id, vec![(peer_a_id, Decimal::ONE)]); + + let peer_a_book = NormalizedOrderbook { + asset_id: peer_a_id, + market: B256::default(), + vwap_bid: dec!(0.55), + vwap_ask: dec!(0.60), // High ask, P = 0.60 + bid_untradeable: false, + ask_untradeable: false, + timestamp: 123456, + }; + + let peer_b_book = NormalizedOrderbook { + asset_id: peer_b_id, + market: B256::default(), + vwap_bid: dec!(0.55), + vwap_ask: dec!(0.60), // High ask, P = 0.60 (Sum P = 1.20) + bid_untradeable: false, + ask_untradeable: false, + timestamp: 123457, + }; + + let gas_per_leg_usdc = dec!(0.05); + + // Inject first book, then second book + actor + .handle_book_update(peer_a_book.clone(), &tx, dec!(0.0), gas_per_leg_usdc) + .await; + actor + .handle_book_update(peer_b_book.clone(), &tx, dec!(0.01), gas_per_leg_usdc) + .await; + + // Check cache clearing + assert!(!actor.orderbook_cache.contains_key(&peer_a_id)); + assert!(!actor.orderbook_cache.contains_key(&peer_b_id)); + + // Wait for signal + let signal = rx.recv().await.unwrap(); + + // High ask prices indicate sum > 1.0 (Contradiction). + // Profit should be strictly positive. + assert!(signal.expected_profit_usdc > dec!(0.0)); + assert_eq!(signal.legs.len(), 2); + } } diff --git a/crates/engine/src/strategies/mod.rs b/crates/engine/src/strategies/mod.rs index 33b940c..2190f4b 100644 --- a/crates/engine/src/strategies/mod.rs +++ b/crates/engine/src/strategies/mod.rs @@ -1,7 +1,5 @@ // Note: ContradictionStrategy is implemented as a Tier-3 strategy -// but is not currently wired into the Engine runtime (EngineActor). -// It will be wired once the mapping/source for mutually-exclusive pairs -// is provided by the Tier-2 Brain API. +// and is wired into the Engine runtime (EngineActor). pub mod contradiction; pub mod implication; pub mod partition; From 98119c8f624b4ab9d7ba1b3d1507191442c608f8 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 17:52:13 -0500 Subject: [PATCH 3/9] fix(engine): resolve asset_id misusage and adjust bid logic in contradiction evaluation - Renamed `condition_id` to `asset_id` across the contradiction evaluation flow in `EngineActor` to ensure consistent terminology and logic. - Fixed incorrect usage of `vwap_ask` for bid-related calculations; replaced with `vwap_bid` to align with strategy logic. - Updated sizing calculation to use `sum_bids` instead of the highest price, ensuring a more precise allocation of liquidity. - Adjusted asset frequency computation to add directional constraints, ensuring order and correctness when processing parent-child pairs. - Updated test order book parameters to match bid/ask logic corrections and reflect valid cases for evaluation. These changes improve the accuracy and readability of the contradiction evaluation and support the intended strategy behavior while addressing inconsistencies in asset handling and computational logic. --- crates/engine/src/actor.rs | 49 +++++++++++++++++++------------------- 1 file changed, 25 insertions(+), 24 deletions(-) diff --git a/crates/engine/src/actor.rs b/crates/engine/src/actor.rs index ddb97ac..7f588e2 100644 --- a/crates/engine/src/actor.rs +++ b/crates/engine/src/actor.rs @@ -611,9 +611,9 @@ impl EngineActor { fee_rate: Decimal, gas_per_leg_usdc: Decimal, ) { - let condition_id = book.asset_id; + let asset_id = book.asset_id; - let peers = match self.contradiction_edges.get(&condition_id) { + let peers = match self.contradiction_edges.get(&asset_id) { Some(p) => p.clone(), None => return, }; @@ -628,10 +628,10 @@ impl EngineActor { continue; } - let ask_a = book.vwap_ask; - let ask_b = peer_book.vwap_ask; + let bid_a = book.vwap_bid; + let bid_b = peer_book.vwap_bid; - if ask_a <= Decimal::ZERO || ask_b <= Decimal::ZERO { + if bid_a <= Decimal::ZERO || bid_b <= Decimal::ZERO { continue; } @@ -643,24 +643,23 @@ impl EngineActor { let weight_b = weight; if let Some(profit) = crate::strategies::contradiction::ContradictionStrategy::check( - condition_id, - ask_a, + asset_id, + bid_a, weight_a, peer_id, - ask_b, + bid_b, weight_b, fee_rate, total_est_gas_usdc, ) { - // Sizing follows partition strategy style: max(ask_a, ask_b) controls risk limit - let highest_price = ask_a.max(ask_b); - let consistent_size_shares = - arbos_core::constants::TARGET_LIQUIDITY / highest_price; + // Sizing follows logic from strategy: TARGET_LIQUIDITY / sum_bids + let sum_bids = bid_a + bid_b; + let consistent_size_shares = arbos_core::constants::TARGET_LIQUIDITY / sum_bids; let signal = ArbSignal { legs: vec![ TradeAction::Sell { - asset_id: condition_id, + asset_id, size: consistent_size_shares, }, TradeAction::Sell { @@ -673,7 +672,7 @@ impl EngineActor { }; info!( - asset_a = %condition_id, + asset_a = %asset_id, asset_b = %peer_id, profit = %profit, "Routing Contradiction ArbSignal to Bot" @@ -686,7 +685,7 @@ impl EngineActor { } // Eject executed legs - self.orderbook_cache.remove(&condition_id); + self.orderbook_cache.remove(&asset_id); self.orderbook_cache.remove(&peer_id); // Break after executing one contradiction for this asset to avoid double-spend @@ -954,12 +953,14 @@ impl EngineActor { current_ts: i64, ) { for (&p, children) in edges { - if Self::is_asset_unexpired(p, timestamps, current_ts) { - *frequency.entry(p).or_insert(0) += 1; - } for &(c, _) in children { - if Self::is_asset_unexpired(c, timestamps, current_ts) { - *frequency.entry(c).or_insert(0) += 1; + if p < c { + if Self::is_asset_unexpired(p, timestamps, current_ts) { + *frequency.entry(p).or_insert(0) += 1; + } + if Self::is_asset_unexpired(c, timestamps, current_ts) { + *frequency.entry(c).or_insert(0) += 1; + } } } } @@ -1166,8 +1167,8 @@ mod tests { let peer_a_book = NormalizedOrderbook { asset_id: peer_a_id, market: B256::default(), - vwap_bid: dec!(0.55), - vwap_ask: dec!(0.60), // High ask, P = 0.60 + vwap_bid: dec!(0.60), // High bid, P = 0.60 + vwap_ask: dec!(0.65), bid_untradeable: false, ask_untradeable: false, timestamp: 123456, @@ -1176,8 +1177,8 @@ mod tests { let peer_b_book = NormalizedOrderbook { asset_id: peer_b_id, market: B256::default(), - vwap_bid: dec!(0.55), - vwap_ask: dec!(0.60), // High ask, P = 0.60 (Sum P = 1.20) + vwap_bid: dec!(0.60), // High bid, P = 0.60 (Sum P = 1.20) + vwap_ask: dec!(0.65), bid_untradeable: false, ask_untradeable: false, timestamp: 123457, From b4bd59bf37ab4356d81334806d66d99e2a3da66a Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 17:52:36 -0500 Subject: [PATCH 4/9] fix(brain): improve logging for end_date parsing and introduce confidence-based partitioning - Enhanced error handling in `_graph_manager` by adding detailed warnings during end_date parsing failures, including `condition_id`, `end_date`, and error details. - Introduced confidence-based logic in partition group mapping: - Built a lookup table for mutually exclusive relationship confidences. - Integrated minimum confidence calculation for grouped partitions. - Updated deterministic hashing to include a `v1_` prefix for future migrations. - Improved confidence tracking for partitions, applying calculated values instead of a hardcoded `1.0`. These changes increase the clarity of logged issues, improve partition grouping accuracy, and lay the foundation for future versioned hashing enhancements. --- services/brain/main.py | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/services/brain/main.py b/services/brain/main.py index c2b26f9..4f6763c 100644 --- a/services/brain/main.py +++ b/services/brain/main.py @@ -512,8 +512,9 @@ async def get_brain_state(session: AsyncSession = Depends(get_session)): # End dates from Polymarket are usually ISO8601 strings dt = datetime.fromisoformat(end_date.replace("Z", "+00:00")) asset_end_timestamps[yes_token] = int(dt.timestamp()) - except Exception: - pass + except Exception as e: + (logger.bind(condition_id=condition_id, end_date=end_date, error=str(e)) + .warning("failed_to_parse_end_date")) # 2. Implications and Contradictions implications_dicts = await _graph_manager.get_all_relationships() @@ -548,6 +549,14 @@ async def get_brain_state(session: AsyncSession = Depends(get_session)): partitions_lists = await _graph_manager.get_partition_groups() partitions = [] + # Build a lookup for mutual exclusion confidences + me_confidences = {} + for rel in implications_dicts: + if rel["logic_type"] == "MUTUALLY_EXCLUSIVE": + c1, c2 = rel["parent_condition_id"], rel["child_condition_id"] + me_confidences[f"{c1}|{c2}"] = rel["confidence"] + me_confidences[f"{c2}|{c1}"] = rel["confidence"] + import hashlib for i, group in enumerate(partitions_lists): if len(group) >= 2: @@ -558,14 +567,25 @@ async def get_brain_state(session: AsyncSession = Depends(get_session)): group_tokens.append(tok) if len(group_tokens) >= 2: + # Calculate minimum confidence among mutually exclusive edges within the group + group_confs = [] + for j in range(len(group)): + for k in range(j + 1, len(group)): + conf = me_confidences.get(f"{group[j]}|{group[k]}") + if conf is not None: + group_confs.append(conf) + + partition_confidence = min(group_confs) if group_confs else 1.0 + # Hash the sorted assets to get a deterministic B256-like condition ID - hash_id = hashlib.sha256("".join(sorted(group_tokens)).encode()).hexdigest() + hash_input = "v1_" + "_".join(sorted(group_tokens)) + hash_id = hashlib.sha256(hash_input.encode()).hexdigest() partitions.append( PartitionMapping( condition_id=f"0x{hash_id}", - expected_outcomes_count=len(group_tokens), + expected_outcomes_count=len(group), assets=group_tokens, - confidence=1.0 # Inherited confidence + confidence=partition_confidence ) ) From 5b05a08559a8f1713d68d7d1b5af724fd17908d7 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 18:09:00 -0500 Subject: [PATCH 5/9] feat(engine): introduce asset-to-partition mapping and refactor partition handling - Added `asset_to_partition` mapping in `EngineActor` to enable O(1) partition lookups for assets. - Refactored `market_outcomes` and `market_expected_outcome_counts` to use `String` keys instead of `B256` for consistent data handling across the actor. - Updated partition-related methods (`get_partition_bids`, `is_partition_exhaustively_tradable`, `check_and_send_partition`) to adopt the new `asset_to_partition` logic. - Enhanced `parse_partitions` to generate `asset_to_partition` mapping alongside market outcomes and expected outcomes. - Deferred order book ejection during contradiction evaluation for correct processing of peer execution. - Improved code clarity and reduced redundancy in partition and contradiction evaluation logic. These changes improve the engine's partition-handling performance, simplify the data structure interface, and ensure more robust execution in complex asset scenarios. --- crates/engine/src/actor.rs | 124 ++++++++++++++++++++----------------- services/brain/main.py | 4 +- 2 files changed, 70 insertions(+), 58 deletions(-) diff --git a/crates/engine/src/actor.rs b/crates/engine/src/actor.rs index 7f588e2..79190b1 100644 --- a/crates/engine/src/actor.rs +++ b/crates/engine/src/actor.rs @@ -13,7 +13,6 @@ use arbos_core::constants::{ }; use serde::Deserialize; -use std::str::FromStr; #[derive(Debug, Deserialize)] pub struct BrainStatePayload { @@ -57,10 +56,13 @@ pub struct EngineActor { pub contradiction_edges: HashMap>, // Peer -> (Peer ID, Default Weight 1.0) // Map of Market (Condition) -> Set of Asset IDs (Outcomes) - pub market_outcomes: HashMap>, + pub market_outcomes: HashMap>, + + // Map of Asset ID -> Partition ID (for O(1) partition lookups) + pub asset_to_partition: HashMap, // Map of Market (Condition) -> Required Number of Outcomes for Exhaustiveness - pub market_expected_outcome_counts: HashMap, + pub market_expected_outcome_counts: HashMap, // Map of Asset ID -> Resolution Deadline (Timestamp in seconds) pub market_end_timestamps: HashMap, @@ -75,6 +77,12 @@ pub struct EngineActor { pub api_client: reqwest::Client, } +pub type PartitionParseResult = ( + HashMap>, + HashMap, + HashMap, +); + impl Default for EngineActor { fn default() -> Self { Self::new() @@ -94,6 +102,7 @@ impl EngineActor { implication_reverse_edges: HashMap::new(), contradiction_edges: HashMap::new(), market_outcomes: HashMap::new(), + asset_to_partition: HashMap::new(), market_expected_outcome_counts: HashMap::new(), market_end_timestamps: HashMap::new(), tracked_assets: HashSet::new(), @@ -493,19 +502,24 @@ impl EngineActor { fee_rate: Decimal, gas_per_leg_usdc: Decimal, ) { - let condition_id = book.market; + let asset_id = book.asset_id; + + let partition_id = match self.asset_to_partition.get(&asset_id) { + Some(pid) => pid.clone(), + None => return, + }; - let bids = match self.get_partition_bids(&condition_id) { + let bids = match self.get_partition_bids(&partition_id) { Some(b) => b, None => return, }; - if !self.is_partition_exhaustively_tradable(&condition_id, &bids) { + if !self.is_partition_exhaustively_tradable(&partition_id, &bids) { return; } self.check_and_send_partition( - condition_id, + partition_id, bids, book.timestamp, bot_tx, @@ -515,11 +529,8 @@ impl EngineActor { .await; } - fn get_partition_bids( - &self, - condition_id: &polymarket_client_sdk::types::B256, - ) -> Option> { - let outcome_assets = self.market_outcomes.get(condition_id)?; + fn get_partition_bids(&self, partition_id: &String) -> Option> { + let outcome_assets = self.market_outcomes.get(partition_id)?; let mut bids = Vec::new(); for asset in outcome_assets { @@ -539,13 +550,13 @@ impl EngineActor { fn is_partition_exhaustively_tradable( &self, - condition_id: &polymarket_client_sdk::types::B256, + partition_id: &String, bids: &[(U256, Decimal)], ) -> bool { // Enforce strict exhaustiveness: Are we receiving the EXACT number of outcomes for this partition? let required_outcome_count = self .market_expected_outcome_counts - .get(condition_id) + .get(partition_id) .copied() .unwrap_or(0); @@ -559,7 +570,7 @@ impl EngineActor { async fn check_and_send_partition( &mut self, - condition_id: polymarket_client_sdk::types::B256, + partition_id: String, bids: Vec<(U256, Decimal)>, timestamp: i64, bot_tx: &Sender, @@ -592,7 +603,7 @@ impl EngineActor { timestamp, }; - info!(market = ?condition_id, leg_count = bids.len(), profit = %profit, "Routing Partition ArbSignal to Bot"); + info!(partition_id = %partition_id, leg_count = bids.len(), profit = %profit, "Routing Partition ArbSignal to Bot"); if let Err(e) = bot_tx.send(signal).await { error!("Failed to route Partition ArbSignal to Bot channel: {}", e); } @@ -613,12 +624,19 @@ impl EngineActor { ) { let asset_id = book.asset_id; - let peers = match self.contradiction_edges.get(&asset_id) { - Some(p) => p.clone(), + let peers_len = match self.contradiction_edges.get(&asset_id) { + Some(p) => p.len(), None => return, }; - for (peer_id, weight) in peers { + let mut executed_peer = None; + + for i in 0..peers_len { + let (peer_id, weight) = { + let peers = self.contradiction_edges.get(&asset_id).unwrap(); + peers[i] + }; + let peer_book = match self.orderbook_cache.get(&peer_id) { Some(b) => b, None => continue, @@ -684,15 +702,19 @@ impl EngineActor { ); } - // Eject executed legs - self.orderbook_cache.remove(&asset_id); - self.orderbook_cache.remove(&peer_id); + // Defer ejection + executed_peer = Some(peer_id); // Break after executing one contradiction for this asset to avoid double-spend // on the same loop tick if it contradicts multiple peers break; } } + + if let Some(peer_id) = executed_peer { + self.orderbook_cache.remove(&asset_id); + self.orderbook_cache.remove(&peer_id); + } } /// The Max_Duration filter. Markets resolving $>30$ days out trap capital if the Oracle takes weeks. @@ -724,7 +746,7 @@ impl EngineActor { // 1. Parse Graph Edges let (new_imp_edges, new_imp_rev) = Self::parse_implications(&payload.implications); - let (new_market_outcomes, new_market_expected) = + let (new_market_outcomes, new_asset_to_partition, new_market_expected) = Self::parse_partitions(&payload.partitions); let new_contradictions = Self::parse_contradictions(&payload.contradictions); let new_timestamps = Self::parse_timestamps(&payload.asset_end_timestamps); @@ -746,6 +768,7 @@ impl EngineActor { self.implication_reverse_edges = new_imp_rev; self.contradiction_edges = new_contradictions; self.market_outcomes = new_market_outcomes; + self.asset_to_partition = new_asset_to_partition; self.market_expected_outcome_counts = new_market_expected; self.market_end_timestamps = new_timestamps; @@ -841,48 +864,37 @@ impl EngineActor { edges } - fn parse_partitions( - partitions: &[crate::actor::PartitionMapping], - ) -> ( - HashMap>, - HashMap, - ) { + fn parse_partitions(partitions: &[crate::actor::PartitionMapping]) -> PartitionParseResult { let mut outcomes = HashMap::new(); + let mut asset_to_partition = HashMap::new(); let mut expected = HashMap::new(); for part in partitions { if part.confidence < MIN_CONFIDENCE_THRESHOLD { continue; } - match polymarket_client_sdk::types::B256::from_str(&part.condition_id) { - Ok(cond_id) => { - expected.insert(cond_id, part.expected_outcomes_count); - - let mut asset_set = HashSet::new(); - for a_str in &part.assets { - match U256::from_str_radix(a_str, 10) { - Ok(a_id) => { - asset_set.insert(a_id); - } - Err(_) => { - warn!( - "Failed to parse Partition Asset U256 ID, dropping asset: {}", - a_str - ); - } - } + + let partition_id = part.condition_id.clone(); + expected.insert(partition_id.clone(), part.expected_outcomes_count); + + let mut asset_set = HashSet::new(); + for a_str in &part.assets { + match U256::from_str_radix(a_str, 10) { + Ok(a_id) => { + asset_set.insert(a_id); + asset_to_partition.insert(a_id, partition_id.clone()); + } + Err(_) => { + warn!( + "Failed to parse Partition Asset U256 ID, dropping asset: {}", + a_str + ); } - outcomes.insert(cond_id, asset_set); - } - Err(_) => { - warn!( - "Failed to parse Partition Condition B256 ID, dropping partition: {}", - part.condition_id - ); } } + outcomes.insert(partition_id, asset_set); } - (outcomes, expected) + (outcomes, asset_to_partition, expected) } fn parse_timestamps(asset_end_timestamps: &HashMap) -> HashMap { @@ -898,7 +910,7 @@ impl EngineActor { fn compute_tracked_assets( edges: &HashMap>, contradictions: &HashMap>, - outcomes: &HashMap>, + outcomes: &HashMap>, timestamps: &HashMap, ) -> HashSet { let current_ts = std::time::SystemTime::now() @@ -968,7 +980,7 @@ impl EngineActor { fn tally_outcome_frequencies( frequency: &mut HashMap, - outcomes: &HashMap>, + outcomes: &HashMap>, timestamps: &HashMap, current_ts: i64, ) { diff --git a/services/brain/main.py b/services/brain/main.py index 4f6763c..7fc5298 100644 --- a/services/brain/main.py +++ b/services/brain/main.py @@ -7,11 +7,12 @@ from __future__ import annotations +import hashlib import sys from contextlib import asynccontextmanager from datetime import UTC, datetime - from typing import Any + from fastapi import Depends, FastAPI, HTTPException, Query, Security from fastapi.security import APIKeyHeader from loguru import logger @@ -557,7 +558,6 @@ async def get_brain_state(session: AsyncSession = Depends(get_session)): me_confidences[f"{c1}|{c2}"] = rel["confidence"] me_confidences[f"{c2}|{c1}"] = rel["confidence"] - import hashlib for i, group in enumerate(partitions_lists): if len(group) >= 2: group_tokens = [] From 96251b07eb8799ba1dd5f194daf959c2e7848418 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 18:20:30 -0500 Subject: [PATCH 6/9] feat(engine/brain): enhance `/state` endpoint with authentication and improve partition bid handling - **Engine Updates:** - Updated `get_partition_bids` and `is_partition_exhaustively_tradable` methods to use `&str` instead of `String` for `partition_id` parameters, enhancing performance with borrowed references. - Added `BRAIN_API_KEY` retrieval and included it as a header in API client requests, securing communication with the Brain service. - **Brain Service Enhancements:** - Secured the `/state` endpoint with admin-level authentication via `Depends(_verify_admin)`. - Added environment variable check for `ADMIN_API_KEY` in test setups. - Expanded testing for `/state`: - Verified missing authentication response (`401 Unauthorized`). - Verified proper response with valid API key and mocked graph data. - **Model Enhancements:** - Introduced and tested new Brain state models: - `ImplicationMapping`, `PartitionMapping`, `ContradictionMapping`, and `BrainStatePayload`. - Confirmed compatibility and correctness of structured state payload serialization. These changes improve the Rust engine's partition handling and ensure secure, robust collaboration between the Brain service and the Rust Engine. --- crates/engine/src/actor.rs | 7 +++-- services/brain/main.py | 2 +- services/brain/tests/test_main.py | 47 +++++++++++++++++++++++++++++ services/brain/tests/test_models.py | 28 +++++++++++++++++ 4 files changed, 81 insertions(+), 3 deletions(-) diff --git a/crates/engine/src/actor.rs b/crates/engine/src/actor.rs index 79190b1..8572e58 100644 --- a/crates/engine/src/actor.rs +++ b/crates/engine/src/actor.rs @@ -529,7 +529,7 @@ impl EngineActor { .await; } - fn get_partition_bids(&self, partition_id: &String) -> Option> { + fn get_partition_bids(&self, partition_id: &str) -> Option> { let outcome_assets = self.market_outcomes.get(partition_id)?; let mut bids = Vec::new(); @@ -550,7 +550,7 @@ impl EngineActor { fn is_partition_exhaustively_tradable( &self, - partition_id: &String, + partition_id: &str, bids: &[(U256, Decimal)], ) -> bool { // Enforce strict exhaustiveness: Are we receiving the EXACT number of outcomes for this partition? @@ -790,9 +790,12 @@ impl EngineActor { return Err(anyhow::anyhow!("BRAIN_API_URL must have a valid host")); } + let brain_api_key = std::env::var("BRAIN_API_KEY").unwrap_or_else(|_| "".to_string()); + let res = self .api_client .get(valid_brain_url.join("state")?) + .header("X-API-Key", brain_api_key) .send() .await?; diff --git a/services/brain/main.py b/services/brain/main.py index 7fc5298..ac3d4e7 100644 --- a/services/brain/main.py +++ b/services/brain/main.py @@ -487,7 +487,7 @@ async def graph_stats(): # ── Graph Sync Endpoint ──────────────────────────────────────────────────────── -@app.get("/state", response_model=BrainStatePayload) +@app.get("/state", response_model=BrainStatePayload, dependencies=[Depends(_verify_admin)]) async def get_brain_state(session: AsyncSession = Depends(get_session)): """Return the entire synchronized graph state for the Rust Engine to poll.""" diff --git a/services/brain/tests/test_main.py b/services/brain/tests/test_main.py index 1f8d918..b52033c 100644 --- a/services/brain/tests/test_main.py +++ b/services/brain/tests/test_main.py @@ -16,6 +16,7 @@ def _mock_env(monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://postgres:password@localhost:5432/arbos") monkeypatch.setenv("WATSONX_APIKEY", "test-key") monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project") + monkeypatch.setenv("ADMIN_API_KEY", "test-admin-key") # Clear cached settings from config import get_settings @@ -82,3 +83,49 @@ async def test_returns_stats(self, client): assert "total_markets" in data assert "total_relationships" in data assert data["total_markets"] == 0 + + +class TestBrainStateEndpoint: + async def test_requires_auth_for_state(self, client): + response = await client.get("/state") + assert response.status_code == 401 + + async def test_returns_state(self, client): + from main import app, get_session + from unittest.mock import AsyncMock + + mock_session = AsyncMock() + mock_result = MagicMock() + mock_result.all.return_value = [ + ("cond_a", "2024-12-31T00:00:00Z", ["yes_a", "no_a"]), + ("cond_b", "2024-12-31T00:00:00Z", ["yes_b", "no_b"]) + ] + mock_session.execute.return_value = mock_result + + async def override_get_session_for_state(): + yield mock_session + + app.dependency_overrides[get_session] = override_get_session_for_state + + with patch("main._graph_manager") as mock_gm: + mock_gm.get_all_relationships = AsyncMock(return_value=[ + { + "parent_condition_id": "cond_a", + "child_condition_id": "cond_b", + "logic_type": "IMPLIES", + "confidence": 0.9, + } + ]) + mock_gm.get_partition_groups = AsyncMock(return_value=[]) + + response = await client.get("/state", headers={"X-API-Key": "test-admin-key"}) + assert response.status_code == 200 + data = response.json() + + assert "implications" in data + assert len(data["implications"]) == 1 + assert data["implications"][0]["parent_asset_id"] == "yes_a" + assert data["implications"][0]["child_asset_id"] == "yes_b" + assert "yes_a" in data["asset_end_timestamps"] + + app.dependency_overrides.pop(get_session, None) diff --git a/services/brain/tests/test_models.py b/services/brain/tests/test_models.py index 2e70fa6..4550007 100644 --- a/services/brain/tests/test_models.py +++ b/services/brain/tests/test_models.py @@ -13,6 +13,10 @@ MarketRead, RelationshipOutput, RelationshipRead, + ImplicationMapping, + PartitionMapping, + ContradictionMapping, + BrainStatePayload, ) @@ -144,3 +148,27 @@ def test_basic(self): condition_id_b="0xdef", ) assert req.condition_id_a == "0xabc" + + +class TestBrainStateModels: + def test_implication_mapping(self): + m = ImplicationMapping(parent_asset_id="A", child_asset_id="B", confidence=0.9) + assert m.parent_asset_id == "A" + + def test_partition_mapping(self): + p = PartitionMapping(condition_id="C", expected_outcomes_count=2, + assets=["A", "B"], confidence=1.0) + assert p.expected_outcomes_count == 2 + + def test_contradiction_mapping(self): + c = ContradictionMapping(asset_a="A", asset_b="B", confidence=0.8) + assert c.asset_b == "B" + + def test_brain_state_payload(self): + payload = BrainStatePayload( + implications=[], + partitions=[], + contradictions=[], + asset_end_timestamps={"A": 1234567890} + ) + assert payload.asset_end_timestamps["A"] == 1234567890 From 39b44230628bdc8381ac742fe0d1212ad699b9b9 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 18:30:01 -0500 Subject: [PATCH 7/9] fix(engine): validate ADMIN_API_KEY during brain state sync - Updated `brain_api_key` retrieval to require `ADMIN_API_KEY` instead of `BRAIN_API_KEY`. - Added error handling to ensure `ADMIN_API_KEY` is present and properly configured to avoid runtime issues during state synchronization. This change improves security and robustness for engine-to-brain communication. --- crates/engine/src/actor.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/engine/src/actor.rs b/crates/engine/src/actor.rs index 8572e58..1aada32 100644 --- a/crates/engine/src/actor.rs +++ b/crates/engine/src/actor.rs @@ -790,7 +790,9 @@ impl EngineActor { return Err(anyhow::anyhow!("BRAIN_API_URL must have a valid host")); } - let brain_api_key = std::env::var("BRAIN_API_KEY").unwrap_or_else(|_| "".to_string()); + let brain_api_key = std::env::var("ADMIN_API_KEY").map_err(|_| { + anyhow::anyhow!("ADMIN_API_KEY environment variable is required to sync brain state") + })?; let res = self .api_client From cbd66320543dd17741a75350c099ad4e8397f1a2 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 18:40:56 -0500 Subject: [PATCH 8/9] refactor(engine): rename market-related fields to partition for consistency - Renamed `market_outcomes` to `partition_outcomes` for alignment with partition-centric terminology. - Updated `market_expected_outcome_counts` to `partition_expected_outcome_counts`, ensuring consistent field naming across the codebase. - Refactored relevant methods and logic to match the updated terminology (`get_partition_bids`, `parse_partitions`, etc.). - Improved environment variable handling for `ADMIN_API_KEY` during API requests to avoid errors when the variable is missing or empty. - Updated the Brain service to use `len(group_tokens)` for outcome count calculation in `PartitionMapping`, ensuring correctness. This refactor standardizes field names, improves clarity, and ensures accuracy in partition-based handling within the Engine and related services. --- crates/engine/src/actor.rs | 41 +++++++++++++++++++------------------- services/brain/main.py | 2 +- 2 files changed, 21 insertions(+), 22 deletions(-) diff --git a/crates/engine/src/actor.rs b/crates/engine/src/actor.rs index 1aada32..ae306dd 100644 --- a/crates/engine/src/actor.rs +++ b/crates/engine/src/actor.rs @@ -55,14 +55,14 @@ pub struct EngineActor { // Map of undirected contradicting mutually-exclusive peers (A implies NOT B, etc.) pub contradiction_edges: HashMap>, // Peer -> (Peer ID, Default Weight 1.0) - // Map of Market (Condition) -> Set of Asset IDs (Outcomes) - pub market_outcomes: HashMap>, + // Map of Partition ID -> Set of Asset IDs (Outcomes) + pub partition_outcomes: HashMap>, // Map of Asset ID -> Partition ID (for O(1) partition lookups) pub asset_to_partition: HashMap, - // Map of Market (Condition) -> Required Number of Outcomes for Exhaustiveness - pub market_expected_outcome_counts: HashMap, + // Map of Partition ID -> Required Number of Outcomes for Exhaustiveness + pub partition_expected_outcome_counts: HashMap, // Map of Asset ID -> Resolution Deadline (Timestamp in seconds) pub market_end_timestamps: HashMap, @@ -101,9 +101,9 @@ impl EngineActor { implication_edges: HashMap::new(), implication_reverse_edges: HashMap::new(), contradiction_edges: HashMap::new(), - market_outcomes: HashMap::new(), + partition_outcomes: HashMap::new(), asset_to_partition: HashMap::new(), - market_expected_outcome_counts: HashMap::new(), + partition_expected_outcome_counts: HashMap::new(), market_end_timestamps: HashMap::new(), tracked_assets: HashSet::new(), current_gas_usdc_per_leg: (ESTIMATED_GAS_UNITS_PER_LEG * ESTIMATED_GWEI_PRICE) @@ -530,7 +530,7 @@ impl EngineActor { } fn get_partition_bids(&self, partition_id: &str) -> Option> { - let outcome_assets = self.market_outcomes.get(partition_id)?; + let outcome_assets = self.partition_outcomes.get(partition_id)?; let mut bids = Vec::new(); for asset in outcome_assets { @@ -555,7 +555,7 @@ impl EngineActor { ) -> bool { // Enforce strict exhaustiveness: Are we receiving the EXACT number of outcomes for this partition? let required_outcome_count = self - .market_expected_outcome_counts + .partition_expected_outcome_counts .get(partition_id) .copied() .unwrap_or(0); @@ -746,7 +746,7 @@ impl EngineActor { // 1. Parse Graph Edges let (new_imp_edges, new_imp_rev) = Self::parse_implications(&payload.implications); - let (new_market_outcomes, new_asset_to_partition, new_market_expected) = + let (new_partition_outcomes, new_asset_to_partition, new_partition_expected) = Self::parse_partitions(&payload.partitions); let new_contradictions = Self::parse_contradictions(&payload.contradictions); let new_timestamps = Self::parse_timestamps(&payload.asset_end_timestamps); @@ -755,7 +755,7 @@ impl EngineActor { let new_tracked_assets = Self::compute_tracked_assets( &new_imp_edges, &new_contradictions, - &new_market_outcomes, + &new_partition_outcomes, &new_timestamps, ); @@ -767,9 +767,9 @@ impl EngineActor { self.implication_edges = new_imp_edges; self.implication_reverse_edges = new_imp_rev; self.contradiction_edges = new_contradictions; - self.market_outcomes = new_market_outcomes; + self.partition_outcomes = new_partition_outcomes; self.asset_to_partition = new_asset_to_partition; - self.market_expected_outcome_counts = new_market_expected; + self.partition_expected_outcome_counts = new_partition_expected; self.market_end_timestamps = new_timestamps; debug!("Successfully synchronized in-memory Graph from Tier 2 Brain API."); @@ -790,16 +790,15 @@ impl EngineActor { return Err(anyhow::anyhow!("BRAIN_API_URL must have a valid host")); } - let brain_api_key = std::env::var("ADMIN_API_KEY").map_err(|_| { - anyhow::anyhow!("ADMIN_API_KEY environment variable is required to sync brain state") - })?; + let mut req = self.api_client.get(valid_brain_url.join("state")?); - let res = self - .api_client - .get(valid_brain_url.join("state")?) - .header("X-API-Key", brain_api_key) - .send() - .await?; + if let Ok(brain_api_key) = std::env::var("ADMIN_API_KEY") + && !brain_api_key.is_empty() + { + req = req.header("X-API-Key", brain_api_key); + } + + let res = req.send().await?; if !res.status().is_success() { return Err(anyhow::anyhow!( diff --git a/services/brain/main.py b/services/brain/main.py index ac3d4e7..4a04614 100644 --- a/services/brain/main.py +++ b/services/brain/main.py @@ -583,7 +583,7 @@ async def get_brain_state(session: AsyncSession = Depends(get_session)): partitions.append( PartitionMapping( condition_id=f"0x{hash_id}", - expected_outcomes_count=len(group), + expected_outcomes_count=len(group_tokens), assets=group_tokens, confidence=partition_confidence ) From 2822cd81e2cf11e66ca66282eec3b7bf107edfd9 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 18:52:41 -0500 Subject: [PATCH 9/9] feat(brain): optimize `/state` endpoint with conditional querying and refined partition processing - **Performance Enhancements:** - Added condition-based filtering to database queries, ensuring only relevant `condition_ids` are processed to reduce data overhead. - Optimized implication and partition logic by pre-collecting `condition_ids` to streamline metadata queries. - **Partition Handling Improvements:** - Excluded partitions with incomplete token mappings or invalid configurations (e.g., non-cliques). - Refined confidence calculations to consider minimum confidence across relationships while ensuring clique validation. - **Logging Enhancements:** - Added detailed debug-level logs for skipped relationships and partitions, highlighting reasons such as missing tokens or invalid configurations. These updates enhance the `/state` endpoint with better performance, more reliable partition processing, and improved visibility into processing behaviors. --- services/brain/main.py | 53 ++++++++++++++++++++++++++++++++++++------ 1 file changed, 46 insertions(+), 7 deletions(-) diff --git a/services/brain/main.py b/services/brain/main.py index 4a04614..c1c428c 100644 --- a/services/brain/main.py +++ b/services/brain/main.py @@ -491,10 +491,32 @@ async def graph_stats(): async def get_brain_state(session: AsyncSession = Depends(get_session)): """Return the entire synchronized graph state for the Rust Engine to poll.""" - # 1. Fetch metadata (tokens and timestamps) + # 1. Fetch relationships and partition groups + implications_dicts = await _graph_manager.get_all_relationships() + partitions_lists = await _graph_manager.get_partition_groups() + + # Collect all needed condition IDs to optimize the DB query + needed_conditions = set() + for rel in implications_dicts: + needed_conditions.add(rel["parent_condition_id"]) + needed_conditions.add(rel["child_condition_id"]) + for group in partitions_lists: + for cid in group: + needed_conditions.add(cid) + + if not needed_conditions: + return BrainStatePayload( + implications=[], partitions=[], contradictions=[], asset_end_timestamps={} + ) + + # 2. Fetch metadata (tokens and timestamps) result = await session.execute( select(Market.condition_id, Market.end_date, Market.clob_token_ids) - .where(Market.active.is_(True), Market.closed.is_(False)) + .where( + Market.active.is_(True), + Market.closed.is_(False), + Market.condition_id.in_(needed_conditions) + ) ) asset_end_timestamps = {} @@ -517,8 +539,7 @@ async def get_brain_state(session: AsyncSession = Depends(get_session)): (logger.bind(condition_id=condition_id, end_date=end_date, error=str(e)) .warning("failed_to_parse_end_date")) - # 2. Implications and Contradictions - implications_dicts = await _graph_manager.get_all_relationships() + # 3. Implications and Contradictions implications = [] contradictions = [] @@ -527,6 +548,12 @@ async def get_brain_state(session: AsyncSession = Depends(get_session)): child_token = condition_to_yes_token.get(rel["child_condition_id"]) if not parent_token or not child_token: + logger.bind( + relationship_id=rel.get("id"), + parent_condition_id=rel["parent_condition_id"], + child_condition_id=rel["child_condition_id"], + logic_type=rel["logic_type"] + ).debug("skipped_relationship_missing_tokens") continue if rel["logic_type"] == "IMPLIES": @@ -546,8 +573,7 @@ async def get_brain_state(session: AsyncSession = Depends(get_session)): ) ) - # 3. Partitions - partitions_lists = await _graph_manager.get_partition_groups() + # 4. Partitions partitions = [] # Build a lookup for mutual exclusion confidences @@ -565,15 +591,28 @@ async def get_brain_state(session: AsyncSession = Depends(get_session)): tok = condition_to_yes_token.get(cid) if tok: group_tokens.append(tok) + else: + (logger.bind(condition_id=cid, group=group) + .debug("partition_member_missing_token")) if len(group_tokens) >= 2: - # Calculate minimum confidence among mutually exclusive edges within the group + # Calculate minimum confidence and verify clique group_confs = [] + is_clique = True for j in range(len(group)): for k in range(j + 1, len(group)): conf = me_confidences.get(f"{group[j]}|{group[k]}") if conf is not None: group_confs.append(conf) + else: + is_clique = False + break + if not is_clique: + break + + if not is_clique: + logger.bind(group=group).debug("skipped_partition_not_clique") + continue partition_confidence = min(group_confs) if group_confs else 1.0