Skip to main content

cairn_runtime/
state.rs

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/// A tool call requested by the model that has not produced a result yet.
27#[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/// The full state of a run, derived purely from its events.
36///
37/// There is no other source of truth: live execution, resume after a crash
38/// and time-travel debugging all go through [`RunState::apply`].
39#[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    /// Results gathered for the current batch of tool calls.
51    pub tool_results: Vec<Content>,
52    /// Validated final answer when the agent has a finish schema.
53    #[serde(default)]
54    pub structured_output: Option<Value>,
55    /// How many final answers have been rejected by the finish schema.
56    #[serde(default)]
57    pub finish_rejections: u32,
58    pub next_seq: u64,
59}
60
61impl RunState {
62    /// Rebuild state from a run's events. Returns `None` for an empty slice.
63    /// Pass a prefix (`&events[..=n]`) to see the state at any point in time.
64    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}