1use std::any::type_name;
2use std::collections::{BTreeMap, BTreeSet};
3use std::sync::{Arc, Mutex};
4
5use serde::de::DeserializeOwned;
6use serde::{Serialize, Serializer};
7use serde_json::{json, Value};
8use thiserror::Error;
9
10use crate::agent::Agent;
11use crate::budget::{BudgetExhaustion, BudgetUsageSnapshot};
12use crate::checkpoint::ResumeObservation;
13use crate::config::ResolvedModelConfig;
14use crate::events::RunEvent;
15use crate::run_config::RunConfig;
16use crate::tools::{ToolContext, ToolOrchestrator, ToolRunOptions};
17use crate::types::{AgentResult, AgentStatus, CompletionReason, Message, Metadata, TaskTokenUsage};
18
19#[derive(Debug, Clone, PartialEq, Serialize)]
20pub struct ApprovalSnapshot {
21 pub interruption_id: String,
22 pub tool_name: String,
23 pub tool_call_id: String,
24 pub arguments: BTreeMap<String, Value>,
25 pub message: String,
26 pub cycle_index: Option<u32>,
27 pub approved: bool,
28}
29
30#[derive(Debug, Error)]
31pub enum FinalOutputError {
32 #[error("run result for agent `{agent_name}` has no final output (status: {status:?})")]
33 Missing {
34 agent_name: String,
35 status: AgentStatus,
36 },
37 #[error("failed to deserialize final output for agent `{agent_name}` as `{target}`: {source}")]
38 Deserialize {
39 agent_name: String,
40 target: &'static str,
41 #[source]
42 source: serde_json::Error,
43 },
44}
45
46#[derive(Clone)]
47pub struct RunResult {
48 agent_name: String,
49 run_id: String,
50 trace_id: String,
51 input: String,
52 new_items: Vec<Message>,
53 events: Vec<RunEvent>,
54 token_usage: TaskTokenUsage,
55 metadata: Metadata,
56 result: AgentResult,
57 resolved: Option<ResolvedModelConfig>,
58 resume_context: Option<RunResumeContext>,
59}
60
61impl RunResult {
62 pub fn new(
63 agent_name: impl Into<String>,
64 result: AgentResult,
65 resolved: ResolvedModelConfig,
66 ) -> Self {
67 let token_usage = result.token_usage.clone();
68 Self {
69 agent_name: agent_name.into(),
70 run_id: String::new(),
71 trace_id: String::new(),
72 input: String::new(),
73 new_items: Vec::new(),
74 events: Vec::new(),
75 token_usage,
76 metadata: Metadata::new(),
77 result,
78 resolved: Some(resolved),
79 resume_context: None,
80 }
81 }
82
83 pub(crate) fn without_resolved_model(
84 agent_name: impl Into<String>,
85 result: AgentResult,
86 ) -> Self {
87 let token_usage = result.token_usage.clone();
88 Self {
89 agent_name: agent_name.into(),
90 run_id: String::new(),
91 trace_id: String::new(),
92 input: String::new(),
93 new_items: Vec::new(),
94 events: Vec::new(),
95 token_usage,
96 metadata: Metadata::new(),
97 result,
98 resolved: None,
99 resume_context: None,
100 }
101 }
102
103 pub fn agent_name(&self) -> &str {
104 &self.agent_name
105 }
106
107 pub fn run_id(&self) -> &str {
108 &self.run_id
109 }
110
111 pub fn trace_id(&self) -> &str {
112 &self.trace_id
113 }
114
115 pub fn input(&self) -> &str {
116 &self.input
117 }
118
119 pub fn new_items(&self) -> &[Message] {
120 &self.new_items
121 }
122
123 pub fn events(&self) -> &[RunEvent] {
124 &self.events
125 }
126
127 pub fn token_usage(&self) -> &TaskTokenUsage {
128 &self.token_usage
129 }
130
131 pub fn metadata(&self) -> &Metadata {
132 &self.metadata
133 }
134
135 pub fn status(&self) -> AgentStatus {
136 self.result.status
137 }
138
139 pub fn completion_reason(&self) -> Option<CompletionReason> {
140 self.result.completion_reason
141 }
142
143 pub fn completion_tool_name(&self) -> Option<&str> {
144 self.result.completion_tool_name.as_deref()
145 }
146
147 pub fn partial_output(&self) -> Option<&str> {
148 self.result.partial_output.as_deref()
149 }
150
151 pub fn budget_usage(&self) -> Option<&BudgetUsageSnapshot> {
152 self.result.budget_usage.as_ref()
153 }
154
155 pub fn budget_exhaustion(&self) -> Option<&BudgetExhaustion> {
156 self.result.budget_exhaustion.as_ref()
157 }
158
159 pub fn checkpoint_key(&self) -> Option<&str> {
160 self.result.checkpoint_key.as_deref()
161 }
162
163 pub fn resume_observation(&self) -> Option<&ResumeObservation> {
164 self.result.resume_observation.as_ref()
165 }
166
167 pub fn final_output(&self) -> Option<&str> {
168 self.result
169 .final_answer
170 .as_deref()
171 .or(self.result.wait_reason.as_deref())
172 .or(self.result.error.as_deref())
173 }
174
175 pub fn deserialize<T>(&self) -> Result<T, FinalOutputError>
176 where
177 T: DeserializeOwned,
178 {
179 let output = self
180 .final_output()
181 .ok_or_else(|| FinalOutputError::Missing {
182 agent_name: self.agent_name.clone(),
183 status: self.status(),
184 })?;
185 serde_json::from_str(output).map_err(|source| FinalOutputError::Deserialize {
186 agent_name: self.agent_name.clone(),
187 target: type_name::<T>(),
188 source,
189 })
190 }
191
192 pub fn result(&self) -> &AgentResult {
193 &self.result
194 }
195
196 pub fn approvals(&self) -> Vec<ApprovalSnapshot> {
197 approval_snapshots(&self.result, &BTreeSet::new())
198 }
199
200 pub fn approval_snapshot(&self) -> Vec<ApprovalSnapshot> {
201 self.approvals()
202 }
203
204 pub fn resolved_model(&self) -> Option<&ResolvedModelConfig> {
205 self.resolved.as_ref()
206 }
207
208 pub fn to_value(&self) -> Value {
209 let result = self.result.to_dict();
210 let status = result.get("status").cloned().unwrap_or(Value::Null);
211 let token_usage = result
212 .get("token_usage")
213 .cloned()
214 .unwrap_or_else(|| json!({}));
215 json!({
216 "input": self.input,
217 "new_items": self.new_items.iter().map(Message::to_dict).collect::<Vec<_>>(),
218 "final_output": self.final_output(),
219 "status": status,
220 "events": self.events,
221 "token_usage": token_usage,
222 "trace_id": self.trace_id,
223 "run_id": self.run_id,
224 "metadata": self.metadata,
225 "agent_name": self.agent_name,
226 "completion_reason": self.result.completion_reason,
227 "completion_tool_name": self.result.completion_tool_name,
228 "partial_output": self.result.partial_output,
229 "budget_usage": self.result.budget_usage,
230 "budget_exhaustion": self.result.budget_exhaustion,
231 "checkpoint_key": self.result.checkpoint_key,
232 "resume_observation": self.result.resume_observation,
233 "resolved_model": self.resolved.as_ref().map(resolved_model_public_value),
234 })
235 }
236
237 pub fn to_dict(&self) -> Value {
238 self.to_value()
239 }
240
241 pub fn with_input(mut self, input: impl Into<String>) -> Self {
242 self.input = input.into();
243 self
244 }
245
246 pub fn with_new_items(mut self, new_items: Vec<Message>) -> Self {
247 self.new_items = new_items;
248 self
249 }
250
251 pub fn with_events(mut self, events: Vec<RunEvent>) -> Self {
252 self.events = events;
253 self
254 }
255
256 pub fn with_token_usage(mut self, token_usage: TaskTokenUsage) -> Self {
257 self.token_usage = token_usage;
258 self
259 }
260
261 pub fn with_metadata(mut self, metadata: Metadata) -> Self {
262 self.metadata = metadata;
263 self
264 }
265
266 pub fn into_state(self) -> Result<RunState, String> {
267 RunState::from_result(self)
268 }
269
270 pub(crate) fn with_resume_context(mut self, context: RunResumeContext) -> Self {
271 self.resume_context = Some(context);
272 self
273 }
274
275 pub(crate) fn with_ids(
276 mut self,
277 run_id: impl Into<String>,
278 trace_id: impl Into<String>,
279 ) -> Self {
280 self.run_id = run_id.into();
281 self.trace_id = trace_id.into();
282 self
283 }
284
285 pub(crate) fn resume_context(&self) -> Option<&RunResumeContext> {
286 self.resume_context.as_ref()
287 }
288}
289
290impl Serialize for RunResult {
291 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
292 where
293 S: Serializer,
294 {
295 self.to_value().serialize(serializer)
296 }
297}
298
299#[derive(Clone)]
300pub struct RunState {
301 result: RunResult,
302 approved_interruption_ids: Vec<String>,
303 approval_consumption: Arc<Mutex<BTreeSet<String>>>,
304}
305
306impl RunState {
307 pub fn from_result(result: RunResult) -> Result<Self, String> {
308 if result.status() != AgentStatus::WaitUser {
309 return Err("only interrupted runs can be converted into RunState".to_string());
310 }
311 Ok(Self {
312 result,
313 approved_interruption_ids: Vec::new(),
314 approval_consumption: Arc::new(Mutex::new(BTreeSet::new())),
315 })
316 }
317
318 pub fn result(&self) -> &RunResult {
319 &self.result
320 }
321
322 pub fn approve(&mut self, interruption_id: &str) -> Result<(), String> {
323 if !self
324 .pending_approval_ids()
325 .iter()
326 .any(|id| id == interruption_id)
327 {
328 return Err(format!("unknown approval interruption: {interruption_id}"));
329 }
330 if !self
331 .approved_interruption_ids
332 .iter()
333 .any(|id| id == interruption_id)
334 {
335 self.approved_interruption_ids
336 .push(interruption_id.to_string());
337 }
338 Ok(())
339 }
340
341 pub fn approved_interruption_ids(&self) -> &[String] {
342 &self.approved_interruption_ids
343 }
344
345 pub fn approvals(&self) -> Vec<ApprovalSnapshot> {
346 let approved = self
347 .approved_interruption_ids
348 .iter()
349 .cloned()
350 .collect::<BTreeSet<_>>();
351 approval_snapshots(self.result.result(), &approved)
352 }
353
354 pub fn approval_snapshot(&self) -> Vec<ApprovalSnapshot> {
355 self.approvals()
356 }
357
358 pub fn pending_approval_ids(&self) -> Vec<String> {
359 self.result
360 .approvals()
361 .into_iter()
362 .map(|snapshot| snapshot.interruption_id)
363 .collect()
364 }
365
366 pub(crate) fn into_inner(self) -> (RunResult, Vec<String>, Arc<Mutex<BTreeSet<String>>>) {
367 (
368 self.result,
369 self.approved_interruption_ids,
370 self.approval_consumption,
371 )
372 }
373}
374
375fn approval_snapshots(
376 result: &AgentResult,
377 approved_interruption_ids: &BTreeSet<String>,
378) -> Vec<ApprovalSnapshot> {
379 let mut seen = BTreeSet::new();
380 let mut snapshots = Vec::new();
381 for cycle in &result.cycles {
382 for tool_result in &cycle.tool_results {
383 let metadata = &tool_result.metadata;
384 let interruption_id = metadata
385 .get("approval_interruption_id")
386 .or_else(|| metadata.get("request_id"))
387 .and_then(Value::as_str)
388 .map(str::trim)
389 .filter(|value| !value.is_empty());
390 let Some(interruption_id) = interruption_id else {
391 continue;
392 };
393 let approval_requested = metadata
394 .get("mode")
395 .and_then(Value::as_str)
396 .is_some_and(|mode| mode == "approval_requested")
397 || metadata
398 .get("approval_required")
399 .and_then(Value::as_bool)
400 .unwrap_or(false);
401 if !approval_requested || !seen.insert(interruption_id.to_string()) {
402 continue;
403 }
404 let arguments = metadata
405 .get("arguments")
406 .and_then(Value::as_object)
407 .map(|arguments| {
408 arguments
409 .iter()
410 .map(|(key, value)| (key.clone(), value.clone()))
411 .collect()
412 })
413 .unwrap_or_default();
414 snapshots.push(ApprovalSnapshot {
415 interruption_id: interruption_id.to_string(),
416 tool_name: metadata
417 .get("tool_name")
418 .and_then(Value::as_str)
419 .unwrap_or_default()
420 .to_string(),
421 tool_call_id: tool_result.tool_call_id.clone(),
422 arguments,
423 message: metadata
424 .get("message")
425 .and_then(Value::as_str)
426 .unwrap_or(&tool_result.content)
427 .to_string(),
428 cycle_index: Some(cycle.index),
429 approved: approved_interruption_ids.contains(interruption_id),
430 });
431 }
432 }
433 snapshots
434}
435
436fn resolved_model_public_value(resolved: &ResolvedModelConfig) -> Value {
437 json!({
438 "backend": resolved.backend,
439 "requested_model": resolved.requested_model,
440 "selected_model": resolved.selected_model,
441 "model_id": resolved.model_id,
442 "endpoint": resolved
443 .endpoint_options
444 .first()
445 .map(|option| option.endpoint.endpoint_id.as_str()),
446 })
447}
448
449#[derive(Clone)]
450pub(crate) struct RunResumeContext {
451 pub agent: Agent,
452 pub input: crate::runner::NormalizedInput,
453 pub config: RunConfig,
454 pub runner: crate::runner::Runner,
455 pub pending_tool_approval: Option<PendingToolApproval>,
456}
457
458#[derive(Clone)]
459pub(crate) struct PendingToolApproval {
460 pub interruption_id: String,
461 pub call: crate::types::ToolCall,
462 pub cycle_index: u32,
463 pub context: ToolContext,
464 pub options: ToolRunOptions,
465 pub orchestrator: ToolOrchestrator,
466 pub task: crate::types::AgentTask,
467 pub hook_manager: crate::runtime::RuntimeHookManager,
468}