use std::sync::{Arc, Mutex};
use serde_json::Value;
use crate::budget::{BudgetEnforcementBoundary, BudgetEvaluator};
use crate::events::{RunEvent, RunEventPayload, ToolStatus};
use crate::result::{PendingToolApproval, RunResult, RunResumeContext, RunState};
use crate::run_config::INITIAL_BUDGET_USAGE_METADATA_KEY;
use crate::tools::ToolLifecycleEvent;
use crate::types::{
last_assistant_output, AgentResult, AgentStatus, CompletionReason, ToolDirective,
};
use super::helpers::terminal_event;
use super::session_blocking::block_on_session;
use super::support::{
apply_cancellation_precedence, apply_optional_output_validation, apply_output_guardrails,
capture_event, effective_event_store, extract_handoff, output_type_validation_error,
SingleRunOutcome,
};
use super::{effective_session_id, NormalizedInput, Runner};
impl Runner {
pub async fn resume(&self, state: RunState) -> Result<RunResult, String> {
Box::pin(self.resume_with_optional_input(state, None)).await
}
pub async fn resume_with_input(
&self,
state: RunState,
input: impl Into<NormalizedInput>,
) -> Result<RunResult, String> {
Box::pin(self.resume_with_optional_input(state, Some(input.into()))).await
}
async fn resume_with_optional_input(
&self,
state: RunState,
input: Option<NormalizedInput>,
) -> Result<RunResult, String> {
let (source, approved_ids, approval_consumption) = state.into_inner();
let Some(resume_context) = source.resume_context().cloned() else {
return Err("run state does not include resume context".to_string());
};
let origin_runner = resume_context.runner.clone();
if let Some(result) = Box::pin(origin_runner.resume_approved_tool_call(
&source,
&resume_context,
&approved_ids,
&approval_consumption,
input.as_ref(),
))
.await
{
return result;
}
let mut config = resume_context.config;
config.initial_messages = Some(source.result().messages.clone());
config.initial_shared_state = source.result().shared_state.clone();
set_initial_budget_usage(&mut config, source.budget_usage())?;
let result = origin_runner
.run_with_config(
&resume_context.agent,
input.unwrap_or(resume_context.input),
config,
)
.await
.map_err(|error| error.to_string())?;
Ok(result)
}
async fn resume_approved_tool_call(
&self,
source: &RunResult,
resume_context: &RunResumeContext,
approved_ids: &[String],
approval_consumption: &Arc<Mutex<std::collections::BTreeSet<String>>>,
resume_input: Option<&NormalizedInput>,
) -> Option<Result<RunResult, String>> {
let approval = match select_approved_tool_context(
resume_context.pending_tool_approval.as_ref(),
approved_ids,
) {
Ok(Some(approval)) => approval.clone(),
Ok(None) => return None,
Err(error) => return Some(Err(error)),
};
if !approval_snapshot_matches_result(source.result(), &approval) {
return Some(Err(
"approved tool call does not match the captured interruption".to_string(),
));
}
if resume_input.is_some() {
return Some(Err(
"input cannot be provided when resuming an approved tool call".to_string(),
));
}
let cancellation_token = resume_context
.config
.cancellation_token
.as_ref()
.or(self.default_run_config.cancellation_token.as_ref());
if cancellation_token.is_some_and(crate::runtime::CancellationToken::is_cancelled) {
let mut cancelled = source.result().clone();
cancelled.status = AgentStatus::Failed;
cancelled.completion_reason = Some(CompletionReason::Cancelled);
cancelled.completion_tool_name = None;
cancelled.partial_output = cancelled
.partial_output
.or_else(|| last_assistant_output(&cancelled.cycles));
cancelled.error = Some(
cancellation_token
.and_then(crate::runtime::CancellationToken::reason)
.unwrap_or_else(|| "run cancelled".to_string()),
);
cancelled.budget_exhaustion = None;
cancelled.final_answer = None;
cancelled.wait_reason = None;
return Some(self.finalize_approval_terminal(
source,
resume_context,
&approval.interruption_id,
cancelled,
source.new_items().to_vec(),
cancellation_token,
None,
Vec::new(),
));
}
{
let mut consumed = approval_consumption
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if !consumed.insert(approval.interruption_id.clone()) {
return Some(Err("approval_already_consumed".to_string()));
}
}
let resumed_run_id = format!("run_{}", uuid::Uuid::new_v4().simple());
let budget_limits = resume_context
.config
.budget_limits
.clone()
.or_else(|| self.default_run_config.budget_limits.clone());
let host_cost_meter = resume_context
.config
.host_cost_meter
.clone()
.or_else(|| self.default_run_config.host_cost_meter.clone());
let mut budget_evaluator = match budget_limits.filter(|limits| limits.has_limits()) {
Some(limits) => {
match BudgetEvaluator::new(limits, host_cost_meter, source.budget_usage().cloned())
{
Ok(evaluator) => Some(Box::new(evaluator)),
Err(error) => return Some(Err(error)),
}
}
None => None,
};
let mut context = approval.context.clone();
context.shared_state = source.result().shared_state.clone();
let call = approval.call.clone();
let lifecycle_observations = Arc::new(Mutex::new(Vec::<ToolLifecycleEvent>::new()));
let lifecycle_observations_for_callback = lifecycle_observations.clone();
let options = approval
.options
.clone()
.lifecycle_callback(Arc::new(move |event| {
lifecycle_observations_for_callback
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push(event);
}));
let execution = approval
.orchestrator
.run_one_with_approval_and_metadata_deferred(
call.clone(),
&mut context,
options,
|_call, _requirement, _context, _metadata| None,
)
.await
.map_err(|error| error.to_string());
let mut execution = match execution {
Ok(execution) => execution,
Err(error) => {
let persistence = persist_approval_lifecycle_events(
self,
source,
resume_context,
&resumed_run_id,
approval.cycle_index,
&lifecycle_observations,
);
return Some(persistence.and(Err(error)));
}
};
let tool_result = execution.result().clone();
let mut tool_result = approval.hook_manager.apply_after_tool_call(
&approval.task,
approval.cycle_index,
&call,
&context,
tool_result,
);
let behavior_reason = crate::runtime::tool_call_runner::apply_tool_use_behavior(
&approval.task,
&call,
&mut tool_result,
);
execution.replace_result(tool_result);
let tool_result = execution.complete();
let resume_tool_events = match persist_approval_lifecycle_events(
self,
source,
resume_context,
&resumed_run_id,
approval.cycle_index,
&lifecycle_observations,
) {
Ok(events) => events,
Err(error) => return Some(Err(error)),
};
let mut agent_result = source.result().clone();
agent_result.shared_state = context.shared_state.clone();
if let Some(cycle) = agent_result
.cycles
.iter_mut()
.find(|cycle| cycle.index == approval.cycle_index)
{
if let Some(existing) = cycle.tool_results.iter_mut().find(|existing| {
existing.tool_call_id == call.id
&& existing
.metadata
.get("approval_interruption_id")
.and_then(Value::as_str)
== Some(approval.interruption_id.as_str())
}) {
*existing = tool_result.clone();
} else {
cycle.tool_results.push(tool_result.clone());
}
}
let tool_message = tool_result.to_message();
agent_result.messages.retain(|message| {
!(message.role == crate::types::MessageRole::Tool
&& message.tool_call_id.as_deref() == Some(call.id.as_str()))
});
agent_result.messages.push(tool_message.clone());
if let Some(session) = resume_context.config.session.as_ref() {
let session_items =
crate::sessions::SessionItem::from_message(&tool_message).map(|item| vec![item]);
let Some(session_items) = session_items else {
return Some(Err(
"approved resume messages cannot be persisted to session".to_string(),
));
};
if let Err(error) = block_on_session(session.add_items(session_items)) {
return Some(Err(error));
}
}
let mut new_items = source
.new_items()
.iter()
.filter(|message| {
!(message.role == crate::types::MessageRole::Tool
&& message.tool_call_id.as_deref() == Some(call.id.as_str()))
})
.cloned()
.collect::<Vec<_>>();
new_items.push(tool_message);
let mut resume_budget_events = Vec::new();
if let Some(evaluator) = &mut budget_evaluator {
let observed_exhaustion = evaluator.tool_batch_complete(false);
let snapshot = evaluator.snapshot();
let cancelled =
cancellation_token.is_some_and(crate::runtime::CancellationToken::is_cancelled);
let exhaustion = (!cancelled).then_some(observed_exhaustion).flatten();
agent_result.budget_usage = Some(snapshot.clone());
agent_result.budget_exhaustion = exhaustion.clone();
if exhaustion.is_some() {
agent_result.status = AgentStatus::Failed;
agent_result.completion_reason = Some(CompletionReason::BudgetExhausted);
agent_result.completion_tool_name = None;
agent_result.partial_output = last_assistant_output(&agent_result.cycles);
agent_result.final_answer = None;
agent_result.wait_reason = None;
agent_result.error = Some("Run budget exhausted.".to_string());
}
let payload = match exhaustion.clone() {
Some(budget_exhaustion) => RunEventPayload::BudgetExhausted {
enforcement_boundary: BudgetEnforcementBoundary::ToolBatchComplete,
budget_usage: snapshot,
budget_exhaustion,
},
None => RunEventPayload::BudgetSnapshot {
enforcement_boundary: BudgetEnforcementBoundary::ToolBatchComplete,
budget_usage: snapshot,
},
};
let mut budget_event = RunEvent::new(
&resumed_run_id,
source.trace_id(),
source.agent_name(),
Some(approval.cycle_index),
payload,
);
let session_id = effective_session_id(&self.default_run_config, &resume_context.config);
if let Some(session_id) = session_id.as_deref() {
budget_event = budget_event.with_session_id(session_id);
}
let (event_store, event_store_fail_closed) =
effective_event_store(&self.default_run_config, &resume_context.config);
if let Err(error) = capture_event(
None,
None,
event_store.as_ref(),
event_store_fail_closed,
budget_event.clone(),
) {
return Some(Err(error));
}
resume_budget_events.push(budget_event);
}
if agent_result.completion_reason != Some(CompletionReason::BudgetExhausted)
&& tool_result.directive == ToolDirective::Continue
{
let mut config = resume_context.config.clone();
config.initial_messages = Some(agent_result.messages.clone());
config.initial_shared_state = agent_result.shared_state.clone();
config.trace_id = Some(source.trace_id().to_string());
if let Err(error) =
set_initial_budget_usage(&mut config, agent_result.budget_usage.as_ref())
{
return Some(Err(error));
}
let mut prior_events = Vec::new();
if !resume_budget_events.is_empty() {
prior_events.extend_from_slice(source.events());
}
prior_events.extend(resume_tool_events);
prior_events.extend(resume_budget_events);
let result = self
.run_with_config_and_run_id(
&resume_context.agent,
NormalizedInput::from(source.input().to_string()),
config,
resumed_run_id,
)
.await
.map(move |result| {
let mut events = prior_events;
events.extend_from_slice(result.events());
let mut metadata = result.metadata().clone();
metadata.insert("resumed".to_string(), Value::Bool(true));
metadata.insert(
"approved_interruption_id".to_string(),
Value::String(approval.interruption_id.clone()),
);
result.with_events(events).with_metadata(metadata)
});
return Some(result);
}
if agent_result.completion_reason != Some(CompletionReason::BudgetExhausted) {
let completion_reason = behavior_reason.unwrap_or(match tool_result.directive {
ToolDirective::Finish => CompletionReason::ToolFinish,
ToolDirective::WaitUser => CompletionReason::WaitUser,
ToolDirective::Continue => unreachable!(),
});
agent_result.completion_reason = Some(completion_reason);
agent_result.completion_tool_name = Some(call.name.clone());
agent_result.error = None;
match tool_result.directive {
ToolDirective::Finish => {
agent_result.status = AgentStatus::Completed;
agent_result.partial_output = None;
agent_result.final_answer =
Some(crate::runtime::extract_final_message(&tool_result));
agent_result.wait_reason = None;
}
ToolDirective::WaitUser => {
agent_result.status = AgentStatus::WaitUser;
agent_result.partial_output = last_assistant_output(&agent_result.cycles);
agent_result.final_answer = None;
agent_result.wait_reason =
Some(crate::runtime::extract_wait_reason(&tool_result));
}
ToolDirective::Continue => unreachable!(),
}
}
let guardrail_context = context
.run_context
.clone()
.unwrap_or_else(|| crate::RunContext {
run_id: source.run_id().to_string(),
agent_name: resume_context.agent.name().to_string(),
metadata: source.metadata().clone(),
..crate::RunContext::default()
});
agent_result =
apply_output_guardrails(&resume_context.agent, &guardrail_context, agent_result);
agent_result = apply_cancellation_precedence(agent_result, cancellation_token);
let output_type_validation_error =
output_type_validation_error(&resume_context.agent, &agent_result);
let (validated_result, output_validation_error) = apply_optional_output_validation(
&resume_context.agent,
&guardrail_context,
agent_result,
output_type_validation_error,
);
agent_result = validated_result;
let resumed = match self.finalize_approval_terminal(
source,
resume_context,
&approval.interruption_id,
agent_result,
new_items,
cancellation_token,
Some(resumed_run_id),
resume_tool_events
.into_iter()
.chain(resume_budget_events)
.collect(),
) {
Ok(resumed) => resumed,
Err(error) => return Some(Err(error)),
};
if let Some(error) = output_validation_error {
return Some(Err(error));
}
let Some(handoff) = extract_handoff(resumed.result()) else {
return Some(Ok(resumed));
};
let event_collector = Arc::new(Mutex::new(resumed.events().to_vec()));
let initial_outcome = SingleRunOutcome {
result: resumed,
handoff: Some(handoff),
};
let runner = self.clone();
let agent = resume_context.agent.clone();
let input = resume_context.input.clone();
let config = resume_context.config.clone();
Some(
tokio::task::spawn_blocking(move || {
runner.run_agent_chain_with_initial(
&agent,
input,
config,
Some(event_collector),
None,
None,
Some(initial_outcome),
None,
)
})
.await
.map_err(|error| format!("resume handoff task failed: {error}"))
.and_then(|result| result),
)
}
#[allow(clippy::too_many_arguments)] fn finalize_approval_terminal(
&self,
source: &RunResult,
resume_context: &RunResumeContext,
interruption_id: &str,
agent_result: AgentResult,
new_items: Vec<crate::types::Message>,
cancellation_token: Option<&crate::runtime::CancellationToken>,
resumed_run_id: Option<String>,
additional_events: Vec<RunEvent>,
) -> Result<RunResult, String> {
let resumed_run_id =
resumed_run_id.unwrap_or_else(|| format!("run_{}", uuid::Uuid::new_v4().simple()));
let mut events = source.events().to_vec();
events.extend(additional_events);
let mut resumed = RunResult::new(
resume_context.agent.name().to_string(),
agent_result,
source
.resolved_model()
.cloned()
.expect("interrupted runs have a resolved model"),
)
.with_ids(&resumed_run_id, source.trace_id())
.with_input(source.input())
.with_new_items(new_items)
.with_events(events)
.with_metadata({
let mut metadata = source.metadata().clone();
metadata.insert("resumed".to_string(), Value::Bool(true));
metadata.insert(
"approved_interruption_id".to_string(),
Value::String(interruption_id.to_string()),
);
metadata
})
.with_resume_context(resume_context.clone());
let event_collector = Arc::new(Mutex::new(resumed.events().to_vec()));
let session_id = effective_session_id(&self.default_run_config, &resume_context.config);
let (event_store, event_store_fail_closed) =
effective_event_store(&self.default_run_config, &resume_context.config);
capture_event(
Some(&event_collector),
None,
event_store.as_ref(),
event_store_fail_closed,
terminal_event(
resumed.result(),
resumed.run_id(),
resumed.trace_id(),
resume_context.agent.name(),
session_id.as_deref(),
cancellation_token,
),
)?;
let events = event_collector
.lock()
.map(|events| events.clone())
.unwrap_or_default();
resumed = resumed.with_events(events);
Ok(resumed)
}
}
fn persist_approval_lifecycle_events(
runner: &Runner,
source: &RunResult,
resume_context: &RunResumeContext,
resumed_run_id: &str,
cycle_index: u32,
observations: &Arc<Mutex<Vec<ToolLifecycleEvent>>>,
) -> Result<Vec<RunEvent>, String> {
let observations = std::mem::take(
&mut *observations
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner),
);
let session_id = effective_session_id(&runner.default_run_config, &resume_context.config);
let (event_store, event_store_fail_closed) =
effective_event_store(&runner.default_run_config, &resume_context.config);
let mut events = Vec::with_capacity(observations.len());
for observation in observations {
let mut event = approval_lifecycle_run_event(
observation,
resumed_run_id,
source.trace_id(),
source.agent_name(),
cycle_index,
);
if let Some(session_id) = session_id.as_deref() {
event = event.with_session_id(session_id);
}
capture_event(
None,
None,
event_store.as_ref(),
event_store_fail_closed,
event.clone(),
)?;
events.push(event);
}
Ok(events)
}
fn approval_lifecycle_run_event(
observation: ToolLifecycleEvent,
run_id: &str,
trace_id: &str,
agent_name: &str,
cycle_index: u32,
) -> RunEvent {
match observation {
ToolLifecycleEvent::Planned {
call,
tool_metadata,
} => RunEvent::tool_call_planned(
run_id,
trace_id,
agent_name,
cycle_index,
call.id,
call.name,
Value::Object(call.arguments.into_iter().collect()),
)
.with_tool_metadata(tool_metadata.as_ref()),
ToolLifecycleEvent::Started {
call,
tool_metadata,
} => RunEvent::tool_call_started(
run_id,
trace_id,
agent_name,
cycle_index,
call.id,
call.name,
Value::Object(call.arguments.into_iter().collect()),
)
.with_tool_metadata(tool_metadata.as_ref()),
ToolLifecycleEvent::Completed {
call,
result,
execution_started,
duration_ms,
tool_metadata,
} => {
let status = match result.status {
crate::types::ToolResultStatus::Success => ToolStatus::Success,
crate::types::ToolResultStatus::Error => ToolStatus::Error,
crate::types::ToolResultStatus::WaitResponse => ToolStatus::WaitResponse,
crate::types::ToolResultStatus::Running => ToolStatus::Running,
crate::types::ToolResultStatus::PendingCompress => ToolStatus::PendingCompress,
};
RunEvent::new(
run_id,
trace_id,
agent_name,
Some(cycle_index),
RunEventPayload::ToolCallCompleted {
tool_call_id: result.tool_call_id.clone(),
tool_name: call.name,
status,
directive: result.directive,
error_code: result.error_code.clone(),
execution_started,
duration_ms,
},
)
.with_tool_metadata(tool_metadata.as_ref())
.with_metadata(
"tool_arguments",
Value::Object(call.arguments.into_iter().collect()),
)
.with_metadata(
"metadata",
Value::Object(result.metadata.into_iter().collect()),
)
.with_metadata("content", Value::String(result.content))
}
}
}
fn set_initial_budget_usage(
config: &mut crate::run_config::RunConfig,
usage: Option<&crate::budget::BudgetUsageSnapshot>,
) -> Result<(), String> {
match usage {
Some(usage) => {
let value = serde_json::to_value(usage)
.map_err(|error| format!("failed to serialize resumed budget usage: {error}"))?;
config
.metadata
.insert(INITIAL_BUDGET_USAGE_METADATA_KEY.to_string(), value);
}
None => {
config.metadata.remove(INITIAL_BUDGET_USAGE_METADATA_KEY);
}
}
Ok(())
}
fn select_approved_tool_context<'a>(
pending: Option<&'a PendingToolApproval>,
approved_ids: &[String],
) -> Result<Option<&'a PendingToolApproval>, String> {
if approved_ids.is_empty() {
return Ok(None);
}
let pending = pending.ok_or_else(|| {
"approved tool call is missing its captured interruption context".to_string()
})?;
if !approved_ids.iter().any(|id| id == &pending.interruption_id) {
return Err("approved tool call is missing its captured interruption context".to_string());
}
Ok(Some(pending))
}
fn approval_snapshot_matches_result(result: &AgentResult, approval: &PendingToolApproval) -> bool {
result.cycles.iter().any(|cycle| {
cycle.index == approval.cycle_index
&& cycle.tool_calls.iter().any(|call| call == &approval.call)
&& cycle.tool_results.iter().any(|tool_result| {
tool_result.tool_call_id == approval.call.id
&& tool_result
.metadata
.get("approval_interruption_id")
.and_then(Value::as_str)
== Some(approval.interruption_id.as_str())
&& tool_result
.metadata
.get("tool_name")
.and_then(Value::as_str)
== Some(approval.call.name.as_str())
&& tool_result.metadata.get("arguments")
== Some(&Value::Object(
approval.call.arguments.clone().into_iter().collect(),
))
})
})
}
#[cfg(test)]
mod tests {
use super::select_approved_tool_context;
#[test]
fn approved_id_without_captured_context_fails_closed() {
let error = match select_approved_tool_context(None, &["approval_1".to_string()]) {
Ok(_) => panic!("missing context must fail"),
Err(error) => error,
};
assert_eq!(
error,
"approved tool call is missing its captured interruption context"
);
}
#[test]
fn conversational_resume_without_approved_id_needs_no_approval_context() {
assert!(select_approved_tool_context(None, &[])
.expect("conversational resume")
.is_none());
}
}