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 error_code(&self) -> Option<&str> {
176 self.result.error_code.as_deref()
177 }
178
179 pub fn deserialize<T>(&self) -> Result<T, FinalOutputError>
180 where
181 T: DeserializeOwned,
182 {
183 let output = self
184 .final_output()
185 .ok_or_else(|| FinalOutputError::Missing {
186 agent_name: self.agent_name.clone(),
187 status: self.status(),
188 })?;
189 serde_json::from_str(output).map_err(|source| FinalOutputError::Deserialize {
190 agent_name: self.agent_name.clone(),
191 target: type_name::<T>(),
192 source,
193 })
194 }
195
196 pub fn result(&self) -> &AgentResult {
197 &self.result
198 }
199
200 pub fn approvals(&self) -> Vec<ApprovalSnapshot> {
201 approval_snapshots(&self.result, &BTreeSet::new())
202 }
203
204 pub fn approval_snapshot(&self) -> Vec<ApprovalSnapshot> {
205 self.approvals()
206 }
207
208 pub fn resolved_model(&self) -> Option<&ResolvedModelConfig> {
209 self.resolved.as_ref()
210 }
211
212 pub fn to_value(&self) -> Value {
213 let result = self.result.to_dict();
214 let status = result.get("status").cloned().unwrap_or(Value::Null);
215 let token_usage = result
216 .get("token_usage")
217 .cloned()
218 .unwrap_or_else(|| json!({}));
219 let mut payload = json!({
220 "input": self.input,
221 "new_items": self.new_items.iter().map(Message::to_dict).collect::<Vec<_>>(),
222 "final_output": self.final_output(),
223 "status": status,
224 "events": self.events,
225 "token_usage": token_usage,
226 "trace_id": self.trace_id,
227 "run_id": self.run_id,
228 "metadata": self.metadata,
229 "agent_name": self.agent_name,
230 "completion_reason": self.result.completion_reason,
231 "completion_tool_name": self.result.completion_tool_name,
232 "partial_output": self.result.partial_output,
233 "budget_usage": self.result.budget_usage,
234 "budget_exhaustion": self.result.budget_exhaustion,
235 "checkpoint_key": self.result.checkpoint_key,
236 "resume_observation": self.result.resume_observation,
237 "resolved_model": self.resolved.as_ref().map(resolved_model_public_value),
238 });
239 if let Some(error_code) = self.error_code() {
240 payload["error_code"] = Value::String(error_code.to_string());
241 }
242 payload
243 }
244
245 pub fn to_dict(&self) -> Value {
246 self.to_value()
247 }
248
249 pub fn with_input(mut self, input: impl Into<String>) -> Self {
250 self.input = input.into();
251 self
252 }
253
254 pub fn with_new_items(mut self, new_items: Vec<Message>) -> Self {
255 self.new_items = new_items;
256 self
257 }
258
259 pub fn with_events(mut self, events: Vec<RunEvent>) -> Self {
260 self.events = events;
261 self
262 }
263
264 pub fn with_token_usage(mut self, token_usage: TaskTokenUsage) -> Self {
265 self.token_usage = token_usage;
266 self
267 }
268
269 pub fn with_metadata(mut self, metadata: Metadata) -> Self {
270 self.metadata = metadata;
271 self
272 }
273
274 pub fn into_state(self) -> Result<RunState, String> {
275 RunState::from_result(self)
276 }
277
278 pub(crate) fn with_resume_context(mut self, context: RunResumeContext) -> Self {
279 self.resume_context = Some(context);
280 self
281 }
282
283 pub(crate) fn with_ids(
284 mut self,
285 run_id: impl Into<String>,
286 trace_id: impl Into<String>,
287 ) -> Self {
288 self.run_id = run_id.into();
289 self.trace_id = trace_id.into();
290 self
291 }
292
293 pub(crate) fn resume_context(&self) -> Option<&RunResumeContext> {
294 self.resume_context.as_ref()
295 }
296}
297
298impl Serialize for RunResult {
299 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
300 where
301 S: Serializer,
302 {
303 self.to_value().serialize(serializer)
304 }
305}
306
307#[derive(Clone)]
308pub struct RunState {
309 result: RunResult,
310 approved_interruption_ids: Vec<String>,
311 approval_consumption: Arc<Mutex<BTreeSet<String>>>,
312}
313
314impl RunState {
315 pub fn from_result(result: RunResult) -> Result<Self, String> {
316 if result.status() != AgentStatus::WaitUser {
317 return Err("only interrupted runs can be converted into RunState".to_string());
318 }
319 Ok(Self {
320 result,
321 approved_interruption_ids: Vec::new(),
322 approval_consumption: Arc::new(Mutex::new(BTreeSet::new())),
323 })
324 }
325
326 pub fn result(&self) -> &RunResult {
327 &self.result
328 }
329
330 pub fn approve(&mut self, interruption_id: &str) -> Result<(), String> {
331 if !self
332 .pending_approval_ids()
333 .iter()
334 .any(|id| id == interruption_id)
335 {
336 return Err(format!("unknown approval interruption: {interruption_id}"));
337 }
338 if !self
339 .approved_interruption_ids
340 .iter()
341 .any(|id| id == interruption_id)
342 {
343 self.approved_interruption_ids
344 .push(interruption_id.to_string());
345 }
346 Ok(())
347 }
348
349 pub fn approved_interruption_ids(&self) -> &[String] {
350 &self.approved_interruption_ids
351 }
352
353 pub fn approvals(&self) -> Vec<ApprovalSnapshot> {
354 let approved = self
355 .approved_interruption_ids
356 .iter()
357 .cloned()
358 .collect::<BTreeSet<_>>();
359 approval_snapshots(self.result.result(), &approved)
360 }
361
362 pub fn approval_snapshot(&self) -> Vec<ApprovalSnapshot> {
363 self.approvals()
364 }
365
366 pub fn pending_approval_ids(&self) -> Vec<String> {
367 self.result
368 .approvals()
369 .into_iter()
370 .map(|snapshot| snapshot.interruption_id)
371 .collect()
372 }
373
374 pub(crate) fn into_inner(self) -> (RunResult, Vec<String>, Arc<Mutex<BTreeSet<String>>>) {
375 (
376 self.result,
377 self.approved_interruption_ids,
378 self.approval_consumption,
379 )
380 }
381}
382
383fn approval_snapshots(
384 result: &AgentResult,
385 approved_interruption_ids: &BTreeSet<String>,
386) -> Vec<ApprovalSnapshot> {
387 let mut seen = BTreeSet::new();
388 let mut snapshots = Vec::new();
389 for cycle in &result.cycles {
390 for tool_result in &cycle.tool_results {
391 let metadata = &tool_result.metadata;
392 let interruption_id = metadata
393 .get("approval_interruption_id")
394 .or_else(|| metadata.get("request_id"))
395 .and_then(Value::as_str)
396 .map(str::trim)
397 .filter(|value| !value.is_empty());
398 let Some(interruption_id) = interruption_id else {
399 continue;
400 };
401 let approval_requested = metadata
402 .get("mode")
403 .and_then(Value::as_str)
404 .is_some_and(|mode| mode == "approval_requested")
405 || metadata
406 .get("approval_required")
407 .and_then(Value::as_bool)
408 .unwrap_or(false);
409 if !approval_requested || !seen.insert(interruption_id.to_string()) {
410 continue;
411 }
412 let arguments = metadata
413 .get("arguments")
414 .and_then(Value::as_object)
415 .map(|arguments| {
416 arguments
417 .iter()
418 .map(|(key, value)| (key.clone(), value.clone()))
419 .collect()
420 })
421 .unwrap_or_default();
422 snapshots.push(ApprovalSnapshot {
423 interruption_id: interruption_id.to_string(),
424 tool_name: metadata
425 .get("tool_name")
426 .and_then(Value::as_str)
427 .unwrap_or_default()
428 .to_string(),
429 tool_call_id: tool_result.tool_call_id.clone(),
430 arguments,
431 message: metadata
432 .get("message")
433 .and_then(Value::as_str)
434 .unwrap_or(&tool_result.content)
435 .to_string(),
436 cycle_index: Some(cycle.index),
437 approved: approved_interruption_ids.contains(interruption_id),
438 });
439 }
440 }
441 snapshots
442}
443
444fn resolved_model_public_value(resolved: &ResolvedModelConfig) -> Value {
445 json!({
446 "backend": resolved.backend,
447 "requested_model": resolved.requested_model,
448 "selected_model": resolved.selected_model,
449 "model_id": resolved.model_id,
450 "endpoint": resolved
451 .endpoint_options
452 .first()
453 .map(|option| option.endpoint.endpoint_id.as_str()),
454 })
455}
456
457#[derive(Clone)]
458pub(crate) struct RunResumeContext {
459 pub agent: Agent,
460 pub input: crate::runner::NormalizedInput,
461 pub config: RunConfig,
462 pub runner: crate::runner::Runner,
463 pub pending_tool_approval: Option<PendingToolApproval>,
464}
465
466#[derive(Clone)]
467pub(crate) struct PendingToolApproval {
468 pub interruption_id: String,
469 pub call: crate::types::ToolCall,
470 pub cycle_index: u32,
471 pub context: ToolContext,
472 pub options: ToolRunOptions,
473 pub orchestrator: ToolOrchestrator,
474 pub task: crate::types::AgentTask,
475 pub hook_manager: crate::runtime::RuntimeHookManager,
476}