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