use super::*;
use crate::tool_dispatch::schedule_tool_batch;
pub(in crate::runtime) struct ToolBatchRunOutcome {
pub launches: Vec<crate::runtime::ToolCallLaunch>,
pub triggers: Vec<crate::tool_dispatch::ToolTriggerEffectOutcome>,
}
async fn run_one_tool_call(
index: usize,
prepared_tool: crate::PreparedToolCall,
invocation: crate::RuntimeInvocation,
context: crate::RuntimeExecutionContext<'_>,
) -> crate::runtime::ToolCallLaunch {
Box::pin(context.execute_prepared_tool_call_launch(prepared_tool, index, Some(invocation)))
.await
}
fn cancelled_tool_call_launch(
call_id: String,
tool_name: String,
args: serde_json::Value,
replay: Option<crate::llm::types::ProviderReplayMeta>,
) -> crate::runtime::ToolCallLaunch {
crate::runtime::ToolCallLaunch::Done {
result: cancelled_completed_tool_call(call_id, tool_name, args, replay),
}
}
fn cancelled_completed_tool_call(
call_id: String,
tool_name: String,
args: serde_json::Value,
replay: Option<crate::llm::types::ProviderReplayMeta>,
) -> crate::sansio::CompletedToolCall {
let output =
crate::ToolCallOutput::cancelled(crate::ToolCancellation::runtime("tool call cancelled"));
crate::sansio::CompletedToolCall {
call_id: call_id.clone(),
tool_name: tool_name.clone(),
args,
model_return: crate::ModelToolReturn {
call_id,
tool_name,
parts: vec![crate::ModelToolReturnPart::text(
"[Tool execution cancelled]\ntool call cancelled".to_string(),
)],
},
output,
duration_ms: 0,
replay,
}
}
impl RuntimeTurnDriver<'_> {
pub(super) async fn invoke_turn_tool_calls_effect(
&mut self,
machine: &mut TurnMachine,
id: crate::sansio::EffectId,
calls: Vec<crate::sansio::PendingToolCall>,
event_tx: &mpsc::Sender<RuntimeStreamEvent>,
cancel: &CancellationToken,
) -> Result<Vec<crate::sansio::CompletedToolCall>, RuntimeEffectControllerError> {
let (tool_event_tx, mut tool_event_rx) = tokio::sync::mpsc::channel::<SessionEvent>(64);
let runtime_event_tx = event_tx.clone();
let tool_event_forwarder = tokio::spawn(async move {
while let Some(event) = tool_event_rx.recv().await {
send_session_event(&runtime_event_tx, event).await;
}
});
let prepare_context = self
.execution_context(
tool_event_tx.clone(),
Arc::new(crate::ChronologicalProjection::default()),
)
.map_err(|err| {
RuntimeEffectControllerError::new("tool_catalog_resolution_failed", err.to_string())
})?;
let mut results = Vec::with_capacity(calls.len());
for (index, call) in calls.into_iter().enumerate() {
let call_id = call.call_id.clone();
let replay = call.replay.clone();
let prepared = match prepare_context.prepare_tool_call(call).await {
crate::tool_dispatch::ToolPreparationOutcome::Prepared(prepared) => prepared,
crate::tool_dispatch::ToolPreparationOutcome::Completed(outcome) => {
let completed = prepare_context
.complete_tool_call(
index,
call_id.clone(),
replay,
*outcome,
crate::TurnActivityId::new(format!("tool:{call_id}")),
)
.await
.completed;
results.push(completed);
continue;
}
};
let parent_invocation =
self.turn_effect_invocation(machine, id, RuntimeEffectKind::ToolCall)?;
let invocation = crate::runtime::causal::child_tool_effect_invocation(
&parent_invocation,
id,
&prepared.call_id,
);
let prepared_for_completion = prepared.clone();
let outcome = self
.execute_typed_turn_effect(
machine,
event_tx,
cancel,
RuntimeEffectEnvelope::new(
invocation,
RuntimeEffectCommand::ToolCall { call: prepared },
),
RuntimeEffectOutcome::into_tool_call_effect,
)
.await?;
match outcome.launch {
crate::runtime::ToolCallLaunch::Done { result } => results.push(result),
crate::runtime::ToolCallLaunch::Pending {
key,
pending,
duration_ms,
} => {
let resolution = self
.await_pending_tool_completion(
machine,
id,
&prepared_for_completion.call_id,
key,
&pending,
event_tx,
cancel,
)
.await?;
let dispatch_outcome = prepare_context
.pending_completion_dispatch_outcome(
prepared_for_completion.tool_name.clone(),
prepared_for_completion.args.clone(),
resolution,
duration_ms,
)
.await;
let completed = prepare_context
.complete_tool_call(
index,
call_id.clone(),
replay,
dispatch_outcome,
crate::TurnActivityId::new(format!("tool:{call_id}")),
)
.await
.completed;
send_turn_activity(
event_tx,
crate::TurnActivityId::new(format!("tool:{call_id}")),
crate::TurnEvent::ToolCallCompleted {
call_id: Some(call_id.clone()),
name: completed.tool_name.clone(),
args: completed.args.clone(),
output: completed.output.clone(),
duration_ms: completed.duration_ms,
},
)
.await;
results.push(completed);
}
}
}
drop(prepare_context);
drop(tool_event_tx);
let _ = tool_event_forwarder.await;
Ok(results)
}
pub(in crate::runtime) async fn run_tool_calls(
&mut self,
pending_tools: Vec<(crate::PreparedToolCall, crate::RuntimeInvocation)>,
event_tx: &mpsc::Sender<RuntimeStreamEvent>,
cancel: &CancellationToken,
) -> Result<ToolBatchRunOutcome, crate::RuntimeEffectControllerError> {
let (tool_event_tx, mut tool_event_rx) = tokio::sync::mpsc::channel::<SessionEvent>(64);
let (turn_event_tx, mut turn_event_rx) = tokio::sync::mpsc::channel::<TurnActivity>(64);
let runtime_event_tx = event_tx.clone();
let tool_event_forwarder = tokio::spawn(async move {
while let Some(event) = tool_event_rx.recv().await {
send_session_event(&runtime_event_tx, event).await;
}
});
let runtime_event_tx = event_tx.clone();
let turn_event_forwarder = tokio::spawn(async move {
while let Some(event) = turn_event_rx.recv().await {
let _ = runtime_event_tx.send(RuntimeStreamEvent::Turn(event)).await;
}
});
let context = match self.execution_context(
tool_event_tx.clone(),
Arc::new(crate::ChronologicalProjection::default()),
) {
Ok(context) => context.with_turn_event_sender(turn_event_tx.clone()),
Err(err) => {
drop(tool_event_tx);
drop(turn_event_tx);
let _ = tool_event_forwarder.await;
let _ = turn_event_forwarder.await;
return Err(crate::RuntimeEffectControllerError::new(
"tool_catalog_resolution_failed",
err.to_string(),
));
}
};
let indexed_tools = pending_tools.into_iter().enumerate().collect::<Vec<_>>();
let tool_cancel = cancel.child_token();
let outcomes = schedule_tool_batch(
indexed_tools,
|(index, _)| *index,
|(_, (pending_tool, _))| context.tool_scheduling(&pending_tool.tool_name),
{
let context = context.clone();
let cancel = cancel.clone();
let tool_cancel = tool_cancel.clone();
move |(index, (pending_tool, parent_invocation))| {
let context = context.clone().with_cancellation_token(tool_cancel.clone());
let cancel = cancel.clone();
let tool_cancel = tool_cancel.clone();
let cancelled_tool = pending_tool.clone();
async move {
let tool_call =
run_one_tool_call(index, pending_tool, parent_invocation, context);
tokio::pin!(tool_call);
tokio::select! {
biased;
_ = cancel.cancelled() => {
tool_cancel.cancel();
let grace = tokio::time::sleep(std::time::Duration::from_millis(50));
tokio::pin!(grace);
tokio::select! {
biased;
outcome = &mut tool_call => outcome,
_ = &mut grace => cancelled_tool_call_launch(
cancelled_tool.call_id,
cancelled_tool.tool_name,
cancelled_tool.args,
cancelled_tool.replay,
),
}
}
outcome = &mut tool_call => outcome,
}
}
}
},
)
.await;
let triggers = context.drain_tool_trigger_outcomes().map_err(|err| {
crate::RuntimeEffectControllerError::new("tool_trigger_outcome_drain", err.to_string())
})?;
drop(context);
drop(tool_event_tx);
drop(turn_event_tx);
let _ = tool_event_forwarder.await;
let _ = turn_event_forwarder.await;
Ok(ToolBatchRunOutcome {
launches: outcomes,
triggers,
})
}
async fn await_pending_tool_completion(
&mut self,
machine: &mut TurnMachine,
parent_effect_id: crate::sansio::EffectId,
call_id: &str,
key: crate::AwaitEventKey,
_pending: &crate::PendingCompletion,
event_tx: &mpsc::Sender<RuntimeStreamEvent>,
cancel: &CancellationToken,
) -> Result<crate::Resolution, RuntimeEffectControllerError> {
let parent =
self.turn_effect_invocation(machine, parent_effect_id, RuntimeEffectKind::ToolCall)?;
let invocation = crate::runtime::causal::child_effect_invocation(
&parent,
format!("{}:{call_id}:await", parent_effect_id.0),
RuntimeEffectKind::AwaitEvent,
format!("{call_id}:await"),
);
let _ = event_tx;
let scoped_effect_controller = self.scoped_effect_controller.clone();
let deadline = _pending
.deadline
.map(|duration| std::time::Instant::now() + duration);
let outcome = scoped_effect_controller
.controller()
.execute_effect(
RuntimeEffectEnvelope::new(invocation, RuntimeEffectCommand::AwaitEvent { key }),
crate::RuntimeEffectLocalExecutor::await_event(cancel.clone(), deadline),
)
.await?;
RuntimeEffectOutcome::into_await_event(outcome)
}
}