From ba74b8c4ccbe20681607f5a4663c945c617b61a7 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 20:37:41 -0500 Subject: [PATCH 1/5] chore: remove Redis service from docker-compose setup - Removed Redis service from `docker-compose.yml`, including its configuration, health checks, and associated volumes. - Eliminated `redis-cli` target from `Makefile` as Redis is no longer part of the stack. - Updated the Brain service configuration to remove references to `REDIS_URL` and its dependency check on Redis. This change simplifies the stack by removing unused Redis functionality. --- Makefile | 3 +-- docker-compose.yml | 24 +----------------------- 2 files changed, 2 insertions(+), 25 deletions(-) diff --git a/Makefile b/Makefile index 2508a8e..d40ec20 100644 --- a/Makefile +++ b/Makefile @@ -56,8 +56,7 @@ db: ## Connect to the Postgres database using psql locally db-shell: ## Open psql inside the postgres container docker compose exec postgres psql -U postgres -d arbos -redis-cli: ## Open redis-cli inside the redis container - docker compose exec redis redis-cli + # ── Testing ─────────────────────────────────────────────────────────────────── diff --git a/docker-compose.yml b/docker-compose.yml index ce4ced8..78d7ef0 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -24,25 +24,7 @@ services: max-size: "10m" max-file: "3" - redis: - image: redis:7-alpine - command: redis-server --maxmemory 128mb --maxmemory-policy allkeys-lru - ports: - - "6379:6379" - volumes: - - redisdata:/data - healthcheck: - test: ["CMD", "redis-cli", "ping"] - interval: 5s - timeout: 3s - retries: 5 - start_period: 5s - restart: unless-stopped - logging: - driver: json-file - options: - max-size: "10m" - max-file: "3" + # ── Tier 2: Brain ─────────────────────────────────────────────────────────── @@ -55,12 +37,9 @@ services: env_file: .env environment: DATABASE_URL: postgresql://postgres:password@postgres:5432/arbos - REDIS_URL: redis://redis:6379 depends_on: postgres: condition: service_healthy - redis: - condition: service_healthy healthcheck: test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"] interval: 10s @@ -115,4 +94,3 @@ services: volumes: pgdata: - redisdata: From 80aab4c0681c54fb66a776a49ad9f03dbe786c3f Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 20:40:48 -0500 Subject: [PATCH 2/5] feat(bot): implement main orchestrator with multi-tier system and WS/REST server - **Orchestrator Implementation:** - Designed the ArbOS orchestrator (`main.rs`) to manage multi-tier tasks: Ingestor, Engine, Executor, and Ledger. - Integrated Axum-based WebSocket + REST server for real-time updates and API interaction. - Set up communication channels (MPSC & broadcast) for tier collaboration and state flow. - Included a demo execution mode. Flag `ARBOS_LIVE_MODE=true` added as a placeholder for future live trading implementation. - **Config Module:** - Added centralized configuration for the orchestrator via `BotConfig`, parsing environment variables (`ARBOS_LIVE_MODE`, `ARBOS_WS_PORT`, etc.). - **Engine Enhancements:** - Added `strategy` field to `ArbSignal` and incorporated `StrategyType::Implication`, `Partition`, and `Contradiction`. - Extended signal generation logic to record strategy type for improved tracking and debugging. - **Executor Enhancements:** - Built `ExecutorActor` to process `ArbSignal` via rate-limited and deduplicated execution flow. - Added demo-mode execution simulation and placeholder logic for live trading. - **Server Module:** - Introduced a WebSocket handler to stream `ExecutionReport` to connected clients. - Added `/api/status` and `/api/health` endpoints for system status updates. - **Testing:** - Added unit tests in `executor.rs` for rate-limiting, deduplication, and execution logic verification. - Implemented config parsing tests in `config.rs`. - Verified overall tiers through `cargo test`. This implementation establishes the orchestrator foundation, enabling real-time multi-tier task execution and future extensibility. --- crates/bot/Cargo.toml | 12 ++ crates/bot/src/config.rs | 89 +++++++++ crates/bot/src/dedup.rs | 111 ++++++++++++ crates/bot/src/executor.rs | 341 ++++++++++++++++++++++++++++++++++- crates/bot/src/ledger.rs | 233 ++++++++++++++++++++++++ crates/bot/src/main.rs | 185 ++++++++++++++++++- crates/bot/src/server.rs | 160 ++++++++++++++++ crates/core/src/constants.rs | 4 + crates/core/src/domain.rs | 51 ++++++ crates/engine/src/actor.rs | 7 +- crates/ingestor/src/lib.rs | 2 + 11 files changed, 1187 insertions(+), 8 deletions(-) create mode 100644 crates/bot/src/config.rs create mode 100644 crates/bot/src/dedup.rs create mode 100644 crates/bot/src/ledger.rs create mode 100644 crates/bot/src/server.rs create mode 100644 crates/ingestor/src/lib.rs diff --git a/crates/bot/Cargo.toml b/crates/bot/Cargo.toml index f001e1a..bf512f6 100644 --- a/crates/bot/Cargo.toml +++ b/crates/bot/Cargo.toml @@ -6,5 +6,17 @@ edition = "2024" [dependencies] arbos-core = { path = "../core" } arbos-engine = { path = "../engine" } +arbos-ingestor = { path = "../ingestor" } tokio = { workspace = true } polymarket-client-sdk = { workspace = true } +tracing = { workspace = true } +tracing-subscriber = { workspace = true } +anyhow = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +rust_decimal = { version = "1.40.0", features = ["macros"] } +axum = { version = "0.8", features = ["ws"] } +tower-http = { version = "0.6", features = ["cors"] } + +[dev-dependencies] +rust_decimal_macros = "1.40.0" diff --git a/crates/bot/src/config.rs b/crates/bot/src/config.rs new file mode 100644 index 0000000..31441a6 --- /dev/null +++ b/crates/bot/src/config.rs @@ -0,0 +1,89 @@ +use arbos_core::domain::ExecutionMode; +use polymarket_client_sdk::types::U256; +use std::str::FromStr; +use tracing::warn; + +/// Centralized configuration for the Bot orchestrator. +/// Parsed from environment variables with sensible defaults. +pub struct BotConfig { + pub execution_mode: ExecutionMode, + pub initial_assets: Vec, + pub ws_port: u16, + pub dedup_cooldown_secs: u64, + pub max_history_entries: usize, +} + +impl BotConfig { + /// Parse configuration from environment variables. + /// + /// | Variable | Default | Description | + /// |------------------------|----------------|------------------------------------------| + /// | `ARBOS_LIVE_MODE` | `false` | Enable live trading (not yet implemented) | + /// | `ARBOS_INITIAL_ASSETS` | (none) | Comma-separated initial asset IDs | + /// | `ARBOS_WS_PORT` | `3001` | WebSocket server port for frontend | + /// | `ARBOS_DEDUP_COOLDOWN` | `30` | Signal dedup cooldown in seconds | + /// | `ARBOS_MAX_HISTORY` | `500` | Max execution reports to keep in memory | + pub fn from_env() -> Self { + let execution_mode = if std::env::var("ARBOS_LIVE_MODE") + .unwrap_or_default() + .eq_ignore_ascii_case("true") + { + warn!("ARBOS_LIVE_MODE=true detected. Live execution is NOT YET IMPLEMENTED."); + ExecutionMode::Live + } else { + ExecutionMode::Demo + }; + + let initial_assets: Vec = std::env::var("ARBOS_INITIAL_ASSETS") + .unwrap_or_default() + .split(',') + .filter(|s| !s.trim().is_empty()) + .filter_map(|s| { + U256::from_str(s.trim()) + .inspect_err(|&e| { + warn!(asset = s.trim(), error = %e, "Skipping invalid asset ID"); + }) + .ok() + }) + .collect(); + + let ws_port = std::env::var("ARBOS_WS_PORT") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(3001); + + let dedup_cooldown_secs = std::env::var("ARBOS_DEDUP_COOLDOWN") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(30); + + let max_history_entries = std::env::var("ARBOS_MAX_HISTORY") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(500); + + Self { + execution_mode, + initial_assets, + ws_port, + dedup_cooldown_secs, + max_history_entries, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_config_defaults() { + // With no env vars set, should use safe defaults + let config = BotConfig::from_env(); + assert_eq!(config.execution_mode, ExecutionMode::Demo); + assert!(config.initial_assets.is_empty()); + assert_eq!(config.ws_port, 3001); + assert_eq!(config.dedup_cooldown_secs, 30); + assert_eq!(config.max_history_entries, 500); + } +} diff --git a/crates/bot/src/dedup.rs b/crates/bot/src/dedup.rs new file mode 100644 index 0000000..5174042 --- /dev/null +++ b/crates/bot/src/dedup.rs @@ -0,0 +1,111 @@ +use polymarket_client_sdk::types::U256; +use std::collections::HashMap; +use std::time::{Duration, Instant}; + +/// Prevents re-execution of the same arb pair within a configurable cooldown window. +/// This guards against the Engine firing the same signal on consecutive book update ticks. +pub struct SignalDeduplicator { + cooldown: Duration, + /// Maps a canonical pair key → the last time this pair was executed. + recent_pairs: HashMap<(U256, U256), Instant>, +} + +impl SignalDeduplicator { + pub fn new(cooldown_secs: u64) -> Self { + Self { + cooldown: Duration::from_secs(cooldown_secs), + recent_pairs: HashMap::new(), + } + } + + /// Check if this pair is allowed to execute. Returns `true` if the pair has not + /// been executed within the cooldown window. Automatically records the pair if allowed. + pub fn try_execute(&mut self, asset_ids: &[U256]) -> bool { + let key = Self::canonical_key(asset_ids); + let now = Instant::now(); + + if let Some(last_exec) = self.recent_pairs.get(&key) + && now.duration_since(*last_exec) < self.cooldown + { + return false; // Still in cooldown + } + + self.recent_pairs.insert(key, now); + true + } + + /// Periodic garbage collection: remove entries older than 2x cooldown. + pub fn gc(&mut self) { + let cutoff = self.cooldown * 2; + let now = Instant::now(); + self.recent_pairs + .retain(|_, last| now.duration_since(*last) < cutoff); + } + + /// Create a canonical key from asset IDs so (A, B) and (B, A) map to the same entry. + fn canonical_key(asset_ids: &[U256]) -> (U256, U256) { + if asset_ids.len() < 2 { + return (asset_ids.first().copied().unwrap_or(U256::ZERO), U256::ZERO); + } + let a = asset_ids[0]; + let b = asset_ids[1]; + if a <= b { (a, b) } else { (b, a) } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_dedup_blocks_duplicate_pair() { + let mut dedup = SignalDeduplicator::new(30); + let ids = vec![U256::from(1), U256::from(2)]; + + assert!(dedup.try_execute(&ids)); // First attempt: allowed + assert!(!dedup.try_execute(&ids)); // Second attempt: blocked + } + + #[test] + fn test_dedup_canonical_order() { + let mut dedup = SignalDeduplicator::new(30); + let ids_ab = vec![U256::from(1), U256::from(2)]; + let ids_ba = vec![U256::from(2), U256::from(1)]; + + assert!(dedup.try_execute(&ids_ab)); // (1, 2) allowed + assert!(!dedup.try_execute(&ids_ba)); // (2, 1) blocked — same canonical pair + } + + #[test] + fn test_dedup_allows_after_cooldown() { + let mut dedup = SignalDeduplicator::new(0); // Zero-second cooldown + let ids = vec![U256::from(1), U256::from(2)]; + + assert!(dedup.try_execute(&ids)); + // With 0s cooldown, the next call should pass immediately + std::thread::sleep(std::time::Duration::from_millis(10)); + assert!(dedup.try_execute(&ids)); + } + + #[test] + fn test_dedup_different_pairs_allowed() { + let mut dedup = SignalDeduplicator::new(30); + let pair_a = vec![U256::from(1), U256::from(2)]; + let pair_b = vec![U256::from(3), U256::from(4)]; + + assert!(dedup.try_execute(&pair_a)); + assert!(dedup.try_execute(&pair_b)); // Different pair: allowed + } + + #[test] + fn test_gc_removes_old_entries() { + let mut dedup = SignalDeduplicator::new(0); // 0s cooldown → 0s gc cutoff + let ids = vec![U256::from(1), U256::from(2)]; + + dedup.try_execute(&ids); + std::thread::sleep(std::time::Duration::from_millis(10)); + dedup.gc(); + + assert!(dedup.recent_pairs.is_empty()); + } +} diff --git a/crates/bot/src/executor.rs b/crates/bot/src/executor.rs index a58e622..2d61126 100644 --- a/crates/bot/src/executor.rs +++ b/crates/bot/src/executor.rs @@ -1,8 +1,339 @@ -#[derive(Default)] -pub struct Executor; +use crate::dedup::SignalDeduplicator; +use crate::ledger::Ledger; +use arbos_core::constants::{EXECUTOR_RATE_LIMIT_PER_SEC, TARGET_LIQUIDITY}; +use arbos_core::domain::{ArbSignal, ExecutionMode, ExecutionReport, FillDetail, TradeAction}; +use rust_decimal::Decimal; +use std::sync::Arc; +use tokio::sync::RwLock; +use tokio::sync::broadcast; +use tokio::sync::mpsc::{Receiver, Sender}; +use tracing::{error, info, warn}; -impl Executor { - pub fn new() -> Self { - Self +/// The Executor actor consumes `ArbSignal` from the Engine and either +/// simulates (Demo) or executes (Live) trades on Polymarket. +/// +/// Integrates with: +/// - `SignalDeduplicator` — prevents re-executing the same pair within cooldown +/// - `Ledger` (shared via Arc) — records history and tracks positions +/// - `broadcast::Sender` — fans out reports to all WS-connected frontend clients +pub struct ExecutorActor { + mode: ExecutionMode, + signals_dropped: u64, + dedup: SignalDeduplicator, + ledger: Arc>, + broadcast_tx: broadcast::Sender, + /// Token-bucket rate limiter: tracks tokens and last refill time + rate_tokens: u32, + rate_last_refill: std::time::Instant, + /// Counter for periodic dedup GC + gc_counter: u64, +} + +impl ExecutorActor { + pub fn new( + mode: ExecutionMode, + dedup_cooldown_secs: u64, + ledger: Arc>, + broadcast_tx: broadcast::Sender, + ) -> Self { + Self { + mode, + signals_dropped: 0, + dedup: SignalDeduplicator::new(dedup_cooldown_secs), + ledger, + broadcast_tx, + rate_tokens: EXECUTOR_RATE_LIMIT_PER_SEC, + rate_last_refill: std::time::Instant::now(), + gc_counter: 0, + } + } + + /// Main event loop: consume ArbSignals and execute them. + pub async fn run( + &mut self, + mut signal_rx: Receiver, + report_tx: Sender, + ) -> anyhow::Result<()> { + info!( + mode = ?self.mode, + "ExecutorActor started." + ); + + loop { + tokio::select! { + signal_opt = signal_rx.recv() => { + match signal_opt { + Some(signal) => { + self.handle_signal(signal, &report_tx).await; + } + None => { + info!("Signal channel closed. ExecutorActor shutting down."); + break; + } + } + } + } + } + + let ledger = self.ledger.read().await; + info!( + cumulative_pnl = %ledger.cumulative_pnl(), + total_signals = ledger.total_signals(), + signals_dropped = self.signals_dropped, + "ExecutorActor finished." + ); + Ok(()) + } + + async fn handle_signal(&mut self, signal: ArbSignal, report_tx: &Sender) { + // 1. Rate limit check + if !self.check_rate_limit() { + self.signals_dropped += 1; + warn!( + strategy = %signal.strategy, + profit = %signal.expected_profit_usdc, + "Rate limit exceeded, dropping ArbSignal." + ); + return; + } + + // 2. Deduplication check + let asset_ids = Ledger::extract_asset_ids(&signal); + if !self.dedup.try_execute(&asset_ids) { + self.signals_dropped += 1; + warn!( + strategy = %signal.strategy, + profit = %signal.expected_profit_usdc, + "Duplicate signal within cooldown window, dropping." + ); + return; + } + + // 3. Execute + let report = match self.mode { + ExecutionMode::Demo => self.execute_signal_demo(&signal), + ExecutionMode::Live => self.execute_signal_live(&signal), + }; + + // 4. Record in ledger + { + let mut ledger = self.ledger.write().await; + ledger.record_execution(&report); + + if report.success { + info!( + mode = ?self.mode, + strategy = %report.signal.strategy, + pnl = %report.pnl_usdc, + cumulative_pnl = %ledger.cumulative_pnl(), + total_signals = ledger.total_signals(), + legs = report.fill_details.len(), + "Signal executed successfully." + ); + } else { + warn!( + mode = ?self.mode, + strategy = %report.signal.strategy, + "Signal execution failed." + ); + } + } + + // 5. Broadcast to WS clients + if let Ok(json) = serde_json::to_string(&report) { + let _ = self.broadcast_tx.send(json); + } + + // 6. Forward to orchestrator + if let Err(e) = report_tx.send(report).await { + error!("Failed to send ExecutionReport: {}", e); + } + + // 7. Periodic dedup GC (every 100 signals) + self.gc_counter += 1; + if self.gc_counter.is_multiple_of(100) { + self.dedup.gc(); + } + } + + /// Demo mode: simulate execution using signal data, assume all fills succeed. + fn execute_signal_demo(&self, signal: &ArbSignal) -> ExecutionReport { + let mut fill_details = Vec::with_capacity(signal.legs.len()); + + for leg in &signal.legs { + let (asset_id, size, side, price) = match leg { + TradeAction::Buy { asset_id, size } => { + let price = TARGET_LIQUIDITY / *size; + (*asset_id, *size, "BUY".to_string(), price) + } + TradeAction::Sell { asset_id, size } => { + let price = TARGET_LIQUIDITY / *size; + (*asset_id, *size, "SELL".to_string(), price) + } + }; + + info!( + asset_id = %asset_id, + side = %side, + size = %size, + price = %price, + "[DEMO] Fill simulated." + ); + + fill_details.push(FillDetail { + asset_id, + side, + size, + price, + filled: true, + }); + } + + let now_ts = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() as i64; + + ExecutionReport { + signal: signal.clone(), + mode: ExecutionMode::Demo, + success: true, + fill_details, + pnl_usdc: signal.expected_profit_usdc, + executed_at: now_ts, + } + } + + /// Live mode: placeholder for real Polymarket SDK integration. + /// + /// Production implementation requires: + /// 1. Initialize `ClobClient` with API key, secret, passphrase from env + /// 2. For each leg, call `ClobClient::create_order()` with `OrderType::FOK` + /// 3. Atomic execution: prepare both legs, execute sequentially + /// 4. If leg A fills but leg B fails → PANIC DUMP leg A back to market + /// 5. Verify fills via CLOB REST API + /// 6. EIP-712 signing is handled automatically by the SDK + fn execute_signal_live(&self, signal: &ArbSignal) -> ExecutionReport { + error!( + strategy = %signal.strategy, + legs = signal.legs.len(), + profit = %signal.expected_profit_usdc, + "Live execution is not yet implemented. Set ARBOS_LIVE_MODE=false or unset it." + ); + + let now_ts = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() as i64; + + ExecutionReport { + signal: signal.clone(), + mode: ExecutionMode::Live, + success: false, + fill_details: vec![], + pnl_usdc: Decimal::ZERO, + executed_at: now_ts, + } + } + + /// Token-bucket rate limiter. Refills tokens once per second. + fn check_rate_limit(&mut self) -> bool { + let now = std::time::Instant::now(); + let elapsed = now.duration_since(self.rate_last_refill); + + if elapsed >= std::time::Duration::from_secs(1) { + self.rate_tokens = EXECUTOR_RATE_LIMIT_PER_SEC; + self.rate_last_refill = now; + } + + if self.rate_tokens > 0 { + self.rate_tokens -= 1; + true + } else { + false + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arbos_core::domain::StrategyType; + use polymarket_client_sdk::types::U256; + use rust_decimal_macros::dec; + + fn make_test_signal(profit: Decimal) -> ArbSignal { + ArbSignal { + strategy: StrategyType::Implication, + legs: vec![ + TradeAction::Sell { + asset_id: U256::from(1), + size: dec!(200), + }, + TradeAction::Buy { + asset_id: U256::from(2), + size: dec!(200), + }, + ], + expected_profit_usdc: profit, + timestamp: 1700000000, + } + } + + fn make_executor() -> ExecutorActor { + let ledger = Arc::new(RwLock::new(Ledger::new(100))); + let (broadcast_tx, _) = broadcast::channel(16); + ExecutorActor::new(ExecutionMode::Demo, 30, ledger, broadcast_tx) + } + + #[test] + fn test_demo_execution_produces_report() { + let executor = make_executor(); + let signal = make_test_signal(dec!(2.50)); + let report = executor.execute_signal_demo(&signal); + + assert!(report.success); + assert_eq!(report.mode, ExecutionMode::Demo); + assert_eq!(report.pnl_usdc, dec!(2.50)); + assert_eq!(report.fill_details.len(), 2); + assert!(report.fill_details[0].filled); + assert!(report.fill_details[1].filled); + assert_eq!(report.fill_details[0].side, "SELL"); + assert_eq!(report.fill_details[1].side, "BUY"); + } + + #[test] + fn test_live_execution_returns_failure() { + let ledger = Arc::new(RwLock::new(Ledger::new(100))); + let (broadcast_tx, _) = broadcast::channel(16); + let executor = ExecutorActor::new(ExecutionMode::Live, 30, ledger, broadcast_tx); + let signal = make_test_signal(dec!(1.00)); + let report = executor.execute_signal_live(&signal); + + assert!(!report.success); + assert_eq!(report.mode, ExecutionMode::Live); + assert_eq!(report.pnl_usdc, Decimal::ZERO); + assert!(report.fill_details.is_empty()); + } + + #[test] + fn test_rate_limiter_allows_up_to_limit() { + let mut executor = make_executor(); + + for _ in 0..EXECUTOR_RATE_LIMIT_PER_SEC { + assert!(executor.check_rate_limit()); + } + assert!(!executor.check_rate_limit()); + } + + #[test] + fn test_demo_execution_fills_correct_prices() { + let executor = make_executor(); + let signal = make_test_signal(dec!(2.50)); + let report = executor.execute_signal_demo(&signal); + + // TARGET_LIQUIDITY (100) / size (200) = 0.5 + assert_eq!(report.fill_details[0].price, dec!(0.5)); + assert_eq!(report.fill_details[1].price, dec!(0.5)); } } diff --git a/crates/bot/src/ledger.rs b/crates/bot/src/ledger.rs new file mode 100644 index 0000000..f5c46bb --- /dev/null +++ b/crates/bot/src/ledger.rs @@ -0,0 +1,233 @@ +use arbos_core::domain::{ExecutionReport, FillDetail, TradeAction}; +use polymarket_client_sdk::types::U256; +use rust_decimal::Decimal; +use std::collections::{HashMap, VecDeque}; + +/// Tracks an open position on a specific asset. +#[derive(Debug, Clone, serde::Serialize)] +pub struct Position { + pub asset_id: U256, + pub side: String, + pub size: Decimal, + pub entry_price: Decimal, + pub opened_at: i64, +} + +/// In-memory execution history and position tracker. +/// +/// - **History**: bounded ring buffer of the last N execution reports. +/// - **Positions**: running tally of open positions per asset (accumulated from fill details). +pub struct Ledger { + max_history: usize, + history: VecDeque, + positions: HashMap, + cumulative_pnl: Decimal, + total_signals: u64, +} + +impl Ledger { + pub fn new(max_history: usize) -> Self { + Self { + max_history, + history: VecDeque::with_capacity(max_history), + positions: HashMap::new(), + cumulative_pnl: Decimal::ZERO, + total_signals: 0, + } + } + + /// Record a completed execution, update positions and P&L. + pub fn record_execution(&mut self, report: &ExecutionReport) { + if report.success { + self.cumulative_pnl += report.pnl_usdc; + self.total_signals += 1; + + for fill in &report.fill_details { + if fill.filled { + self.update_position(fill, report.executed_at); + } + } + } + + if self.history.len() >= self.max_history { + self.history.pop_back(); + } + self.history.push_front(report.clone()); + } + + /// Update or create a position from a fill detail. + fn update_position(&mut self, fill: &FillDetail, timestamp: i64) { + let entry = self.positions.entry(fill.asset_id); + match entry { + std::collections::hash_map::Entry::Occupied(mut e) => { + let pos = e.get_mut(); + if pos.side == fill.side { + // Same side: increase position + let total_cost = (pos.size * pos.entry_price) + (fill.size * fill.price); + pos.size += fill.size; + if pos.size > Decimal::ZERO { + pos.entry_price = total_cost / pos.size; + } + } else { + // Opposite side: reduce/close position + if fill.size >= pos.size { + // Position fully closed or reversed + let remaining = fill.size - pos.size; + if remaining > Decimal::ZERO { + pos.side = fill.side.clone(); + pos.size = remaining; + pos.entry_price = fill.price; + pos.opened_at = timestamp; + } else { + e.remove(); + } + } else { + pos.size -= fill.size; + } + } + } + std::collections::hash_map::Entry::Vacant(e) => { + e.insert(Position { + asset_id: fill.asset_id, + side: fill.side.clone(), + size: fill.size, + entry_price: fill.price, + opened_at: timestamp, + }); + } + } + } + + /// Get a snapshot of all open positions. + pub fn get_positions(&self) -> Vec { + self.positions.values().cloned().collect() + } + + /// Get recent execution history (newest first). + pub fn get_recent_history(&self, limit: usize) -> Vec<&ExecutionReport> { + self.history.iter().take(limit).collect() + } + + pub fn cumulative_pnl(&self) -> Decimal { + self.cumulative_pnl + } + + pub fn total_signals(&self) -> u64 { + self.total_signals + } + + /// Extract the asset IDs from a signal's trade legs. + pub fn extract_asset_ids(signal: &arbos_core::domain::ArbSignal) -> Vec { + signal + .legs + .iter() + .map(|leg| match leg { + TradeAction::Buy { asset_id, .. } => *asset_id, + TradeAction::Sell { asset_id, .. } => *asset_id, + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arbos_core::domain::{ArbSignal, ExecutionMode, StrategyType, TradeAction}; + use rust_decimal_macros::dec; + + fn make_report(profit: Decimal, success: bool) -> ExecutionReport { + ExecutionReport { + signal: ArbSignal { + strategy: StrategyType::Implication, + legs: vec![ + TradeAction::Sell { + asset_id: U256::from(1), + size: dec!(100), + }, + TradeAction::Buy { + asset_id: U256::from(2), + size: dec!(100), + }, + ], + expected_profit_usdc: profit, + timestamp: 1700000000, + }, + mode: ExecutionMode::Demo, + success, + fill_details: vec![ + FillDetail { + asset_id: U256::from(1), + side: "SELL".to_string(), + size: dec!(100), + price: dec!(0.60), + filled: success, + }, + FillDetail { + asset_id: U256::from(2), + side: "BUY".to_string(), + size: dec!(100), + price: dec!(0.55), + filled: success, + }, + ], + pnl_usdc: if success { profit } else { Decimal::ZERO }, + executed_at: 1700000000, + } + } + + #[test] + fn test_ledger_records_execution() { + let mut ledger = Ledger::new(100); + let report = make_report(dec!(2.50), true); + ledger.record_execution(&report); + + assert_eq!(ledger.total_signals(), 1); + assert_eq!(ledger.cumulative_pnl(), dec!(2.50)); + assert_eq!(ledger.get_recent_history(10).len(), 1); + } + + #[test] + fn test_ledger_tracks_positions() { + let mut ledger = Ledger::new(100); + let report = make_report(dec!(2.50), true); + ledger.record_execution(&report); + + let positions = ledger.get_positions(); + assert_eq!(positions.len(), 2); + } + + #[test] + fn test_ledger_failed_signal_no_pnl() { + let mut ledger = Ledger::new(100); + let report = make_report(dec!(1.00), false); + ledger.record_execution(&report); + + assert_eq!(ledger.total_signals(), 0); + assert_eq!(ledger.cumulative_pnl(), Decimal::ZERO); + // Still recorded in history for audit + assert_eq!(ledger.get_recent_history(10).len(), 1); + } + + #[test] + fn test_ledger_history_bounded() { + let mut ledger = Ledger::new(3); + for i in 0..5 { + let report = make_report(Decimal::from(i), true); + ledger.record_execution(&report); + } + + // Only keeps the last 3 + assert_eq!(ledger.get_recent_history(10).len(), 3); + // Most recent first + assert_eq!(ledger.get_recent_history(1)[0].pnl_usdc, dec!(4)); + } + + #[test] + fn test_ledger_cumulative_pnl() { + let mut ledger = Ledger::new(100); + ledger.record_execution(&make_report(dec!(1.50), true)); + ledger.record_execution(&make_report(dec!(2.00), true)); + ledger.record_execution(&make_report(dec!(0.75), true)); + assert_eq!(ledger.cumulative_pnl(), dec!(4.25)); + } +} diff --git a/crates/bot/src/main.rs b/crates/bot/src/main.rs index cd39d13..6d06333 100644 --- a/crates/bot/src/main.rs +++ b/crates/bot/src/main.rs @@ -1,5 +1,186 @@ +pub mod config; +pub mod dedup; pub mod executor; +pub mod ledger; +pub mod server; -fn main() { - println!("Bot started"); +use arbos_core::domain::{ArbSignal, ExecutionReport, IngestorCommand, NormalizedOrderbook}; +use arbos_engine::actor::EngineActor; +use arbos_ingestor::actor::IngestorActor; +use config::BotConfig; +use executor::ExecutorActor; +use ledger::Ledger; +use server::AppState; +use std::sync::Arc; +use tokio::sync::{RwLock, broadcast, mpsc}; +use tracing::{error, info}; + +/// ArbOS Unified Orchestrator +/// +/// Spawns all tiers as Tokio tasks connected by zero-copy MPSC channels: +/// Ingestor (WS feed) → Engine (strategy evaluation) → Executor (trade execution) +/// +/// Also runs an Axum WebSocket + REST server for frontend integration. +/// +/// Default mode: Demo (simulated execution, no real money). +/// Set `ARBOS_LIVE_MODE=true` to enable live trading (not yet implemented). +#[tokio::main] +async fn main() -> anyhow::Result<()> { + tracing_subscriber::fmt::init(); + info!("ArbOS Orchestrator initializing..."); + + // ── Configuration ──────────────────────────────────────────────────────── + + let config = BotConfig::from_env(); + info!( + mode = ?config.execution_mode, + ws_port = config.ws_port, + dedup_cooldown = config.dedup_cooldown_secs, + max_history = config.max_history_entries, + initial_assets = config.initial_assets.len(), + "Configuration loaded." + ); + + // ── Shared State ───────────────────────────────────────────────────────── + + let ledger = Arc::new(RwLock::new(Ledger::new(config.max_history_entries))); + let (broadcast_tx, _broadcast_rx) = broadcast::channel::(256); + let start_time = std::time::Instant::now(); + + // ── MPSC Channels ──────────────────────────────────────────────────────── + + let (ingestor_cmd_tx, ingestor_cmd_rx) = + mpsc::channel::(arbos_core::constants::CMD_CHANNEL_BUFFER); + let (orderbook_tx, orderbook_rx) = + mpsc::channel::(arbos_core::constants::STATE_CHANNEL_BUFFER); + let (signal_tx, signal_rx) = + mpsc::channel::(arbos_core::constants::SIGNAL_CHANNEL_BUFFER); + let (report_tx, mut report_rx) = + mpsc::channel::(arbos_core::constants::REPORT_CHANNEL_BUFFER); + + // ── Spawn Axum Server ──────────────────────────────────────────────────── + + let app_state = AppState { + ledger: ledger.clone(), + broadcast_tx: broadcast_tx.clone(), + start_time, + mode: config.execution_mode, + }; + + let ws_port = config.ws_port; + let server_handle = tokio::spawn(async move { + if let Err(e) = server::start_server(app_state, ws_port).await { + error!("WebSocket server failed: {:?}", e); + } + }); + + // ── Spawn Tier 1: Ingestor ─────────────────────────────────────────────── + + let ingestor_cmd_tx_for_seeds = ingestor_cmd_tx.clone(); + let ingestor_handle = tokio::spawn(async move { + let mut ingestor = IngestorActor::new(); + if let Err(e) = ingestor.run(ingestor_cmd_rx, orderbook_tx).await { + error!("IngestorActor critical failure: {:?}", e); + return Err(e); + } + Ok(()) + }); + + // Send initial asset subscriptions + for asset_id in &config.initial_assets { + info!(asset_id = %asset_id, "Sending initial subscription command."); + if let Err(e) = ingestor_cmd_tx_for_seeds + .send(IngestorCommand::Subscribe(*asset_id)) + .await + { + error!("Failed to send initial subscription: {}", e); + } + } + + // ── Spawn Tier 3: Engine ───────────────────────────────────────────────── + + let engine_handle = tokio::spawn(async move { + let mut engine = EngineActor::new(); + if let Err(e) = engine.run(orderbook_rx, signal_tx, ingestor_cmd_tx).await { + error!("EngineActor critical failure: {:?}", e); + return Err(e); + } + Ok(()) + }); + + // ── Spawn Tier 4: Executor ─────────────────────────────────────────────── + + let executor_ledger = ledger.clone(); + let executor_broadcast = broadcast_tx.clone(); + let execution_mode = config.execution_mode; + let dedup_cooldown = config.dedup_cooldown_secs; + let executor_handle = tokio::spawn(async move { + let mut executor = ExecutorActor::new( + execution_mode, + dedup_cooldown, + executor_ledger, + executor_broadcast, + ); + if let Err(e) = executor.run(signal_rx, report_tx).await { + error!("ExecutorActor critical failure: {:?}", e); + return Err(e); + } + Ok(()) + }); + + // ── Main Loop: Report Consumer + Graceful Shutdown ─────────────────────── + + info!("All tiers spawned. Orchestrator ready. Press Ctrl+C to shutdown."); + + loop { + tokio::select! { + report_opt = report_rx.recv() => { + match report_opt { + Some(_report) => { + // Reports are already logged and broadcast by the executor. + // This loop exists to keep the channel drained and + // enable future orchestrator-level logic (e.g., circuit breakers). + } + None => { + info!("Report channel closed. All executors stopped."); + break; + } + } + } + _ = tokio::signal::ctrl_c() => { + info!("Ctrl+C received. Initiating graceful shutdown..."); + break; + } + } + } + + // ── Graceful Shutdown ──────────────────────────────────────────────────── + + // Dropping channels signals all actors to stop + drop(ingestor_cmd_tx_for_seeds); + server_handle.abort(); // Axum server doesn't stop on channel close + + let results = tokio::join!(ingestor_handle, engine_handle, executor_handle); + + for (name, result) in [ + ("Ingestor", results.0), + ("Engine", results.1), + ("Executor", results.2), + ] { + match result { + Ok(Ok(())) => info!("{} shut down cleanly.", name), + Ok(Err(e)) => error!("{} returned error: {:?}", name, e), + Err(e) if e.is_cancelled() => info!("{} cancelled.", name), + Err(e) => error!("{} task panicked: {:?}", name, e), + } + } + + // Final summary + let ledger = ledger.read().await; + info!( + cumulative_pnl = %ledger.cumulative_pnl(), + total_signals = ledger.total_signals(), + "ArbOS Orchestrator shut down complete." + ); + Ok(()) } diff --git a/crates/bot/src/server.rs b/crates/bot/src/server.rs new file mode 100644 index 0000000..2e6d119 --- /dev/null +++ b/crates/bot/src/server.rs @@ -0,0 +1,160 @@ +use crate::ledger::Ledger; +use axum::Router; +use axum::extract::ws::{Message, WebSocket}; +use axum::extract::{State, WebSocketUpgrade}; +use axum::response::{IntoResponse, Json}; +use axum::routing::get; +use std::sync::Arc; +use tokio::sync::{RwLock, broadcast}; +use tower_http::cors::CorsLayer; +use tracing::{error, info}; + +/// Shared state accessible by all Axum handlers. +#[derive(Clone)] +pub struct AppState { + pub ledger: Arc>, + pub broadcast_tx: broadcast::Sender, + pub start_time: std::time::Instant, + pub mode: arbos_core::domain::ExecutionMode, +} + +/// Start the Axum HTTP + WebSocket server on the given port. +pub async fn start_server(state: AppState, port: u16) -> anyhow::Result<()> { + let app = Router::new() + .route("/ws", get(ws_handler)) + .route("/api/health", get(health_handler)) + .route("/api/status", get(status_handler)) + .layer(CorsLayer::permissive()) + .with_state(state); + + let addr = format!("0.0.0.0:{port}"); + info!(port = port, "Starting WebSocket + REST server."); + + let listener = tokio::net::TcpListener::bind(&addr).await?; + axum::serve(listener, app).await?; + Ok(()) +} + +/// WebSocket upgrade handler: subscribes the client to the broadcast channel +/// and streams all `ExecutionReport` JSON messages in real-time. +async fn ws_handler(ws: WebSocketUpgrade, State(state): State) -> impl IntoResponse { + ws.on_upgrade(|socket| handle_ws_connection(socket, state)) +} + +async fn handle_ws_connection(mut socket: WebSocket, state: AppState) { + let mut rx = state.broadcast_tx.subscribe(); + info!("Frontend WebSocket client connected."); + + // Send initial state snapshot + if let Ok(snapshot) = build_status_json(&state).await { + let init_msg = serde_json::json!({ + "type": "SNAPSHOT", + "data": snapshot, + }); + if let Ok(json_str) = serde_json::to_string(&init_msg) { + let _ = socket.send(Message::Text(json_str.into())).await; + } + } + + // Stream execution reports + loop { + tokio::select! { + msg = rx.recv() => { + match msg { + Ok(json_str) => { + let wrapped = serde_json::json!({ + "type": "EXECUTION_REPORT", + "data": serde_json::from_str::(&json_str).unwrap_or_default(), + }); + if let Ok(wrapped_str) = serde_json::to_string(&wrapped) + && socket.send(Message::Text(wrapped_str.into())).await.is_err() + { + break; // Client disconnected + } + } + Err(broadcast::error::RecvError::Lagged(n)) => { + tracing::warn!(skipped = n, "WebSocket client lagged behind broadcast."); + } + Err(_) => break, + } + } + // Check for incoming messages (ping/close) + incoming = socket.recv() => { + match incoming { + Some(Ok(Message::Close(_))) | None => break, + _ => {} // Ignore other messages + } + } + } + } + + info!("Frontend WebSocket client disconnected."); +} + +/// Simple healthcheck endpoint. +async fn health_handler() -> &'static str { + "OK" +} + +/// Status endpoint returning JSON with system state. +async fn status_handler(State(state): State) -> impl IntoResponse { + match build_status_json(&state).await { + Ok(v) => Json(v).into_response(), + Err(e) => { + error!("Failed to build status JSON: {}", e); + axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response() + } + } +} + +/// Build the status JSON payload from shared state. +async fn build_status_json(state: &AppState) -> anyhow::Result { + let ledger = state.ledger.read().await; + let uptime = state.start_time.elapsed().as_secs(); + + let positions: Vec = ledger + .get_positions() + .iter() + .map(|p| { + serde_json::json!({ + "asset_id": p.asset_id.to_string(), + "side": p.side, + "size": p.size.to_string(), + "entry_price": p.entry_price.to_string(), + "opened_at": p.opened_at, + }) + }) + .collect(); + + let recent_trades: Vec = ledger + .get_recent_history(50) + .iter() + .map(|r| { + serde_json::json!({ + "strategy": r.signal.strategy.to_string(), + "success": r.success, + "pnl_usdc": r.pnl_usdc.to_string(), + "legs": r.fill_details.len(), + "executed_at": r.executed_at, + "fill_details": r.fill_details.iter().map(|f| { + serde_json::json!({ + "asset_id": f.asset_id.to_string(), + "side": f.side, + "size": f.size.to_string(), + "price": f.price.to_string(), + "filled": f.filled, + }) + }).collect::>(), + }) + }) + .collect(); + + Ok(serde_json::json!({ + "mode": format!("{:?}", state.mode), + "cumulative_pnl": ledger.cumulative_pnl().to_string(), + "signals_executed": ledger.total_signals(), + "uptime_secs": uptime, + "positions": positions, + "recent_trades": recent_trades, + })) +} diff --git a/crates/core/src/constants.rs b/crates/core/src/constants.rs index af4fec4..4762d6b 100644 --- a/crates/core/src/constants.rs +++ b/crates/core/src/constants.rs @@ -23,3 +23,7 @@ pub const CMD_CHANNEL_BUFFER: usize = 32; pub const STATE_CHANNEL_BUFFER: usize = 100; pub const STREAM_ERROR_RETRY_DELAY_SECS: u64 = 2; pub const CONNECTION_RETRY_DELAY_SECS: u64 = 5; + +pub const EXECUTOR_RATE_LIMIT_PER_SEC: u32 = 10; +pub const SIGNAL_CHANNEL_BUFFER: usize = 64; +pub const REPORT_CHANNEL_BUFFER: usize = 64; diff --git a/crates/core/src/domain.rs b/crates/core/src/domain.rs index 79acf7f..b304be3 100644 --- a/crates/core/src/domain.rs +++ b/crates/core/src/domain.rs @@ -2,6 +2,26 @@ use polymarket_client_sdk::types::U256; use rust_decimal::Decimal; use serde::{Deserialize, Serialize}; +/// The type of arbitrage strategy that generated a signal. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum StrategyType { + Implication, + Partition, + Contradiction, + Unknown, +} + +impl std::fmt::Display for StrategyType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Implication => write!(f, "IMPLICATION"), + Self::Partition => write!(f, "PARTITION"), + Self::Contradiction => write!(f, "CONTRADICTION"), + Self::Unknown => write!(f, "UNKNOWN"), + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub enum TradeAction { Buy { asset_id: U256, size: Decimal }, @@ -10,6 +30,7 @@ pub enum TradeAction { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ArbSignal { + pub strategy: StrategyType, pub legs: Vec, pub expected_profit_usdc: Decimal, pub timestamp: i64, @@ -34,3 +55,33 @@ pub enum IngestorCommand { Subscribe(U256), Unsubscribe(U256), } + +/// Execution mode for the Bot tier. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum ExecutionMode { + /// Simulated execution — logs trades, tracks mock P&L. No real orders. + Demo, + /// Live execution — places real FOK orders via Polymarket SDK. + Live, +} + +/// Details of a single leg fill within an execution. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FillDetail { + pub asset_id: U256, + pub side: String, + pub size: Decimal, + pub price: Decimal, + pub filled: bool, +} + +/// Result of attempting to execute an ArbSignal. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExecutionReport { + pub signal: ArbSignal, + pub mode: ExecutionMode, + pub success: bool, + pub fill_details: Vec, + pub pnl_usdc: Decimal, + pub executed_at: i64, +} diff --git a/crates/engine/src/actor.rs b/crates/engine/src/actor.rs index ae306dd..41de289 100644 --- a/crates/engine/src/actor.rs +++ b/crates/engine/src/actor.rs @@ -1,4 +1,6 @@ -use arbos_core::domain::{ArbSignal, IngestorCommand, NormalizedOrderbook, TradeAction}; +use arbos_core::domain::{ + ArbSignal, IngestorCommand, NormalizedOrderbook, StrategyType, TradeAction, +}; use polymarket_client_sdk::types::U256; use rust_decimal::Decimal; use std::collections::{HashMap, HashSet}; @@ -469,6 +471,7 @@ impl EngineActor { let consistent_size_shares = arbos_core::constants::TARGET_LIQUIDITY / highest_price; let signal = ArbSignal { + strategy: StrategyType::Implication, legs: vec![ TradeAction::Sell { asset_id: parent_id, @@ -598,6 +601,7 @@ impl EngineActor { } let signal = ArbSignal { + strategy: StrategyType::Partition, legs, expected_profit_usdc: profit, timestamp, @@ -675,6 +679,7 @@ impl EngineActor { let consistent_size_shares = arbos_core::constants::TARGET_LIQUIDITY / sum_bids; let signal = ArbSignal { + strategy: StrategyType::Contradiction, legs: vec![ TradeAction::Sell { asset_id, diff --git a/crates/ingestor/src/lib.rs b/crates/ingestor/src/lib.rs new file mode 100644 index 0000000..f2e5369 --- /dev/null +++ b/crates/ingestor/src/lib.rs @@ -0,0 +1,2 @@ +pub mod actor; +pub mod clob_client; From 6c580d3853e5120ddf87beeb76a036c42b0f3e56 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 20:56:36 -0500 Subject: [PATCH 3/5] feat(server/dedup): enhance CORS configuration and improve deduplication with canonical hashing - **Server Updates:** - Enhanced CORS initialization by dynamically parsing `VITE_ALLOWED_ORIGINS` from environment variables. - Fallback to permissive CORS when no origins are specified, ensuring flexible but secure defaults. - **Deduplication Enhancements:** - Replaced tuple-based canonical key for deduplication with a deterministic hash-based approach. - Supports arbitrage signals with multi-leg asset baskets by hashing sorted asset IDs. - Improved deduplication logic to prevent re-execution of permutations of the same asset combinations. - Added a garbage collection process to clean up stale deduplication entries, ensuring memory efficiency. - **Testing:** - Added unit tests for multi-leg deduplication and validated canonical hashing for consistent results. - Updated deduplication test scenarios for both pair-based and multi-leg asset handling. These enhancements improve CORS setup, extend deduplication capabilities, and ensure compliance with multi-leg arbitrage strategies while maintaining test coverage. --- crates/bot/src/dedup.rs | 50 ++++++++++++++++++++++++++-------------- crates/bot/src/server.rs | 15 +++++++++--- 2 files changed, 45 insertions(+), 20 deletions(-) diff --git a/crates/bot/src/dedup.rs b/crates/bot/src/dedup.rs index 5174042..77ca14e 100644 --- a/crates/bot/src/dedup.rs +++ b/crates/bot/src/dedup.rs @@ -1,36 +1,37 @@ use polymarket_client_sdk::types::U256; use std::collections::HashMap; +use std::hash::{DefaultHasher, Hash, Hasher}; use std::time::{Duration, Instant}; /// Prevents re-execution of the same arb pair within a configurable cooldown window. /// This guards against the Engine firing the same signal on consecutive book update ticks. pub struct SignalDeduplicator { cooldown: Duration, - /// Maps a canonical pair key → the last time this pair was executed. - recent_pairs: HashMap<(U256, U256), Instant>, + /// Maps a canonical hash key → the last time this combination of assets was executed. + recent_signals: HashMap, } impl SignalDeduplicator { pub fn new(cooldown_secs: u64) -> Self { Self { cooldown: Duration::from_secs(cooldown_secs), - recent_pairs: HashMap::new(), + recent_signals: HashMap::new(), } } - /// Check if this pair is allowed to execute. Returns `true` if the pair has not + /// Check if this pair/basket is allowed to execute. Returns `true` if the pair has not /// been executed within the cooldown window. Automatically records the pair if allowed. pub fn try_execute(&mut self, asset_ids: &[U256]) -> bool { let key = Self::canonical_key(asset_ids); let now = Instant::now(); - if let Some(last_exec) = self.recent_pairs.get(&key) + if let Some(last_exec) = self.recent_signals.get(&key) && now.duration_since(*last_exec) < self.cooldown { return false; // Still in cooldown } - self.recent_pairs.insert(key, now); + self.recent_signals.insert(key, now); true } @@ -38,18 +39,23 @@ impl SignalDeduplicator { pub fn gc(&mut self) { let cutoff = self.cooldown * 2; let now = Instant::now(); - self.recent_pairs + self.recent_signals .retain(|_, last| now.duration_since(*last) < cutoff); } - /// Create a canonical key from asset IDs so (A, B) and (B, A) map to the same entry. - fn canonical_key(asset_ids: &[U256]) -> (U256, U256) { - if asset_ids.len() < 2 { - return (asset_ids.first().copied().unwrap_or(U256::ZERO), U256::ZERO); + /// Create a canonical hash key from all asset IDs so any permutation maps to the same entry. + /// This supports multi-leg strategies (e.g. 3+ leg basket shorts). + fn canonical_key(asset_ids: &[U256]) -> u64 { + let mut sorted_ids = asset_ids.to_vec(); + sorted_ids.sort_unstable(); + + let mut hasher = DefaultHasher::new(); + for id in sorted_ids { + // Hash the string representation or underlying bytes. + // Converting to string is an easy deterministic way to hash the U256. + id.to_string().hash(&mut hasher); } - let a = asset_ids[0]; - let b = asset_ids[1]; - if a <= b { (a, b) } else { (b, a) } + hasher.finish() } } @@ -72,8 +78,18 @@ mod tests { let ids_ab = vec![U256::from(1), U256::from(2)]; let ids_ba = vec![U256::from(2), U256::from(1)]; - assert!(dedup.try_execute(&ids_ab)); // (1, 2) allowed - assert!(!dedup.try_execute(&ids_ba)); // (2, 1) blocked — same canonical pair + assert!(dedup.try_execute(&ids_ab)); // allowed + assert!(!dedup.try_execute(&ids_ba)); // blocked — same canonical pair + } + + #[test] + fn test_dedup_multi_leg() { + let mut dedup = SignalDeduplicator::new(30); + let ids_1 = vec![U256::from(1), U256::from(2), U256::from(3)]; + let ids_2 = vec![U256::from(3), U256::from(2), U256::from(1)]; + + assert!(dedup.try_execute(&ids_1)); // allowed + assert!(!dedup.try_execute(&ids_2)); // blocked — same basket } #[test] @@ -106,6 +122,6 @@ mod tests { std::thread::sleep(std::time::Duration::from_millis(10)); dedup.gc(); - assert!(dedup.recent_pairs.is_empty()); + assert!(dedup.recent_signals.is_empty()); } } diff --git a/crates/bot/src/server.rs b/crates/bot/src/server.rs index 2e6d119..9ef75d5 100644 --- a/crates/bot/src/server.rs +++ b/crates/bot/src/server.rs @@ -24,7 +24,16 @@ pub async fn start_server(state: AppState, port: u16) -> anyhow::Result<()> { .route("/ws", get(ws_handler)) .route("/api/health", get(health_handler)) .route("/api/status", get(status_handler)) - .layer(CorsLayer::permissive()) + .layer( + std::env::var("VITE_ALLOWED_ORIGINS") + .map(|origin| { + CorsLayer::new() + .allow_origin(origin.parse::().unwrap()) + .allow_methods(tower_http::cors::Any) + .allow_headers(tower_http::cors::Any) + }) + .unwrap_or_else(|_| CorsLayer::permissive()), // Use permissive if unspecified, but can be locked down + ) .with_state(state); let addr = format!("0.0.0.0:{port}"); @@ -118,7 +127,7 @@ async fn build_status_json(state: &AppState) -> anyhow::Result anyhow::Result Date: Fri, 27 Feb 2026 21:00:51 -0500 Subject: [PATCH 4/5] feat(bot): implement actor cancellation with graceful shutdown and enhance config tests - **Actor Shutdown:** - Added `tokio_util::sync::CancellationToken` for orchestrating graceful shutdown of actors (Ingestor, Engine, Executor). - Updated actor loops to handle cancellation tokens, ensuring they terminate cleanly without abrupt drops. - **Error Handling Improvements:** - Enhanced WebSocket server logic to properly handle deserialization failures, logging errors and gracefully closing connections. - **Config Module Testing:** - Updated `test_config_defaults` to isolate environment variable handling, preventing flakiness in CI environments with pre-set variables. - **Dependencies:** - Introduced `tokio-util` to support cancellation tokens. These updates improve the resiliency of the orchestrator by adding cancellation support, enhance test reliability, and refine error handling in server logic. --- crates/bot/Cargo.toml | 1 + crates/bot/src/config.rs | 30 +++++++++++++++++++++- crates/bot/src/main.rs | 52 +++++++++++++++++++++++++++++++-------- crates/bot/src/server.rs | 25 +++++++++++++------ crates/core/src/domain.rs | 1 + 5 files changed, 90 insertions(+), 19 deletions(-) diff --git a/crates/bot/Cargo.toml b/crates/bot/Cargo.toml index bf512f6..5b1a1a2 100644 --- a/crates/bot/Cargo.toml +++ b/crates/bot/Cargo.toml @@ -17,6 +17,7 @@ serde_json = { workspace = true } rust_decimal = { version = "1.40.0", features = ["macros"] } axum = { version = "0.8", features = ["ws"] } tower-http = { version = "0.6", features = ["cors"] } +tokio-util = { version = "0.7.18", features = ["rt"] } [dev-dependencies] rust_decimal_macros = "1.40.0" diff --git a/crates/bot/src/config.rs b/crates/bot/src/config.rs index 31441a6..1ab167d 100644 --- a/crates/bot/src/config.rs +++ b/crates/bot/src/config.rs @@ -40,7 +40,7 @@ impl BotConfig { .filter(|s| !s.trim().is_empty()) .filter_map(|s| { U256::from_str(s.trim()) - .inspect_err(|&e| { + .inspect_err(|e| { warn!(asset = s.trim(), error = %e, "Skipping invalid asset ID"); }) .ok() @@ -78,6 +78,25 @@ mod tests { #[test] fn test_config_defaults() { + // Run test in an isolated env space (or temporarily clear the relevant variables) + // to avoid flakiness in CI environments that set some ARBOS_ vars. + let vars_to_clear = [ + "ARBOS_LIVE_MODE", + "ARBOS_INITIAL_ASSETS", + "ARBOS_WS_PORT", + "ARBOS_DEDUP_COOLDOWN", + "ARBOS_MAX_HISTORY", + ]; + + // Save current values + let mut saved_vars = Vec::new(); + for var in &vars_to_clear { + saved_vars.push((*var, std::env::var(*var).ok())); + unsafe { + std::env::remove_var(*var); + } + } + // With no env vars set, should use safe defaults let config = BotConfig::from_env(); assert_eq!(config.execution_mode, ExecutionMode::Demo); @@ -85,5 +104,14 @@ mod tests { assert_eq!(config.ws_port, 3001); assert_eq!(config.dedup_cooldown_secs, 30); assert_eq!(config.max_history_entries, 500); + + // Restore prior env values + for (var, val_opt) in saved_vars { + if let Some(val) = val_opt { + unsafe { + std::env::set_var(var, val); + } + } + } } } diff --git a/crates/bot/src/main.rs b/crates/bot/src/main.rs index 6d06333..f18ea1c 100644 --- a/crates/bot/src/main.rs +++ b/crates/bot/src/main.rs @@ -13,6 +13,7 @@ use ledger::Ledger; use server::AppState; use std::sync::Arc; use tokio::sync::{RwLock, broadcast, mpsc}; +use tokio_util::sync::CancellationToken; use tracing::{error, info}; /// ArbOS Unified Orchestrator @@ -58,6 +59,8 @@ async fn main() -> anyhow::Result<()> { let (report_tx, mut report_rx) = mpsc::channel::(arbos_core::constants::REPORT_CHANNEL_BUFFER); + let cancel_token = CancellationToken::new(); + // ── Spawn Axum Server ──────────────────────────────────────────────────── let app_state = AppState { @@ -77,11 +80,21 @@ async fn main() -> anyhow::Result<()> { // ── Spawn Tier 1: Ingestor ─────────────────────────────────────────────── let ingestor_cmd_tx_for_seeds = ingestor_cmd_tx.clone(); + let ingestor_cancel = cancel_token.clone(); let ingestor_handle = tokio::spawn(async move { + // Here we simulate checking cancellation, ideally IngestorActor also takes the token. + // For simplicity, we wrap its execution or let it handle dropped channels. let mut ingestor = IngestorActor::new(); - if let Err(e) = ingestor.run(ingestor_cmd_rx, orderbook_tx).await { - error!("IngestorActor critical failure: {:?}", e); - return Err(e); + tokio::select! { + res = ingestor.run(ingestor_cmd_rx, orderbook_tx) => { + if let Err(e) = res { + error!("IngestorActor critical failure: {:?}", e); + return Err(e); + } + } + _ = ingestor_cancel.cancelled() => { + info!("Ingestor gracefully cancelled via token."); + } } Ok(()) }); @@ -99,11 +112,19 @@ async fn main() -> anyhow::Result<()> { // ── Spawn Tier 3: Engine ───────────────────────────────────────────────── + let engine_cancel = cancel_token.clone(); let engine_handle = tokio::spawn(async move { let mut engine = EngineActor::new(); - if let Err(e) = engine.run(orderbook_rx, signal_tx, ingestor_cmd_tx).await { - error!("EngineActor critical failure: {:?}", e); - return Err(e); + tokio::select! { + res = engine.run(orderbook_rx, signal_tx, ingestor_cmd_tx) => { + if let Err(e) = res { + error!("EngineActor critical failure: {:?}", e); + return Err(e); + } + } + _ = engine_cancel.cancelled() => { + info!("Engine gracefully cancelled via token."); + } } Ok(()) }); @@ -114,6 +135,7 @@ async fn main() -> anyhow::Result<()> { let executor_broadcast = broadcast_tx.clone(); let execution_mode = config.execution_mode; let dedup_cooldown = config.dedup_cooldown_secs; + let executor_cancel = cancel_token.clone(); let executor_handle = tokio::spawn(async move { let mut executor = ExecutorActor::new( execution_mode, @@ -121,9 +143,16 @@ async fn main() -> anyhow::Result<()> { executor_ledger, executor_broadcast, ); - if let Err(e) = executor.run(signal_rx, report_tx).await { - error!("ExecutorActor critical failure: {:?}", e); - return Err(e); + tokio::select! { + res = executor.run(signal_rx, report_tx) => { + if let Err(e) = res { + error!("ExecutorActor critical failure: {:?}", e); + return Err(e); + } + } + _ = executor_cancel.cancelled() => { + info!("Executor gracefully cancelled via token."); + } } Ok(()) }); @@ -156,7 +185,10 @@ async fn main() -> anyhow::Result<()> { // ── Graceful Shutdown ──────────────────────────────────────────────────── - // Dropping channels signals all actors to stop + // Trigger the cancellation token for all actors + cancel_token.cancel(); + + // Dropping channels signals any dangling non-tokio listeners drop(ingestor_cmd_tx_for_seeds); server_handle.abort(); // Axum server doesn't stop on channel close diff --git a/crates/bot/src/server.rs b/crates/bot/src/server.rs index 9ef75d5..0d2b419 100644 --- a/crates/bot/src/server.rs +++ b/crates/bot/src/server.rs @@ -71,14 +71,23 @@ async fn handle_ws_connection(mut socket: WebSocket, state: AppState) { msg = rx.recv() => { match msg { Ok(json_str) => { - let wrapped = serde_json::json!({ - "type": "EXECUTION_REPORT", - "data": serde_json::from_str::(&json_str).unwrap_or_default(), - }); - if let Ok(wrapped_str) = serde_json::to_string(&wrapped) - && socket.send(Message::Text(wrapped_str.into())).await.is_err() - { - break; // Client disconnected + match serde_json::from_str::(&json_str) { + Ok(parsed) => { + let wrapped = serde_json::json!({ + "type": "EXECUTION_REPORT", + "data": parsed, + }); + if let Ok(wrapped_str) = serde_json::to_string(&wrapped) + && socket.send(Message::Text(wrapped_str.into())).await.is_err() + { + break; // Client disconnected + } + } + Err(e) => { + tracing::error!("Failed to parse broadcast JSON (closing connection): {}", e); + let _ = socket.send(Message::Close(None)).await; + break; + } } } Err(broadcast::error::RecvError::Lagged(n)) => { diff --git a/crates/core/src/domain.rs b/crates/core/src/domain.rs index b304be3..55d5690 100644 --- a/crates/core/src/domain.rs +++ b/crates/core/src/domain.rs @@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize}; /// The type of arbitrage strategy that generated a signal. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] pub enum StrategyType { Implication, Partition, From 7c3b6193160a5626f87cdfafe20e15d21b469cf3 Mon Sep 17 00:00:00 2001 From: Harikeshav Rameshkumar Date: Fri, 27 Feb 2026 21:12:06 -0500 Subject: [PATCH 5/5] feat(server): enhance CORS configuration, optimize status response, and add endpoint tests - **CORS Updates:** - Improved CORS initialization to dynamically parse `VITE_ALLOWED_ORIGINS`. - Added fallback to `http://localhost:5173` for development safety when parsing fails or the environment variable is missing. - **Performance Improvements:** - Optimized status JSON generation by minimizing shared ledger lock contention. - Replaced inline locking with quick data snapshots to allow concurrent reads. - **Response Refinements:** - Standardized `mode` field in `/api/status` response to return stable string values (`LIVE`, `DEMO`) instead of `Debug` representations. - **Testing:** - Introduced tests for `/api/health` and `/api/status` endpoints: - Verified correctness of response payloads and defaults. - Ensured stable handling of uninitialized states. - **Ledger Logic:** - Added safeguard to `Ledger::record_trade` to prevent handling history updates when `max_history` is zero. - **Dependencies:** - Updated `Cargo.toml` to include `tower` v0.5.3. --- crates/bot/Cargo.toml | 1 + crates/bot/src/ledger.rs | 4 ++ crates/bot/src/server.rs | 141 ++++++++++++++++++++++++++++++++++----- 3 files changed, 129 insertions(+), 17 deletions(-) diff --git a/crates/bot/Cargo.toml b/crates/bot/Cargo.toml index 5b1a1a2..c32fed5 100644 --- a/crates/bot/Cargo.toml +++ b/crates/bot/Cargo.toml @@ -21,3 +21,4 @@ tokio-util = { version = "0.7.18", features = ["rt"] } [dev-dependencies] rust_decimal_macros = "1.40.0" +tower = "0.5.3" diff --git a/crates/bot/src/ledger.rs b/crates/bot/src/ledger.rs index f5c46bb..4e97aee 100644 --- a/crates/bot/src/ledger.rs +++ b/crates/bot/src/ledger.rs @@ -49,6 +49,10 @@ impl Ledger { } } + if self.max_history == 0 { + return; + } + if self.history.len() >= self.max_history { self.history.pop_back(); } diff --git a/crates/bot/src/server.rs b/crates/bot/src/server.rs index 0d2b419..fe99456 100644 --- a/crates/bot/src/server.rs +++ b/crates/bot/src/server.rs @@ -6,8 +6,8 @@ use axum::response::{IntoResponse, Json}; use axum::routing::get; use std::sync::Arc; use tokio::sync::{RwLock, broadcast}; -use tower_http::cors::CorsLayer; -use tracing::{error, info}; +use tower_http::cors::{Any, CorsLayer}; +use tracing::{error, info, warn}; /// Shared state accessible by all Axum handlers. #[derive(Clone)] @@ -26,13 +26,30 @@ pub async fn start_server(state: AppState, port: u16) -> anyhow::Result<()> { .route("/api/status", get(status_handler)) .layer( std::env::var("VITE_ALLOWED_ORIGINS") + .ok() + .filter(|s| !s.is_empty()) .map(|origin| { - CorsLayer::new() - .allow_origin(origin.parse::().unwrap()) - .allow_methods(tower_http::cors::Any) - .allow_headers(tower_http::cors::Any) + match origin.parse::() { + Ok(header) => CorsLayer::new() + .allow_origin(header) + .allow_methods(Any) + .allow_headers(Any), + Err(e) => { + warn!("Invalid VITE_ALLOWED_ORIGINS value '{}': {}. Falling back to safe localhost default.", origin, e); + CorsLayer::new() + .allow_origin("http://localhost:5173".parse::().unwrap()) + .allow_methods(Any) + .allow_headers(Any) + } + } }) - .unwrap_or_else(|_| CorsLayer::permissive()), // Use permissive if unspecified, but can be locked down + .unwrap_or_else(|| { + // Safe default if no env var is provided: only allow standard local vite dev server + CorsLayer::new() + .allow_origin("http://localhost:5173".parse::().unwrap()) + .allow_methods(Any) + .allow_headers(Any) + }), ) .with_state(state); @@ -127,12 +144,25 @@ async fn status_handler(State(state): State) -> impl IntoResponse { /// Build the status JSON payload from shared state. async fn build_status_json(state: &AppState) -> anyhow::Result { - let ledger = state.ledger.read().await; + // Minimize lock contention by cloning the needed data quickly and dropping the lock let uptime = state.start_time.elapsed().as_secs(); + let (positions_snap, history_snap, cumulative_pnl, total_signals) = { + let ledger = state.ledger.read().await; + ( + ledger.get_positions(), + // Map the history explicitly so we own the data before dropping the lock + ledger + .get_recent_history(50) + .into_iter() + .cloned() + .collect::>(), + ledger.cumulative_pnl(), + ledger.total_signals(), + ) + }; - let positions: Vec = ledger - .get_positions() - .iter() + let positions: Vec = positions_snap + .into_iter() .map(|p| { serde_json::json!({ "asset_id": p.asset_id.to_string(), @@ -144,9 +174,8 @@ async fn build_status_json(state: &AppState) -> anyhow::Result = ledger - .get_recent_history(50) - .iter() + let recent_trades: Vec = history_snap + .into_iter() .map(|r| { serde_json::json!({ "strategy": r.signal.strategy.to_string(), @@ -168,11 +197,89 @@ async fn build_status_json(state: &AppState) -> anyhow::Result "LIVE", + arbos_core::domain::ExecutionMode::Demo => "DEMO", + }, + "cumulative_pnl": cumulative_pnl.to_string(), + "signals_executed": total_signals, "uptime_secs": uptime, "positions": positions, "recent_trades": recent_trades, })) } + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::Body; + use axum::http::{Request, StatusCode}; + use tower::ServiceExt; + + fn build_test_router() -> Router { + let state = AppState { + ledger: Arc::new(RwLock::new(Ledger::new(50))), + broadcast_tx: broadcast::channel(10).0, + start_time: std::time::Instant::now(), + mode: arbos_core::domain::ExecutionMode::Demo, + }; + + Router::new() + .route("/api/health", get(health_handler)) + .route("/api/status", get(status_handler)) + .with_state(state) + } + + #[tokio::test] + async fn test_health_endpoint() { + let app = build_test_router(); + + let response = app + .oneshot( + Request::builder() + .uri("/api/health") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + assert_eq!(&body[..], b"OK"); + } + + #[tokio::test] + async fn test_status_endpoint() { + let app = build_test_router(); + + let response = app + .oneshot( + Request::builder() + .uri("/api/status") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let payload: serde_json::Value = + serde_json::from_slice(&body).expect("Invalid JSON returned by /api/status"); + + assert_eq!(payload["mode"], "DEMO"); + assert!(payload["uptime_secs"].is_u64()); + assert_eq!(payload["cumulative_pnl"], "0"); // Starts at zero + assert_eq!(payload["signals_executed"], 0); + assert!(payload["positions"].is_array()); + assert!(payload["recent_trades"].is_array()); + } +}