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 #[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 #[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}