1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4use serde_json::{Value, json};
5
6use crate::event::{Event, EventKind};
7use crate::message::{Content, Message, Role};
8use crate::model::Usage;
9
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11#[serde(tag = "status", rename_all = "snake_case")]
12pub enum Status {
13 Running,
14 AwaitingApproval,
15 Completed { output: String },
16 Failed { error: String },
17 BudgetExceeded { reason: String },
18}
19
20impl Status {
21 pub fn is_terminal(&self) -> bool {
22 matches!(self, Status::Completed { .. } | Status::Failed { .. } | Status::BudgetExceeded { .. })
23 }
24}
25
26#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
28pub struct PendingCall {
29 pub id: String,
30 pub name: String,
31 pub input: Value,
32 pub awaiting_approval: bool,
33}
34
35#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
40pub struct RunState {
41 pub run_id: String,
42 pub agent: String,
43 pub status: Status,
44 pub messages: Vec<Message>,
45 pub steps: u32,
46 pub usage: Usage,
47 pub cost_usd: f64,
48 pub pending: Vec<PendingCall>,
49 pub approvals: HashMap<String, bool>,
50 pub tool_results: Vec<Content>,
52 #[serde(default)]
54 pub structured_output: Option<Value>,
55 #[serde(default)]
57 pub finish_rejections: u32,
58 pub next_seq: u64,
59}
60
61impl RunState {
62 pub fn from_events(events: &[Event]) -> Option<Self> {
65 let first = events.first()?;
66 let mut state = RunState {
67 run_id: first.run_id.clone(),
68 agent: String::new(),
69 status: Status::Running,
70 messages: Vec::new(),
71 steps: 0,
72 usage: Usage::default(),
73 cost_usd: 0.0,
74 pending: Vec::new(),
75 approvals: HashMap::new(),
76 tool_results: Vec::new(),
77 structured_output: None,
78 finish_rejections: 0,
79 next_seq: 0,
80 };
81 for event in events {
82 state.apply(event);
83 }
84 Some(state)
85 }
86
87 pub fn apply(&mut self, event: &Event) {
88 debug_assert_eq!(event.seq, self.next_seq, "events must be applied in order");
89 self.next_seq = event.seq + 1;
90
91 match &event.kind {
92 EventKind::RunStarted { agent, input } => {
93 self.agent = agent.clone();
94 self.messages = vec![input.clone()];
95 self.status = Status::Running;
96 }
97 EventKind::ModelCalled { response, cost_usd, .. } => {
98 self.steps += 1;
99 self.usage.add(&response.total_usage());
100 self.cost_usd += cost_usd.unwrap_or(0.0);
101 self.messages.push(Message::assistant(response.content.clone()));
102 self.pending = response
103 .content
104 .iter()
105 .filter_map(|c| match c {
106 Content::ToolUse { id, name, input } => Some(PendingCall {
107 id: id.clone(),
108 name: name.clone(),
109 input: input.clone(),
110 awaiting_approval: false,
111 }),
112 _ => None,
113 })
114 .collect();
115 }
116 EventKind::ModelFailed { usage, cost_usd, .. } => {
117 self.usage.add(usage);
118 self.cost_usd += cost_usd.unwrap_or(0.0);
119 }
120 EventKind::ToolDenied { call_id, reason, .. } => {
121 self.resolve(call_id, json!({ "error": format!("tool call denied: {reason}") }), true);
122 }
123 EventKind::ApprovalRequested { call_id, .. } => {
124 if let Some(p) = self.pending.iter_mut().find(|p| &p.id == call_id) {
125 p.awaiting_approval = true;
126 }
127 self.status = Status::AwaitingApproval;
128 }
129 EventKind::ApprovalDecided { call_id, approved, .. } => {
130 self.approvals.insert(call_id.clone(), *approved);
131 if let Some(p) = self.pending.iter_mut().find(|p| &p.id == call_id) {
132 p.awaiting_approval = false;
133 }
134 if !self.pending.iter().any(|p| p.awaiting_approval) {
135 self.status = Status::Running;
136 }
137 }
138 EventKind::ToolCompleted { call_id, output, is_error, .. } => {
139 self.resolve(call_id, output.clone(), *is_error);
140 }
141 EventKind::BudgetExceeded { reason } => {
142 self.status = Status::BudgetExceeded { reason: reason.clone() };
143 }
144 EventKind::RunCompleted { output, structured } => {
145 self.status = Status::Completed { output: output.clone() };
146 self.structured_output = structured.clone();
147 }
148 EventKind::FinishRejected { reason, .. } => {
149 self.finish_rejections += 1;
150 self.messages.push(Message::user(format!(
151 "Your final answer was rejected: {reason}\n\
152 Reply again with only a JSON object that matches the required schema."
153 )));
154 }
155 EventKind::RunFailed { error } => {
156 self.status = Status::Failed { error: error.clone() };
157 }
158 }
159 }
160
161 fn resolve(&mut self, call_id: &str, content: Value, is_error: bool) {
162 self.pending.retain(|p| p.id != call_id);
163 self.tool_results.push(Content::ToolResult { tool_use_id: call_id.to_owned(), content, is_error });
164 if self.pending.is_empty() {
165 let results = std::mem::take(&mut self.tool_results);
166 self.messages.push(Message { role: Role::User, content: results });
167 }
168 }
169
170 pub fn awaiting_approval(&self) -> Vec<PendingCall> {
171 self.pending.iter().filter(|p| p.awaiting_approval).cloned().collect()
172 }
173}