Skip to main content

schwab_cli/agent/
state.rs

1use std::collections::HashMap;
2use std::fs;
3use std::path::Path;
4
5use anyhow::Result;
6use chrono::{DateTime, NaiveDate, Utc};
7use serde::{Deserialize, Serialize};
8use serde_json::{json, Value};
9
10use crate::options::types::StrategyKind;
11
12#[derive(Debug, Clone, Serialize, Deserialize, Default)]
13pub struct AgentState {
14    pub agent_id: String,
15    pub last_tick: Option<DateTime<Utc>>,
16    pub trades_today: u32,
17    pub trades_day: Option<NaiveDate>,
18    pub open_positions: HashMap<String, TrackedPosition>,
19    pub last_actions: Vec<AgentAction>,
20    #[serde(default)]
21    pub pending_order_ids: Vec<String>,
22    #[serde(default)]
23    pub pending_orders: Vec<PendingOrder>,
24    #[serde(default)]
25    pub tick_count: u64,
26    #[serde(default)]
27    pub last_llm_review_tick: Option<u64>,
28    #[serde(default)]
29    pub llm_review_count: u64,
30    #[serde(default)]
31    pub last_llm_summary: Option<Value>,
32    #[serde(default)]
33    pub last_session: Option<String>,
34    #[serde(default)]
35    pub regular_tick_count: u64,
36    #[serde(default)]
37    pub last_overnight_digest_at: Option<DateTime<Utc>>,
38    #[serde(default)]
39    pub open_playbook: Option<Value>,
40    /// Last EQO regular-session open flag from agent tick (Schwab hours).
41    #[serde(default)]
42    pub last_market_open: Option<bool>,
43    #[serde(default)]
44    pub last_auth_reminder_level: Option<String>,
45    #[serde(default)]
46    pub last_auth_reminder_at: Option<DateTime<Utc>>,
47    /// Last LLM review pushed to Telegram (for digest dedup).
48    #[serde(default)]
49    pub last_telegram_llm_at: Option<DateTime<Utc>>,
50    #[serde(default)]
51    pub last_telegram_llm_digest_key: Option<String>,
52    /// Paper-trading ledger when running with --simulate (separate state file).
53    #[serde(default, skip_serializing_if = "Option::is_none")]
54    pub sim: Option<crate::agent::sim::SimLedger>,
55}
56
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct TrackedPosition {
59    pub position_id: String,
60    pub account_hash: String,
61    pub underlying: String,
62    pub expiry: String,
63    pub strategy: String,
64    pub opened_at: DateTime<Utc>,
65    pub entry_credit: Option<f64>,
66    pub max_loss_usd: f64,
67    /// Spread quantity (each leg at Schwab should match this count).
68    #[serde(default = "default_one")]
69    pub contracts: u32,
70    /// Strategy params for sim marks / vertical reconstruction.
71    #[serde(default, skip_serializing_if = "Option::is_none")]
72    pub entry_params: Option<Value>,
73}
74
75#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
76#[serde(rename_all = "snake_case")]
77pub enum PendingOrderAction {
78    Entry,
79    Exit,
80}
81
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct PendingOrder {
84    pub order_id: String,
85    pub account_hash: String,
86    pub action: PendingOrderAction,
87    pub position_id: String,
88    pub reserved_risk_usd: f64,
89    pub submitted_at: DateTime<Utc>,
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub last_status: Option<String>,
92    #[serde(default, skip_serializing_if = "Option::is_none")]
93    pub detail: Option<Value>,
94}
95
96fn default_one() -> u32 {
97    1
98}
99
100#[derive(Debug, Clone, Serialize, Deserialize)]
101pub struct AgentAction {
102    pub at: DateTime<Utc>,
103    pub action: String,
104    pub detail: Value,
105}
106
107pub fn load_state(path: &Path) -> Result<AgentState> {
108    if !path.exists() {
109        return Ok(AgentState::default());
110    }
111    let content = fs::read_to_string(path)?;
112    Ok(serde_json::from_str(&content)?)
113}
114
115pub fn save_state(path: &Path, state: &AgentState) -> Result<()> {
116    if let Some(parent) = path.parent() {
117        fs::create_dir_all(parent)?;
118    }
119    let content = serde_json::to_string_pretty(state)?;
120    fs::write(path, content)?;
121    Ok(())
122}
123
124impl AgentState {
125    pub fn reset_daily_if_needed(&mut self, today: NaiveDate) {
126        if self.trades_day != Some(today) {
127            self.trades_today = 0;
128            self.trades_day = Some(today);
129        }
130    }
131
132    pub fn record_action(&mut self, action: &str, detail: Value) {
133        self.last_actions.push(AgentAction {
134            at: Utc::now(),
135            action: action.to_string(),
136            detail,
137        });
138        if self.last_actions.len() > 100 {
139            let drain = self.last_actions.len() - 100;
140            self.last_actions.drain(0..drain);
141        }
142    }
143
144    pub fn count_open_for_strategy(&self, account_hash: &str, strategy: StrategyKind) -> u32 {
145        self.open_positions
146            .values()
147            .filter(|p| p.account_hash == account_hash && p.strategy == strategy.as_str())
148            .count() as u32
149    }
150
151    pub fn pending_entry_count(&self) -> u32 {
152        self.pending_orders
153            .iter()
154            .filter(|p| p.action == PendingOrderAction::Entry)
155            .count() as u32
156    }
157
158    pub fn pending_count(&self) -> usize {
159        self.pending_orders.len().max(self.pending_order_ids.len())
160    }
161
162    pub fn trades_capacity_used(&self) -> u32 {
163        self.trades_today.saturating_add(self.pending_entry_count())
164    }
165
166    pub fn open_risk_usd(&self) -> f64 {
167        self.open_positions
168            .values()
169            .map(|p| p.max_loss_usd.max(0.0))
170            .sum()
171    }
172
173    pub fn pending_entry_risk_usd(&self) -> f64 {
174        self.pending_orders
175            .iter()
176            .filter(|p| p.action == PendingOrderAction::Entry)
177            .map(|p| p.reserved_risk_usd.max(0.0))
178            .sum()
179    }
180
181    pub fn reserved_risk_usd(&self) -> f64 {
182        self.open_risk_usd() + self.pending_entry_risk_usd()
183    }
184
185    pub fn has_pending_position(&self, position_id: &str) -> bool {
186        self.pending_orders
187            .iter()
188            .any(|p| p.position_id == position_id)
189    }
190
191    pub fn add_pending_order(&mut self, pending: PendingOrder) {
192        if !self
193            .pending_order_ids
194            .iter()
195            .any(|id| id == &pending.order_id)
196        {
197            self.pending_order_ids.push(pending.order_id.clone());
198        }
199        if let Some(existing) = self
200            .pending_orders
201            .iter_mut()
202            .find(|p| p.order_id == pending.order_id)
203        {
204            *existing = pending;
205        } else {
206            self.pending_orders.push(pending);
207        }
208    }
209
210    pub fn remove_pending_order(&mut self, order_id: &str) -> Option<PendingOrder> {
211        self.pending_order_ids.retain(|id| id != order_id);
212        let idx = self
213            .pending_orders
214            .iter()
215            .position(|p| p.order_id == order_id)?;
216        Some(self.pending_orders.remove(idx))
217    }
218
219    pub fn clear_legacy_pending_ids(&mut self) {
220        let structured: std::collections::HashSet<&str> = self
221            .pending_orders
222            .iter()
223            .map(|p| p.order_id.as_str())
224            .collect();
225        self.pending_order_ids
226            .retain(|id| structured.contains(id.as_str()));
227    }
228
229    pub fn total_contracts(&self) -> u32 {
230        self.open_positions
231            .values()
232            .map(|p| p.contracts.max(1))
233            .sum()
234    }
235}
236
237pub fn state_summary(state: &AgentState) -> Value {
238    json!({
239        "agent_id": state.agent_id,
240        "last_tick": state.last_tick,
241        "trades_today": state.trades_today,
242        "open_positions": state.open_positions.len(),
243        "tick_count": state.tick_count,
244        "last_llm_review_tick": state.last_llm_review_tick,
245        "last_llm_summary": state.last_llm_summary,
246        "last_session": state.last_session,
247        "regular_tick_count": state.regular_tick_count,
248        "last_overnight_digest_at": state.last_overnight_digest_at,
249        "open_playbook": state.open_playbook,
250        "last_market_open": state.last_market_open,
251        "last_auth_reminder_at": state.last_auth_reminder_at,
252        "pending_orders": state.pending_count(),
253        "reserved_risk_usd": state.reserved_risk_usd(),
254        "pending_orders_detail": state.pending_orders,
255        "recent_actions": state.last_actions.iter().rev().take(10).collect::<Vec<_>>(),
256    })
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    #[test]
264    fn reserved_risk_includes_pending_entries_only() {
265        let mut state = AgentState::default();
266        state.open_positions.insert(
267            "pos".into(),
268            TrackedPosition {
269                position_id: "pos".into(),
270                account_hash: "acct".into(),
271                underlying: "IWM".into(),
272                expiry: "2026-07-31".into(),
273                strategy: "vertical".into(),
274                opened_at: Utc::now(),
275                entry_credit: Some(0.25),
276                max_loss_usd: 175.0,
277                contracts: 1,
278                entry_params: None,
279            },
280        );
281        state.add_pending_order(PendingOrder {
282            order_id: "entry-1".into(),
283            account_hash: "acct".into(),
284            action: PendingOrderAction::Entry,
285            position_id: "pending-entry".into(),
286            reserved_risk_usd: 170.0,
287            submitted_at: Utc::now(),
288            last_status: Some("WORKING".into()),
289            detail: None,
290        });
291        state.add_pending_order(PendingOrder {
292            order_id: "exit-1".into(),
293            account_hash: "acct".into(),
294            action: PendingOrderAction::Exit,
295            position_id: "pos".into(),
296            reserved_risk_usd: 0.0,
297            submitted_at: Utc::now(),
298            last_status: Some("WORKING".into()),
299            detail: None,
300        });
301
302        assert_eq!(state.pending_entry_count(), 1);
303        assert!((state.reserved_risk_usd() - 345.0).abs() < 0.01);
304    }
305}