1use std::sync::{Arc, Mutex};
2
3use serde_json::Value;
4
5use crate::budget::{BudgetEnforcementBoundary, BudgetEvaluator};
6use crate::events::{RunEvent, RunEventPayload};
7use crate::result::{PendingToolApproval, RunResult, RunResumeContext, RunState};
8use crate::run_config::INITIAL_BUDGET_USAGE_METADATA_KEY;
9use crate::types::{
10 last_assistant_output, AgentResult, AgentStatus, CompletionReason, ToolDirective,
11};
12
13use super::helpers::terminal_event;
14use super::session_blocking::block_on_session;
15use super::support::{
16 apply_cancellation_precedence, apply_output_guardrails, capture_event, effective_event_store,
17 extract_handoff, SingleRunOutcome,
18};
19use super::{effective_session_id, NormalizedInput, Runner};
20
21impl Runner {
22 pub async fn resume(&self, state: RunState) -> Result<RunResult, String> {
23 Box::pin(self.resume_with_optional_input(state, None)).await
24 }
25
26 pub async fn resume_with_input(
27 &self,
28 state: RunState,
29 input: impl Into<NormalizedInput>,
30 ) -> Result<RunResult, String> {
31 Box::pin(self.resume_with_optional_input(state, Some(input.into()))).await
32 }
33
34 async fn resume_with_optional_input(
35 &self,
36 state: RunState,
37 input: Option<NormalizedInput>,
38 ) -> Result<RunResult, String> {
39 let (source, approved_ids, approval_consumption) = state.into_inner();
40 let Some(resume_context) = source.resume_context().cloned() else {
41 return Err("run state does not include resume context".to_string());
42 };
43 let origin_runner = resume_context.runner.clone();
44 if let Some(result) = Box::pin(origin_runner.resume_approved_tool_call(
45 &source,
46 &resume_context,
47 &approved_ids,
48 &approval_consumption,
49 input.as_ref(),
50 ))
51 .await
52 {
53 return result;
54 }
55 let mut config = resume_context.config;
56 config.initial_messages = Some(source.result().messages.clone());
57 config.initial_shared_state = source.result().shared_state.clone();
58 set_initial_budget_usage(&mut config, source.budget_usage())?;
59 let result = origin_runner
60 .run_with_config(
61 &resume_context.agent,
62 input.unwrap_or(resume_context.input),
63 config,
64 )
65 .await
66 .map_err(|error| error.to_string())?;
67 Ok(result)
68 }
69
70 async fn resume_approved_tool_call(
71 &self,
72 source: &RunResult,
73 resume_context: &RunResumeContext,
74 approved_ids: &[String],
75 approval_consumption: &Arc<Mutex<std::collections::BTreeSet<String>>>,
76 resume_input: Option<&NormalizedInput>,
77 ) -> Option<Result<RunResult, String>> {
78 let approval = match select_approved_tool_context(
79 resume_context.pending_tool_approval.as_ref(),
80 approved_ids,
81 ) {
82 Ok(Some(approval)) => approval.clone(),
83 Ok(None) => return None,
84 Err(error) => return Some(Err(error)),
85 };
86 if !approval_snapshot_matches_result(source.result(), &approval) {
87 return Some(Err(
88 "approved tool call does not match the captured interruption".to_string(),
89 ));
90 }
91 if resume_input.is_some() {
92 return Some(Err(
93 "input cannot be provided when resuming an approved tool call".to_string(),
94 ));
95 }
96 let cancellation_token = resume_context
97 .config
98 .cancellation_token
99 .as_ref()
100 .or(self.default_run_config.cancellation_token.as_ref());
101 if cancellation_token.is_some_and(crate::runtime::CancellationToken::is_cancelled) {
102 let mut cancelled = source.result().clone();
103 cancelled.status = AgentStatus::Failed;
104 cancelled.completion_reason = Some(CompletionReason::Cancelled);
105 cancelled.completion_tool_name = None;
106 cancelled.partial_output = cancelled
107 .partial_output
108 .or_else(|| last_assistant_output(&cancelled.cycles));
109 cancelled.error = Some(
110 cancellation_token
111 .and_then(crate::runtime::CancellationToken::reason)
112 .unwrap_or_else(|| "run cancelled".to_string()),
113 );
114 cancelled.budget_exhaustion = None;
115 cancelled.final_answer = None;
116 cancelled.wait_reason = None;
117 return Some(self.finalize_approval_terminal(
118 source,
119 resume_context,
120 &approval.interruption_id,
121 cancelled,
122 source.new_items().to_vec(),
123 cancellation_token,
124 None,
125 Vec::new(),
126 ));
127 }
128 {
129 let mut consumed = approval_consumption
130 .lock()
131 .unwrap_or_else(std::sync::PoisonError::into_inner);
132 if !consumed.insert(approval.interruption_id.clone()) {
133 return Some(Err("approval_already_consumed".to_string()));
134 }
135 }
136 let resumed_run_id = format!("run_{}", uuid::Uuid::new_v4().simple());
137 let budget_limits = resume_context
138 .config
139 .budget_limits
140 .clone()
141 .or_else(|| self.default_run_config.budget_limits.clone());
142 let host_cost_meter = resume_context
143 .config
144 .host_cost_meter
145 .clone()
146 .or_else(|| self.default_run_config.host_cost_meter.clone());
147 let mut budget_evaluator = match budget_limits.filter(|limits| limits.has_limits()) {
148 Some(limits) => {
149 match BudgetEvaluator::new(limits, host_cost_meter, source.budget_usage().cloned())
150 {
151 Ok(evaluator) => Some(Box::new(evaluator)),
152 Err(error) => return Some(Err(error)),
153 }
154 }
155 None => None,
156 };
157 let mut context = approval.context.clone();
158 context.shared_state = source.result().shared_state.clone();
159 let call = approval.call.clone();
160 let tool_result = approval
161 .orchestrator
162 .run_one_with_approval(
163 call.clone(),
164 &mut context,
165 approval.options.clone(),
166 |_call, _requirement, _context| None,
167 )
168 .await
169 .map_err(|error| error.to_string());
170 let tool_result = match tool_result {
171 Ok(result) => result,
172 Err(error) => return Some(Err(error)),
173 };
174 let mut tool_result = approval.hook_manager.apply_after_tool_call(
175 &approval.task,
176 approval.cycle_index,
177 &call,
178 &context,
179 tool_result,
180 );
181 let behavior_reason = crate::runtime::tool_call_runner::apply_tool_use_behavior(
182 &approval.task,
183 &call,
184 &mut tool_result,
185 );
186 let mut agent_result = source.result().clone();
187 agent_result.shared_state = context.shared_state.clone();
188 if let Some(cycle) = agent_result
189 .cycles
190 .iter_mut()
191 .find(|cycle| cycle.index == approval.cycle_index)
192 {
193 if let Some(existing) = cycle.tool_results.iter_mut().find(|existing| {
194 existing.tool_call_id == call.id
195 && existing
196 .metadata
197 .get("approval_interruption_id")
198 .and_then(Value::as_str)
199 == Some(approval.interruption_id.as_str())
200 }) {
201 *existing = tool_result.clone();
202 } else {
203 cycle.tool_results.push(tool_result.clone());
204 }
205 }
206 let tool_message = tool_result.to_message();
207 agent_result.messages.retain(|message| {
208 !(message.role == crate::types::MessageRole::Tool
209 && message.tool_call_id.as_deref() == Some(call.id.as_str()))
210 });
211 agent_result.messages.push(tool_message.clone());
212 if let Some(session) = resume_context.config.session.as_ref() {
213 let session_items =
214 crate::sessions::SessionItem::from_message(&tool_message).map(|item| vec![item]);
215 let Some(session_items) = session_items else {
216 return Some(Err(
217 "approved resume messages cannot be persisted to session".to_string(),
218 ));
219 };
220 if let Err(error) = block_on_session(session.add_items(session_items)) {
221 return Some(Err(error));
222 }
223 }
224 let mut new_items = source
225 .new_items()
226 .iter()
227 .filter(|message| {
228 !(message.role == crate::types::MessageRole::Tool
229 && message.tool_call_id.as_deref() == Some(call.id.as_str()))
230 })
231 .cloned()
232 .collect::<Vec<_>>();
233 new_items.push(tool_message);
234
235 let mut resume_budget_events = Vec::new();
236 if let Some(evaluator) = &mut budget_evaluator {
237 let observed_exhaustion = evaluator.tool_batch_complete(false);
238 let snapshot = evaluator.snapshot();
239 let cancelled =
240 cancellation_token.is_some_and(crate::runtime::CancellationToken::is_cancelled);
241 let exhaustion = (!cancelled).then_some(observed_exhaustion).flatten();
242 agent_result.budget_usage = Some(snapshot.clone());
243 agent_result.budget_exhaustion = exhaustion.clone();
244 if exhaustion.is_some() {
245 agent_result.status = AgentStatus::Failed;
246 agent_result.completion_reason = Some(CompletionReason::BudgetExhausted);
247 agent_result.completion_tool_name = None;
248 agent_result.partial_output = last_assistant_output(&agent_result.cycles);
249 agent_result.final_answer = None;
250 agent_result.wait_reason = None;
251 agent_result.error = Some("Run budget exhausted.".to_string());
252 }
253 let payload = match exhaustion.clone() {
254 Some(budget_exhaustion) => RunEventPayload::BudgetExhausted {
255 enforcement_boundary: BudgetEnforcementBoundary::ToolBatchComplete,
256 budget_usage: snapshot,
257 budget_exhaustion,
258 },
259 None => RunEventPayload::BudgetSnapshot {
260 enforcement_boundary: BudgetEnforcementBoundary::ToolBatchComplete,
261 budget_usage: snapshot,
262 },
263 };
264 let mut budget_event = RunEvent::new(
265 &resumed_run_id,
266 source.trace_id(),
267 source.agent_name(),
268 Some(approval.cycle_index),
269 payload,
270 );
271 let session_id = effective_session_id(&self.default_run_config, &resume_context.config);
272 if let Some(session_id) = session_id.as_deref() {
273 budget_event = budget_event.with_session_id(session_id);
274 }
275 let (event_store, event_store_fail_closed) =
276 effective_event_store(&self.default_run_config, &resume_context.config);
277 if let Err(error) = capture_event(
278 None,
279 None,
280 event_store.as_ref(),
281 event_store_fail_closed,
282 budget_event.clone(),
283 ) {
284 return Some(Err(error));
285 }
286 resume_budget_events.push(budget_event);
287 }
288
289 if agent_result.completion_reason != Some(CompletionReason::BudgetExhausted)
290 && tool_result.directive == ToolDirective::Continue
291 {
292 let mut config = resume_context.config.clone();
293 config.initial_messages = Some(agent_result.messages.clone());
294 config.initial_shared_state = agent_result.shared_state.clone();
295 config.trace_id = Some(source.trace_id().to_string());
296 if let Err(error) =
297 set_initial_budget_usage(&mut config, agent_result.budget_usage.as_ref())
298 {
299 return Some(Err(error));
300 }
301 let mut prior_events = Vec::new();
302 if !resume_budget_events.is_empty() {
303 prior_events.extend_from_slice(source.events());
304 prior_events.extend(resume_budget_events);
305 }
306 let result = self
307 .run_with_config(&resume_context.agent, source.input().to_string(), config)
308 .await
309 .map(move |result| {
310 let mut events = prior_events;
311 events.extend_from_slice(result.events());
312 let mut metadata = result.metadata().clone();
313 metadata.insert("resumed".to_string(), Value::Bool(true));
314 metadata.insert(
315 "approved_interruption_id".to_string(),
316 Value::String(approval.interruption_id.clone()),
317 );
318 result.with_events(events).with_metadata(metadata)
319 });
320 return Some(result);
321 }
322
323 if agent_result.completion_reason != Some(CompletionReason::BudgetExhausted) {
324 let completion_reason = behavior_reason.unwrap_or(match tool_result.directive {
325 ToolDirective::Finish => CompletionReason::ToolFinish,
326 ToolDirective::WaitUser => CompletionReason::WaitUser,
327 ToolDirective::Continue => unreachable!(),
328 });
329 agent_result.completion_reason = Some(completion_reason);
330 agent_result.completion_tool_name = Some(call.name.clone());
331 agent_result.error = None;
332 match tool_result.directive {
333 ToolDirective::Finish => {
334 agent_result.status = AgentStatus::Completed;
335 agent_result.partial_output = None;
336 agent_result.final_answer =
337 Some(crate::runtime::extract_final_message(&tool_result));
338 agent_result.wait_reason = None;
339 }
340 ToolDirective::WaitUser => {
341 agent_result.status = AgentStatus::WaitUser;
342 agent_result.partial_output = last_assistant_output(&agent_result.cycles);
343 agent_result.final_answer = None;
344 agent_result.wait_reason =
345 Some(crate::runtime::extract_wait_reason(&tool_result));
346 }
347 ToolDirective::Continue => unreachable!(),
348 }
349 }
350 let guardrail_context = context
351 .run_context
352 .clone()
353 .unwrap_or_else(|| crate::RunContext {
354 run_id: source.run_id().to_string(),
355 agent_name: resume_context.agent.name().to_string(),
356 metadata: source.metadata().clone(),
357 ..crate::RunContext::default()
358 });
359 agent_result =
360 apply_output_guardrails(&resume_context.agent, &guardrail_context, agent_result);
361 agent_result = apply_cancellation_precedence(agent_result, cancellation_token);
362 let output_validation_error = agent_result
363 .final_answer
364 .as_deref()
365 .filter(|_| agent_result.status == AgentStatus::Completed)
366 .and_then(|output| {
367 resume_context
368 .agent
369 .validate_output(output)
370 .err()
371 .map(|error| {
372 format!(
373 "failed to validate final output for agent `{}` as `{}`: {error}",
374 resume_context.agent.name(),
375 resume_context
376 .agent
377 .output_type_name()
378 .unwrap_or("configured output type")
379 )
380 })
381 });
382 let mut resumed = match self.finalize_approval_terminal(
383 source,
384 resume_context,
385 &approval.interruption_id,
386 agent_result,
387 new_items,
388 cancellation_token,
389 Some(resumed_run_id),
390 resume_budget_events,
391 ) {
392 Ok(resumed) => resumed,
393 Err(error) => return Some(Err(error)),
394 };
395 if let Some(error) = output_validation_error {
396 return Some(Err(error));
397 }
398 let session_id = effective_session_id(&self.default_run_config, &resume_context.config);
399 let Some(handoff) = extract_handoff(resumed.result()) else {
400 return Some(Ok(resumed));
401 };
402
403 let event_collector = Arc::new(Mutex::new(resumed.events().to_vec()));
404 let mut legacy_event = RunEvent::new(
405 resumed.run_id(),
406 resumed.trace_id(),
407 &handoff.from_agent,
408 Some(handoff.cycle_index),
409 RunEventPayload::Handoff {
410 source_agent: handoff.from_agent.clone(),
411 target_agent: handoff.to_agent.clone(),
412 tool_call_id: handoff.tool_call_id.clone(),
413 },
414 );
415 if let Some(session_id) = session_id.as_deref() {
416 legacy_event = legacy_event.with_session_id(session_id);
417 }
418 for (key, value) in &handoff.metadata {
419 legacy_event = legacy_event.with_metadata(key, value.clone());
420 }
421 let (event_store, event_store_fail_closed) =
422 effective_event_store(&self.default_run_config, &resume_context.config);
423 if let Err(error) = capture_event(
424 Some(&event_collector),
425 None,
426 event_store.as_ref(),
427 event_store_fail_closed,
428 legacy_event,
429 ) {
430 return Some(Err(error));
431 }
432 let events = event_collector
433 .lock()
434 .map(|events| events.clone())
435 .unwrap_or_default();
436 resumed = resumed.with_events(events);
437 let initial_outcome = SingleRunOutcome {
438 result: resumed,
439 handoff: Some(handoff),
440 };
441 let runner = self.clone();
442 let agent = resume_context.agent.clone();
443 let input = resume_context.input.clone();
444 let config = resume_context.config.clone();
445 Some(
446 tokio::task::spawn_blocking(move || {
447 runner.run_agent_chain_with_initial(
448 &agent,
449 input,
450 config,
451 Some(event_collector),
452 None,
453 None,
454 Some(initial_outcome),
455 )
456 })
457 .await
458 .map_err(|error| format!("resume handoff task failed: {error}"))
459 .and_then(|result| result),
460 )
461 }
462
463 #[allow(clippy::too_many_arguments)] fn finalize_approval_terminal(
465 &self,
466 source: &RunResult,
467 resume_context: &RunResumeContext,
468 interruption_id: &str,
469 agent_result: AgentResult,
470 new_items: Vec<crate::types::Message>,
471 cancellation_token: Option<&crate::runtime::CancellationToken>,
472 resumed_run_id: Option<String>,
473 additional_events: Vec<RunEvent>,
474 ) -> Result<RunResult, String> {
475 let resumed_run_id =
476 resumed_run_id.unwrap_or_else(|| format!("run_{}", uuid::Uuid::new_v4().simple()));
477 let mut events = source.events().to_vec();
478 events.extend(additional_events);
479 let mut resumed = RunResult::new(
480 resume_context.agent.name().to_string(),
481 agent_result,
482 source
483 .resolved_model()
484 .cloned()
485 .expect("interrupted runs have a resolved model"),
486 )
487 .with_ids(&resumed_run_id, source.trace_id())
488 .with_input(source.input())
489 .with_new_items(new_items)
490 .with_events(events)
491 .with_metadata({
492 let mut metadata = source.metadata().clone();
493 metadata.insert("resumed".to_string(), Value::Bool(true));
494 metadata.insert(
495 "approved_interruption_id".to_string(),
496 Value::String(interruption_id.to_string()),
497 );
498 metadata
499 })
500 .with_resume_context(resume_context.clone());
501 let event_collector = Arc::new(Mutex::new(resumed.events().to_vec()));
502 let session_id = effective_session_id(&self.default_run_config, &resume_context.config);
503 let (event_store, event_store_fail_closed) =
504 effective_event_store(&self.default_run_config, &resume_context.config);
505 capture_event(
506 Some(&event_collector),
507 None,
508 event_store.as_ref(),
509 event_store_fail_closed,
510 terminal_event(
511 resumed.result(),
512 resumed.run_id(),
513 resumed.trace_id(),
514 resume_context.agent.name(),
515 session_id.as_deref(),
516 cancellation_token,
517 ),
518 )?;
519 let events = event_collector
520 .lock()
521 .map(|events| events.clone())
522 .unwrap_or_default();
523 resumed = resumed.with_events(events);
524 Ok(resumed)
525 }
526}
527
528fn set_initial_budget_usage(
529 config: &mut crate::run_config::RunConfig,
530 usage: Option<&crate::budget::BudgetUsageSnapshot>,
531) -> Result<(), String> {
532 match usage {
533 Some(usage) => {
534 let value = serde_json::to_value(usage)
535 .map_err(|error| format!("failed to serialize resumed budget usage: {error}"))?;
536 config
537 .metadata
538 .insert(INITIAL_BUDGET_USAGE_METADATA_KEY.to_string(), value);
539 }
540 None => {
541 config.metadata.remove(INITIAL_BUDGET_USAGE_METADATA_KEY);
542 }
543 }
544 Ok(())
545}
546
547fn select_approved_tool_context<'a>(
548 pending: Option<&'a PendingToolApproval>,
549 approved_ids: &[String],
550) -> Result<Option<&'a PendingToolApproval>, String> {
551 if approved_ids.is_empty() {
552 return Ok(None);
553 }
554 let pending = pending.ok_or_else(|| {
555 "approved tool call is missing its captured interruption context".to_string()
556 })?;
557 if !approved_ids.iter().any(|id| id == &pending.interruption_id) {
558 return Err("approved tool call is missing its captured interruption context".to_string());
559 }
560 Ok(Some(pending))
561}
562
563fn approval_snapshot_matches_result(result: &AgentResult, approval: &PendingToolApproval) -> bool {
564 result.cycles.iter().any(|cycle| {
565 cycle.index == approval.cycle_index
566 && cycle.tool_calls.iter().any(|call| call == &approval.call)
567 && cycle.tool_results.iter().any(|tool_result| {
568 tool_result.tool_call_id == approval.call.id
569 && tool_result
570 .metadata
571 .get("approval_interruption_id")
572 .and_then(Value::as_str)
573 == Some(approval.interruption_id.as_str())
574 && tool_result
575 .metadata
576 .get("tool_name")
577 .and_then(Value::as_str)
578 == Some(approval.call.name.as_str())
579 && tool_result.metadata.get("arguments")
580 == Some(&Value::Object(
581 approval.call.arguments.clone().into_iter().collect(),
582 ))
583 })
584 })
585}
586
587#[cfg(test)]
588mod tests {
589 use super::select_approved_tool_context;
590
591 #[test]
592 fn approved_id_without_captured_context_fails_closed() {
593 let error = match select_approved_tool_context(None, &["approval_1".to_string()]) {
594 Ok(_) => panic!("missing context must fail"),
595 Err(error) => error,
596 };
597 assert_eq!(
598 error,
599 "approved tool call is missing its captured interruption context"
600 );
601 }
602
603 #[test]
604 fn conversational_resume_without_approved_id_needs_no_approval_context() {
605 assert!(select_approved_tool_context(None, &[])
606 .expect("conversational resume")
607 .is_none());
608 }
609}