use crate::context::{
self, CompactionStrategy, ContextConfig, ContextTracker, DefaultCompaction, ExecutionLimits,
ExecutionTracker,
};
use crate::provider::{
ModelConfig, ProviderError, StreamConfig, StreamEvent, StreamProvider, ToolDefinition,
};
use crate::types::*;
use std::sync::Arc;
pub type ConvertToLlmFn = Box<dyn Fn(&[AgentMessage]) -> Vec<Message> + Send + Sync>;
pub type TransformContextFn = Box<dyn Fn(Vec<AgentMessage>) -> Vec<AgentMessage> + Send + Sync>;
pub type GetMessagesFn = Box<dyn Fn() -> Vec<AgentMessage> + Send + Sync>;
pub type BeforeTurnFn = Arc<dyn Fn(&[AgentMessage], usize) -> bool + Send + Sync>;
pub type AfterTurnFn = Arc<dyn Fn(&[AgentMessage], &Usage) + Send + Sync>;
pub type OnErrorFn = Arc<dyn Fn(&str) + Send + Sync>;
use tokio::sync::mpsc;
use tracing::warn;
#[non_exhaustive]
pub struct AgentLoopConfig {
pub provider: Arc<dyn StreamProvider>,
pub model: String,
pub api_key: String,
pub thinking_level: ThinkingLevel,
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
pub model_config: Option<ModelConfig>,
pub convert_to_llm: Option<ConvertToLlmFn>,
pub transform_context: Option<TransformContextFn>,
pub get_steering_messages: Option<GetMessagesFn>,
pub get_follow_up_messages: Option<GetMessagesFn>,
pub context_config: Option<ContextConfig>,
pub compaction_strategy: Option<Arc<dyn CompactionStrategy>>,
pub execution_limits: Option<ExecutionLimits>,
pub cache_config: CacheConfig,
pub tool_output_sink: Option<crate::shared_state::SharedState>,
pub tool_execution: ToolExecutionStrategy,
pub tool_middleware: Vec<Arc<dyn ToolMiddleware>>,
pub output_schema: Option<crate::provider::OutputSchema>,
pub retry_config: crate::retry::RetryConfig,
pub before_turn: Option<BeforeTurnFn>,
pub after_turn: Option<AfterTurnFn>,
pub on_error: Option<OnErrorFn>,
pub input_filters: Vec<Arc<dyn InputFilter>>,
pub turn_delay: Option<std::time::Duration>,
pub extensions: Vec<Arc<dyn crate::extension::Extension>>,
pub tree_extensions: Vec<Arc<dyn crate::extension::Extension>>,
pub max_stop_continues: usize,
pub run_label: Option<String>,
pub(crate) inherited_extensions: Vec<Arc<dyn crate::extension::Extension>>,
pub(crate) depth: usize,
pub(crate) delegated_by: Option<String>,
pub(crate) parent_run_id: Option<String>,
}
impl AgentLoopConfig {
pub fn new(provider: Arc<dyn StreamProvider>, model: impl Into<String>) -> Self {
Self {
provider,
model: model.into(),
api_key: String::new(),
thinking_level: ThinkingLevel::Off,
max_tokens: None,
temperature: None,
model_config: None,
convert_to_llm: None,
transform_context: None,
get_steering_messages: None,
get_follow_up_messages: None,
context_config: None,
compaction_strategy: None,
execution_limits: None,
cache_config: CacheConfig::default(),
tool_output_sink: None,
tool_execution: ToolExecutionStrategy::default(),
tool_middleware: Vec::new(),
output_schema: None,
retry_config: crate::retry::RetryConfig::default(),
before_turn: None,
after_turn: None,
on_error: None,
input_filters: Vec::new(),
turn_delay: None,
extensions: Vec::new(),
tree_extensions: Vec::new(),
max_stop_continues: crate::extension::DEFAULT_MAX_STOP_CONTINUES,
run_label: None,
inherited_extensions: Vec::new(),
depth: 0,
delegated_by: None,
parent_run_id: None,
}
}
fn has_extensions(&self) -> bool {
!(self.extensions.is_empty()
&& self.tree_extensions.is_empty()
&& self.inherited_extensions.is_empty())
}
pub fn delegated_from(&mut self, ctx: &ToolContext) -> &mut Self {
self.inherited_extensions = ctx.tree_extensions().to_vec();
self.depth = ctx.delegation_depth();
self.delegated_by = ctx.delegation.call_id.clone();
self.parent_run_id = ctx.delegation.parent_run_id.clone();
if self.run_label.is_none() {
self.run_label = ctx.run_label().map(String::from);
}
self
}
pub(crate) fn tree_for_children(&self) -> Vec<Arc<dyn crate::extension::Extension>> {
self.inherited_extensions
.iter()
.chain(&self.tree_extensions)
.cloned()
.collect()
}
}
fn default_convert_to_llm(messages: &[AgentMessage]) -> Vec<Message> {
messages
.iter()
.filter_map(|m| m.as_llm().cloned())
.collect()
}
pub const AGENT_STOPPED_PREFIX: &str = "[Agent stopped:";
pub const CANCELLED_MARKER: &str = "[Agent stopped: cancelled]";
fn close_turn(tx: &mpsc::UnboundedSender<AgentEvent>, context: &AgentContext) {
if let Some(last) = context.messages.last() {
tx.send(AgentEvent::TurnEnd {
message: last.clone(),
tool_results: Vec::new(),
})
.ok();
}
}
fn push_stop_marker(
text: String,
tx: &mpsc::UnboundedSender<AgentEvent>,
context: &mut AgentContext,
new_messages: &mut Vec<AgentMessage>,
) {
let marker = AgentMessage::Llm(Message::User {
content: vec![Content::Text { text }],
timestamp: now_ms(),
});
tx.send(AgentEvent::MessageStart {
message: marker.clone(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: marker.clone(),
})
.ok();
context.messages.push(marker.clone());
new_messages.push(marker);
}
fn mark_cancelled(
tx: &mpsc::UnboundedSender<AgentEvent>,
context: &mut AgentContext,
new_messages: &mut Vec<AgentMessage>,
) {
if new_messages
.iter()
.any(|m| matches!(m, AgentMessage::Llm(Message::Assistant { .. })))
{
push_stop_marker(CANCELLED_MARKER.to_string(), tx, context, new_messages);
}
}
pub const LOOP_ABORT_PREFIX: &str = "[Agent stopped: repeated tool call —";
pub(crate) const LOOP_NUDGE_PREFIX: &str = "[You have called ";
#[derive(Default)]
struct LoopScope {
decision: DecisionStats,
compaction: CompactionSpend,
prompts: Vec<Message>,
spend: Option<Arc<crate::extension::RunSpend>>,
}
tokio::task_local! {
static LOOP_SCOPE: std::cell::RefCell<LoopScope>;
}
#[cfg_attr(not(feature = "decision"), allow(dead_code))]
pub(crate) fn record_decision(f: impl FnOnce(&mut DecisionStats)) {
let _ = LOOP_SCOPE.try_with(|cell| {
let mut piece = DecisionStats::default();
f(&mut piece);
let mut scope = cell.borrow_mut();
if let Some(spend) = &scope.spend {
spend.add(&piece.usage, piece.cost_usd);
}
scope.decision.merge(&piece);
});
}
pub(crate) fn record_compaction(f: impl FnOnce(&mut CompactionSpend)) -> bool {
LOOP_SCOPE
.try_with(|cell| {
let mut piece = CompactionSpend::default();
f(&mut piece);
let mut scope = cell.borrow_mut();
if let Some(spend) = &scope.spend {
spend.add(&piece.usage, piece.cost_usd);
}
scope.compaction.merge(&piece);
})
.is_ok()
}
#[derive(Default)]
struct ScopeSpend {
decision: DecisionStats,
compaction: CompactionSpend,
}
impl ScopeSpend {
fn fold_into(&self, stats: &mut SessionStats) {
stats.decision.merge(&self.decision);
stats.compaction.merge(&self.compaction);
}
}
pub(crate) fn run_prompts() -> Vec<Message> {
LOOP_SCOPE
.try_with(|cell| cell.borrow().prompts.clone())
.unwrap_or_default()
}
fn user_messages(messages: &[AgentMessage]) -> Vec<Message> {
messages
.iter()
.filter_map(|m| match m {
AgentMessage::Llm(msg @ Message::User { .. }) => Some(msg.clone()),
_ => None,
})
.collect()
}
fn note_run_prompts(messages: &[AgentMessage]) {
let users = user_messages(messages);
if !users.is_empty() {
let _ = LOOP_SCOPE.try_with(|cell| cell.borrow_mut().prompts.extend(users));
}
}
async fn with_loop_scope<T>(
prompts: Vec<Message>,
spend: Arc<crate::extension::RunSpend>,
fut: impl std::future::Future<Output = T>,
) -> (T, ScopeSpend) {
let scope = LoopScope {
prompts,
spend: Some(spend),
..Default::default()
};
LOOP_SCOPE
.scope(std::cell::RefCell::new(scope), async {
let out = fut.await;
let recorded = LOOP_SCOPE.with(|cell| {
let mut scope = cell.borrow_mut();
ScopeSpend {
decision: std::mem::take(&mut scope.decision),
compaction: std::mem::take(&mut scope.compaction),
}
});
(out, recorded)
})
.await
}
const REFUSAL_TOOL_RESULT_TEXT: &str = "Tool call not run: the response was stopped as a refusal \
(declined by the model or stopped by the content filter).";
const CANCELLED_TOOL_RESULT_TEXT: &str = "Tool call not run: the run was cancelled.";
pub async fn agent_loop(
prompts: Vec<AgentMessage>,
context: &mut AgentContext,
config: &AgentLoopConfig,
tx: mpsc::UnboundedSender<AgentEvent>,
cancel: tokio_util::sync::CancellationToken,
) -> Vec<AgentMessage> {
agent_loop_with_stats(prompts, context, config, tx, cancel)
.await
.0
}
pub(crate) async fn agent_loop_with_stats(
prompts: Vec<AgentMessage>,
context: &mut AgentContext,
config: &AgentLoopConfig,
tx: mpsc::UnboundedSender<AgentEvent>,
cancel: tokio_util::sync::CancellationToken,
) -> (Vec<AgentMessage>, SessionStats) {
let prompt_messages = user_messages(&prompts);
let (exts, start_failure) = start_extensions(config, &prompt_messages, &cancel).await;
let observer = EventObserver::new(&exts, tx);
let tx = observer.sender();
tx.send(AgentEvent::AgentStart).ok();
let (outcome, recorded) = with_loop_scope(Vec::new(), exts.spend(), async {
let text = if exts.is_empty() {
String::new()
} else {
prompt_text(&prompts)
};
let prompts = apply_input_filters(prompts, config).await?;
if !exts.is_empty() {
exts.on_input(&crate::extension::InputContext::new(&text, &prompts))
.await?;
}
note_run_prompts(&prompts);
let mut new_messages: Vec<AgentMessage> = prompts.clone();
for prompt in &prompts {
context.messages.push(prompt.clone());
}
tx.send(AgentEvent::TurnStart).ok();
for prompt in &prompts {
tx.send(AgentEvent::MessageStart {
message: prompt.clone(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: prompt.clone(),
})
.ok();
}
let stats = run_with_extensions(
context,
&mut new_messages,
config,
&tx,
&cancel,
&exts,
start_failure,
)
.await;
Ok::<_, String>((new_messages, stats))
})
.await;
let (new_messages, mut stats) = match outcome {
Ok(done) => done,
Err(reason) => {
tx.send(AgentEvent::InputRejected {
reason: reason.clone(),
})
.ok();
exts.finish(&crate::extension::RunOutcome::new(
crate::extension::RunEnd::Rejected {
reason: reason.clone(),
},
None,
))
.await;
let mut stats = SessionStats::default();
recorded.fold_into(&mut stats);
tx.send(AgentEvent::AgentEnd {
messages: vec![],
stats: stats.clone(),
})
.ok();
drop(tx);
observer.finish().await;
return (vec![], stats);
}
};
recorded.fold_into(&mut stats);
exts.finish(&run_outcome(&new_messages, &cancel)).await;
tx.send(AgentEvent::AgentEnd {
messages: new_messages.clone(),
stats: stats.clone(),
})
.ok();
drop(tx);
observer.finish().await;
(new_messages, stats)
}
async fn start_extensions(
config: &AgentLoopConfig,
prompts: &[Message],
cancel: &tokio_util::sync::CancellationToken,
) -> (
Arc<crate::extension::ActiveExtensions>,
Option<crate::extension::Failure>,
) {
use crate::extension::{ActiveExtensions, RunInfo};
if !config.has_extensions() {
return (Arc::new(ActiveExtensions::none()), None);
}
let info = RunInfo {
label: config.run_label.as_deref(),
prompts,
depth: config.depth,
delegated_by: config.delegated_by.as_deref(),
parent_run_id: config.parent_run_id.as_deref(),
cancel,
};
let own: Vec<_> = config
.tree_extensions
.iter()
.chain(&config.extensions)
.cloned()
.collect();
let (exts, failure) = ActiveExtensions::start(&config.inherited_extensions, &own, info).await;
(Arc::new(exts), failure)
}
#[allow(clippy::too_many_arguments)]
async fn run_with_extensions(
context: &mut AgentContext,
new_messages: &mut Vec<AgentMessage>,
config: &AgentLoopConfig,
tx: &mpsc::UnboundedSender<AgentEvent>,
cancel: &tokio_util::sync::CancellationToken,
exts: &crate::extension::ActiveExtensions,
start_failure: Option<crate::extension::Failure>,
) -> SessionStats {
if let Some(failure) = start_failure {
fail_run(&failure, exts, config, tx, context, new_messages, true);
return SessionStats::default();
}
let base_tools = context.tools.len();
if !exts.is_empty() {
let tools = exts.tools().await;
if !tools.is_empty() {
crate::tool_source::merge(&mut context.tools, tools);
}
}
let stats = {
use tracing::Instrument;
run_loop(context, new_messages, config, tx, cancel, exts)
.instrument(tracing::info_span!("agent_loop", model = %config.model))
.await
};
context.tools.truncate(base_tools);
if let Some(failure) = exts.settle().await {
fail_run(&failure, exts, config, tx, context, new_messages, false);
}
stats
}
fn fail_run(
failure: &crate::extension::Failure,
exts: &crate::extension::ActiveExtensions,
config: &AgentLoopConfig,
tx: &mpsc::UnboundedSender<AgentEvent>,
context: &mut AgentContext,
new_messages: &mut Vec<AgentMessage>,
in_turn: bool,
) {
exts.latch_failed();
let text = failure.message();
tracing::error!("{text}");
if !in_turn {
tx.send(AgentEvent::TurnStart).ok();
}
let message = error_message(&config.model, text.clone());
announce(tx, &message);
let am: AgentMessage = message.into();
context.messages.push(am.clone());
new_messages.push(am.clone());
if let Some(ref on_error) = config.on_error {
on_error(&text);
}
tx.send(AgentEvent::TurnEnd {
message: am,
tool_results: Vec::new(),
})
.ok();
}
fn run_outcome(
new_messages: &[AgentMessage],
cancel: &tokio_util::sync::CancellationToken,
) -> crate::extension::RunOutcome {
use crate::extension::{failed_extension, RunEnd, RunOutcome};
let stop_reason = new_messages.iter().rev().find_map(|m| match m {
AgentMessage::Llm(Message::Assistant { stop_reason, .. }) => Some(stop_reason.clone()),
_ => None,
});
let end = match new_messages.last() {
Some(AgentMessage::Llm(Message::User { content, .. })) => match content.first() {
Some(Content::Text { text }) if text == CANCELLED_MARKER => RunEnd::Cancelled,
Some(Content::Text { text }) if text.starts_with(AGENT_STOPPED_PREFIX) => {
RunEnd::Stopped {
reason: text.clone(),
}
}
_ if cancel.is_cancelled() => RunEnd::Cancelled,
_ => RunEnd::Completed,
},
Some(AgentMessage::Llm(Message::Assistant {
stop_reason: StopReason::Aborted,
..
})) => RunEnd::Cancelled,
Some(AgentMessage::Llm(Message::Assistant {
stop_reason: StopReason::Stop | StopReason::Length | StopReason::Refusal,
..
})) => RunEnd::Completed,
Some(AgentMessage::Llm(Message::Assistant {
stop_reason: StopReason::Error,
error_message,
..
})) => {
let error = error_message.clone().unwrap_or_default();
RunEnd::Failed {
extension: failed_extension(&error),
error,
}
}
_ if cancel.is_cancelled() => RunEnd::Cancelled,
_ => RunEnd::Completed,
};
RunOutcome::new(end, stop_reason)
}
struct EventObserver {
tx: mpsc::UnboundedSender<AgentEvent>,
exts: Arc<crate::extension::ActiveExtensions>,
task: Option<(tokio::sync::oneshot::Sender<()>, crate::rt::JoinHandle<()>)>,
}
impl EventObserver {
fn new(
exts: &Arc<crate::extension::ActiveExtensions>,
out: mpsc::UnboundedSender<AgentEvent>,
) -> Self {
if exts.is_empty() {
return Self {
tx: out,
exts: exts.clone(),
task: None,
};
}
let (tx, mut rx) = mpsc::unbounded_channel::<AgentEvent>();
let (flush_tx, mut flush_rx) = mpsc::unbounded_channel::<crate::extension::FlushRequest>();
exts.connect_observer(flush_tx);
let (done_tx, mut done_rx) = tokio::sync::oneshot::channel::<()>();
let observed = exts.clone();
use tracing::instrument::{Instrument, WithSubscriber};
let task = crate::rt::spawn(
async move {
loop {
tokio::select! {
biased;
event = rx.recv() => match event {
Some(event) => {
observed.on_event(&event).await;
out.send(event).ok();
}
None => break,
},
Some(ack) = flush_rx.recv() => {
while let Ok(event) = rx.try_recv() {
observed.on_event(&event).await;
out.send(event).ok();
}
let _ = ack.send(());
}
_ = &mut done_rx => {
while let Ok(event) = rx.try_recv() {
observed.on_event(&event).await;
out.send(event).ok();
}
break;
}
}
}
}
.in_current_span()
.with_current_subscriber(),
);
Self {
tx,
exts: exts.clone(),
task: Some((done_tx, task)),
}
}
fn sender(&self) -> mpsc::UnboundedSender<AgentEvent> {
self.tx.clone()
}
async fn finish(self) {
let Self { tx, exts, task } = self;
drop(tx);
if let Some((done, task)) = task {
let _ = done.send(());
if let Err(e) = task.await {
tracing::error!("extension event observer failed: {e}");
}
}
if let Some(failure) = exts.take_failure() {
tracing::error!(
"{} (on the run's last events: too late to change its outcome)",
failure.message()
);
}
}
}
fn prompt_text(prompts: &[AgentMessage]) -> String {
prompts
.iter()
.filter_map(|m| {
if let AgentMessage::Llm(Message::User { content, .. }) = m {
Some(
content
.iter()
.filter_map(|c| {
if let Content::Text { text } = c {
Some(text.as_str())
} else {
None
}
})
.collect::<Vec<_>>()
.join("\n"),
)
} else {
None
}
})
.collect::<Vec<_>>()
.join("\n")
}
async fn apply_input_filters(
prompts: Vec<AgentMessage>,
config: &AgentLoopConfig,
) -> Result<Vec<AgentMessage>, String> {
let prompts = if !config.input_filters.is_empty() {
let user_text = prompt_text(&prompts);
let mut warnings: Vec<String> = Vec::new();
for filter in &config.input_filters {
let verdict = match filter.as_async() {
Some(async_filter) => {
use futures::FutureExt;
std::panic::AssertUnwindSafe(async_filter.filter(&user_text))
.catch_unwind()
.await
.unwrap_or_else(|_| {
warn!("async input filter panicked; rejecting the input");
FilterResult::Reject("input filter panicked".into())
})
}
None => filter.filter(&user_text),
};
match verdict {
FilterResult::Pass => {}
FilterResult::Warn(w) => warnings.push(w),
FilterResult::Reject(reason) => return Err(reason),
}
}
if !warnings.is_empty() {
let warning_text = warnings
.iter()
.map(|w| format!("[Warning: {}]", w))
.collect::<Vec<_>>()
.join("\n");
let mut modified = prompts;
for msg in modified.iter_mut().rev() {
if let AgentMessage::Llm(Message::User { content, .. }) = msg {
content.push(Content::Text { text: warning_text });
break;
}
}
modified
} else {
prompts
}
} else {
prompts
};
Ok(prompts)
}
pub async fn agent_loop_continue(
context: &mut AgentContext,
config: &AgentLoopConfig,
tx: mpsc::UnboundedSender<AgentEvent>,
cancel: tokio_util::sync::CancellationToken,
) -> Vec<AgentMessage> {
agent_loop_continue_with_stats(context, config, tx, cancel)
.await
.0
}
pub(crate) async fn agent_loop_continue_with_stats(
context: &mut AgentContext,
config: &AgentLoopConfig,
tx: mpsc::UnboundedSender<AgentEvent>,
cancel: tokio_util::sync::CancellationToken,
) -> (Vec<AgentMessage>, SessionStats) {
assert!(
!context.messages.is_empty(),
"Cannot continue: no messages in context"
);
if let Some(last) = context.messages.last() {
assert!(
last.role() != "assistant",
"Cannot continue from assistant message"
);
}
let (exts, start_failure) = start_extensions(config, &[], &cancel).await;
let observer = EventObserver::new(&exts, tx);
let tx = observer.sender();
let mut new_messages: Vec<AgentMessage> = Vec::new();
tx.send(AgentEvent::AgentStart).ok();
tx.send(AgentEvent::TurnStart).ok();
let stats = {
let (mut stats, recorded) = with_loop_scope(
Vec::new(),
exts.spend(),
run_with_extensions(
context,
&mut new_messages,
config,
&tx,
&cancel,
&exts,
start_failure,
),
)
.await;
recorded.fold_into(&mut stats);
stats
};
exts.finish(&run_outcome(&new_messages, &cancel)).await;
tx.send(AgentEvent::AgentEnd {
messages: new_messages.clone(),
stats: stats.clone(),
})
.ok();
drop(tx);
observer.finish().await;
(new_messages, stats)
}
async fn run_loop(
context: &mut AgentContext,
new_messages: &mut Vec<AgentMessage>,
config: &AgentLoopConfig,
tx: &mpsc::UnboundedSender<AgentEvent>,
cancel: &tokio_util::sync::CancellationToken,
exts: &crate::extension::ActiveExtensions,
) -> SessionStats {
let mut stats = SessionStats::default();
let mut first_turn = true;
let mut turn_number: usize = 0;
let mut stop_continues: usize = 0;
let delegation = Delegation {
tree: config.tree_for_children(),
depth: config.depth + 1,
label: config.run_label.clone(),
call_id: None,
parent_run_id: (!exts.is_empty()).then(|| exts.run_id().to_string()),
};
let mut last_context_tokens: Option<usize> = None;
let mut growth_total: usize = 0;
let mut growth_samples: usize = 0;
let mut context_tracker = ContextTracker::new();
let mut tracker = config
.execution_limits
.as_ref()
.map(|limits| ExecutionTracker::new(limits.clone()));
let mut pending: Vec<AgentMessage> = config
.get_steering_messages
.as_ref()
.map(|f| f())
.unwrap_or_default();
loop {
if cancel.is_cancelled() {
mark_cancelled(tx, context, new_messages);
if first_turn {
close_turn(tx, context);
}
return stats;
}
let mut steering_after_tools: Option<Vec<AgentMessage>> = None;
loop {
if cancel.is_cancelled() {
mark_cancelled(tx, context, new_messages);
if first_turn {
close_turn(tx, context);
}
return stats;
}
if !first_turn {
tx.send(AgentEvent::TurnStart).ok();
} else {
first_turn = false;
}
if !pending.is_empty() {
note_run_prompts(&pending);
for msg in pending.drain(..) {
tx.send(AgentEvent::MessageStart {
message: msg.clone(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: msg.clone(),
})
.ok();
context.messages.push(msg.clone());
new_messages.push(msg);
}
}
if let Some(failure) = exts.settle().await {
fail_run(&failure, exts, config, tx, context, new_messages, true);
return stats;
}
if let Some(ref tracker) = tracker {
if let Some(reason) = tracker.check_limits() {
warn!("Execution limit reached: {}", reason);
push_stop_marker(
format!("{AGENT_STOPPED_PREFIX} {}]", reason),
tx,
context,
new_messages,
);
close_turn(tx, context);
return stats;
}
}
if let Some(ref before_turn) = config.before_turn {
if !before_turn(&context.messages, turn_number) {
close_turn(tx, context);
return stats;
}
}
if turn_number > 0 {
if let Some(delay) = config.turn_delay {
crate::rt::sleep(delay).await;
}
}
turn_number += 1;
if let Some(ref ctx_config) = config.context_config {
let estimated = context::total_tokens(&context.messages);
let hybrid = context_tracker.estimate_context_tokens(&context.messages);
let overhead = hybrid.saturating_sub(estimated);
if let Some(previous) = last_context_tokens {
growth_samples += 1;
growth_total += estimated.saturating_sub(previous);
}
let growth_per_turn = if growth_samples > 0 {
growth_total as f64 / growth_samples as f64
} else {
0.0
};
let calibrated;
let effective_config = if overhead > 0 {
let floor = ctx_config.max_context_tokens / 10;
calibrated = ContextConfig {
max_context_tokens: ctx_config
.max_context_tokens
.saturating_sub(overhead)
.max(floor),
system_prompt_tokens: 0,
..ctx_config.clone()
};
tracing::debug!(
"compaction budget calibrated: {} -> {} (measured overhead: {} tokens)",
ctx_config.max_context_tokens,
calibrated.max_context_tokens,
overhead
);
&calibrated
} else {
ctx_config
};
let with_headroom;
let effective_config = {
let ratio = effective_config.effective_target_ratio(growth_per_turn);
if ratio != effective_config.compact_target_ratio {
tracing::debug!(
"compaction target adapted: ratio {} -> {:.3} \
({} tokens/turn growth, {:?} turns of headroom)",
effective_config.compact_target_ratio,
ratio,
growth_per_turn as usize,
effective_config.compact_headroom_turns,
);
with_headroom = ContextConfig {
compact_target_ratio: ratio,
..effective_config.clone()
};
&with_headroom
} else {
effective_config
}
};
let strategy: &dyn CompactionStrategy = config
.compaction_strategy
.as_deref()
.unwrap_or(&DefaultCompaction);
let before_len = context.messages.len();
let before_tokens = estimated;
context.messages =
strategy.compact(std::mem::take(&mut context.messages), effective_config);
let after_tokens = context::total_tokens(&context.messages);
if context.messages.len() != before_len || after_tokens != before_tokens {
context_tracker.reset();
stats.compactions += 1;
}
last_context_tokens = Some(after_tokens);
}
let llm_span = tracing::info_span!(
"llm_stream",
turn = turn_number,
model = %config.model,
tokens_in = tracing::field::Empty,
tokens_out = tracing::field::Empty,
tokens_cached = tracing::field::Empty,
cost_usd = tracing::field::Empty,
error = tracing::field::Empty,
);
let message = {
use tracing::Instrument;
stream_assistant_response(context, config, tx, cancel, exts)
.instrument(llm_span.clone())
.await
};
let message = match message {
Ok(message) => message,
Err(crate::extension::ModelHalt::Stop(reason)) => {
push_stop_marker(
format!("{AGENT_STOPPED_PREFIX} {reason}]"),
tx,
context,
new_messages,
);
close_turn(tx, context);
return stats;
}
Err(crate::extension::ModelHalt::Fail(failure)) => {
fail_run(&failure, exts, config, tx, context, new_messages, true);
return stats;
}
};
if let Message::Assistant {
usage, stop_reason, ..
} = &message
{
llm_span.record("error", *stop_reason == StopReason::Error);
llm_span.record("tokens_in", usage.input);
llm_span.record("tokens_out", usage.output);
llm_span.record("tokens_cached", usage.cache_read);
if let Some(cost) = config.model_config.as_ref().and_then(|mc| mc.cost.as_ref()) {
llm_span.record("cost_usd", cost.cost_usd(usage));
}
}
let message = if structured_output_is_tool_forced(config.model_config.as_ref()) {
unwrap_structured_tool_call(message, config.output_schema.as_ref())
} else {
message
};
let agent_msg: AgentMessage = message.clone().into();
context.messages.push(agent_msg.clone());
new_messages.push(agent_msg.clone());
if let Message::Assistant { usage, .. } = &message {
context_tracker.record_usage(usage, context.messages.len() - 1);
stats.record_turn(
usage,
config.model_config.as_ref().and_then(|mc| mc.cost.as_ref()),
);
}
if let Message::Assistant {
ref stop_reason,
ref error_message,
ref usage,
..
} = message
{
if *stop_reason == StopReason::Error || *stop_reason == StopReason::Aborted {
if *stop_reason == StopReason::Error {
if let Some(ref on_error) = config.on_error {
let err_str = error_message.as_deref().unwrap_or("Unknown error");
on_error(err_str);
}
}
if let Some(ref after_turn) = config.after_turn {
after_turn(&context.messages, usage);
}
tx.send(AgentEvent::TurnEnd {
message: agent_msg,
tool_results: vec![],
})
.ok();
return stats;
}
}
if let Message::Assistant {
stop_reason: StopReason::Refusal,
ref content,
ref usage,
..
} = message
{
let mut tool_results: Vec<Message> = Vec::new();
for c in content {
if let Content::ToolCall {
id,
name,
arguments,
..
} = c
{
tracing::warn!(
tool = name.as_str(),
tool_call_id = id.as_str(),
"tool call not executed: the response was a refusal"
);
let (result, _) = unexecuted_tool_call(
id,
name,
arguments,
REFUSAL_TOOL_RESULT_TEXT.to_string(),
tx,
);
let am: AgentMessage = result.clone().into();
context.messages.push(am.clone());
new_messages.push(am);
tool_results.push(result);
}
}
if let Some(ref after_turn) = config.after_turn {
after_turn(&context.messages, usage);
}
tx.send(AgentEvent::TurnEnd {
message: agent_msg,
tool_results,
})
.ok();
return stats;
}
let tool_calls: Vec<_> = match &message {
Message::Assistant { content, .. } => content
.iter()
.filter_map(|c| match c {
Content::ToolCall {
id,
name,
arguments,
..
} => Some((id.clone(), name.clone(), arguments.clone())),
_ => None,
})
.collect(),
_ => vec![],
};
let mut loop_nudge: Option<AgentMessage> = None;
if !tool_calls.is_empty() {
if let Some(ref mut tracker) = tracker {
let sigs: Vec<(String, serde_json::Value)> = tool_calls
.iter()
.map(|(_, name, args)| (name.clone(), args.clone()))
.collect();
match tracker.record_tool_calls(&sigs) {
context::LoopVerdict::Continue => {}
context::LoopVerdict::Steer {
tool_name,
repetitions,
} => {
warn!(
"loop detection: {tool_name} called {repetitions}x with identical arguments; steering"
);
tx.send(AgentEvent::LoopDetected {
tool_name: tool_name.clone(),
repetitions,
aborted: false,
})
.ok();
let nudge = AgentMessage::Llm(Message::User {
content: vec![Content::Text {
text: format!(
"{LOOP_NUDGE_PREFIX}{tool_name} {repetitions} times with identical arguments. The result will not change — change approach, or say why the repetition is needed.]"
),
}],
timestamp: now_ms(),
});
loop_nudge = Some(nudge);
}
context::LoopVerdict::Abort {
tool_name,
repetitions,
} => {
warn!("loop detection: {tool_name} repeated after steering; stopping");
tx.send(AgentEvent::LoopDetected {
tool_name: tool_name.clone(),
repetitions,
aborted: true,
})
.ok();
let stop = AgentMessage::Llm(Message::User {
content: vec![Content::Text {
text: format!(
"{LOOP_ABORT_PREFIX} {tool_name} was called repeatedly with identical \
arguments after being asked to change approach.]"
),
}],
timestamp: now_ms(),
});
for (id, name, _) in &tool_calls {
let denied = Message::ToolResult {
tool_call_id: id.clone(),
tool_name: name.clone(),
content: vec![Content::Text {
text: "Not executed: the run was stopped for repeating \
this call with identical arguments."
.into(),
}],
is_error: true,
timestamp: now_ms(),
};
let am: AgentMessage = denied.into();
context.messages.push(am.clone());
new_messages.push(am);
}
tx.send(AgentEvent::MessageStart {
message: stop.clone(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: stop.clone(),
})
.ok();
context.messages.push(stop.clone());
new_messages.push(stop);
if let Some(after_turn) = config.after_turn.as_ref() {
let usage = match &message {
Message::Assistant { usage, .. } => usage.clone(),
_ => Usage::default(),
};
after_turn(&context.messages, &usage);
}
tx.send(AgentEvent::TurnEnd {
message: agent_msg,
tool_results: Vec::new(),
})
.ok();
return stats;
}
}
}
}
let has_tool_calls = !tool_calls.is_empty();
let mut tool_results: Vec<Message> = Vec::new();
if has_tool_calls {
exts.sync_events().await;
let execution = execute_tool_calls(
&context.tools,
&tool_calls,
tx,
cancel,
config.get_steering_messages.as_ref(),
&config.tool_execution,
Gate {
middleware: &config.tool_middleware,
history: &context.messages,
extensions: exts,
delegation: &delegation,
},
)
.await;
tool_results = execution.tool_results;
steering_after_tools = execution.steering_messages;
for child in &execution.sub_agent_stats {
stats.sub_agents.record_run(child);
stats.decision.merge(&child.decision);
stats.compaction.merge(&child.compaction);
}
let append_cap = config
.context_config
.as_ref()
.filter(|c| c.truncate_tool_output_on_append);
for result in &tool_results {
let original: AgentMessage = result.clone().into();
let mut am = original.clone();
if let Some(ctx_config) = append_cap {
let (plain, marked) =
context::truncate_tool_output_keyed(original.clone(), ctx_config, None);
am = plain;
if let (Some(sink), false) = (&config.tool_output_sink, marked.is_empty()) {
if let Message::ToolResult { tool_call_id, .. } = result {
let blocks = context::block_texts(&original);
let base = context::tool_output_key(
tool_call_id,
&context::message_text(&original),
);
let mut all_stored = true;
let mut written: Vec<String> = Vec::new();
for idx in &marked {
let Some((_, text)) = blocks.iter().find(|(i, _)| i == idx)
else {
debug_assert!(
false,
"marked block {idx} absent from block_texts"
);
all_stored = false;
break;
};
let key = context::block_key(&base, *idx);
match sink.set(&key, text.clone()).await {
Ok(()) => written.push(key),
Err(e) => {
all_stored = false;
warn!(
tool_call_id = %tool_call_id,
block = idx,
bytes = text.len(),
"could not stash tool output: {e}; no marker in \
this result will offer retrieval"
);
break;
}
}
}
if all_stored {
for key in &written {
if sink.get(key).await.is_none() {
all_stored = false;
warn!(
tool_call_id = %tool_call_id,
key = %key,
stored = written.len(),
"a stashed block did not survive storing its \
siblings — the backend cap cannot hold this \
result; no marker will offer retrieval"
);
break;
}
}
}
if !all_stored {
for key in &written {
if !sink.remove(key).await {
warn!("shared state: rollback could not remove {key}");
}
}
}
if all_stored {
am = context::truncate_tool_output_keyed(
original.clone(),
ctx_config,
Some(&base),
)
.0;
}
}
}
}
context.messages.push(am.clone());
new_messages.push(am);
}
}
if let Some(ref mut tracker) = tracker {
let turn_tokens = match &message {
Message::Assistant { usage, .. } => {
(usage.input + usage.output + usage.cache_read + usage.cache_write) as usize
}
_ => context::message_tokens(&agent_msg),
};
tracker.record_turn(turn_tokens);
}
if let Some(ref after_turn) = config.after_turn {
let usage = match &message {
Message::Assistant { usage, .. } => usage.clone(),
_ => Usage::default(),
};
after_turn(&context.messages, &usage);
}
if let Some(nudge) = loop_nudge.take() {
tx.send(AgentEvent::MessageStart {
message: nudge.clone(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: nudge.clone(),
})
.ok();
context.messages.push(nudge.clone());
new_messages.push(nudge);
}
tx.send(AgentEvent::TurnEnd {
message: agent_msg,
tool_results,
})
.ok();
if let Some(steering) = steering_after_tools.take() {
if !steering.is_empty() {
pending = steering;
continue;
}
}
pending = config
.get_steering_messages
.as_ref()
.map(|f| f())
.unwrap_or_default();
if !has_tool_calls && pending.is_empty() {
break;
}
}
let follow_ups = config
.get_follow_up_messages
.as_ref()
.map(|f| f())
.unwrap_or_default();
if !follow_ups.is_empty() {
pending = follow_ups;
continue;
}
if !exts.is_empty() {
if let Some(failure) = exts.settle().await {
fail_run(&failure, exts, config, tx, context, new_messages, false);
return stats;
}
if let Some(message) =
check_stop(exts, config, tx, context, new_messages, &mut stop_continues).await
{
pending = vec![message];
continue;
}
}
break;
}
stats
}
async fn check_stop(
exts: &crate::extension::ActiveExtensions,
config: &AgentLoopConfig,
tx: &mpsc::UnboundedSender<AgentEvent>,
context: &mut AgentContext,
new_messages: &mut Vec<AgentMessage>,
stop_continues: &mut usize,
) -> Option<AgentMessage> {
use crate::extension::{Failure, StopContext, StopGate, EXTENSION_MESSAGE_PREFIX};
let answer = match context.messages.last() {
Some(AgentMessage::Llm(m @ Message::Assistant { stop_reason, .. }))
if *stop_reason == StopReason::Stop =>
{
m.clone()
}
_ => return None,
};
let stop = StopContext::new(&answer, &context.messages, *stop_continues);
match exts.on_stop(&stop).await {
StopGate::Accept => None,
StopGate::Fail(failure) => {
fail_run(&failure, exts, config, tx, context, new_messages, false);
None
}
StopGate::Continue { messages, required } => {
if *stop_continues >= config.max_stop_continues {
if let Some(required) = required {
let failure = Failure {
name: required,
reason: format!(
"the answer was still not accepted after {} continues",
config.max_stop_continues
),
};
fail_run(&failure, exts, config, tx, context, new_messages, false);
} else {
let names: Vec<&str> = messages.iter().map(|(n, _)| n.as_str()).collect();
warn!(
extensions = ?names,
"extensions still ask to continue after max_stop_continues ({}); \
accepting the answer without their approval",
config.max_stop_continues
);
}
return None;
}
*stop_continues += 1;
let text = messages
.iter()
.map(|(name, message)| format!("{EXTENSION_MESSAGE_PREFIX}{name}] {message}"))
.collect::<Vec<_>>()
.join("\n");
Some(AgentMessage::Llm(Message::user(text)))
}
}
}
async fn stream_assistant_response(
context: &AgentContext,
config: &AgentLoopConfig,
tx: &mpsc::UnboundedSender<AgentEvent>,
cancel: &tokio_util::sync::CancellationToken,
exts: &crate::extension::ActiveExtensions,
) -> Result<Message, crate::extension::ModelHalt> {
let messages = if let Some(transform) = &config.transform_context {
transform(context.messages.clone())
} else {
context.messages.clone()
};
let convert = config.convert_to_llm.as_ref();
let llm_messages = match convert {
Some(f) => f(&messages),
None => default_convert_to_llm(&messages),
};
let tool_defs: Vec<ToolDefinition> = crate::extension::tool_definitions(&context.tools);
let mut llm_messages = llm_messages;
if !exts.is_empty() {
let prompts = run_prompts();
let turn = TurnContext::new(
&context.system_prompt,
&llm_messages,
&tool_defs,
&config.model,
)
.with_run_prompts(&prompts);
match exts.before_model(&turn).await {
Ok(notes) => {
if !notes.is_empty() {
let latest_user = llm_messages.iter_mut().rev().find_map(|m| match m {
Message::User { content, .. } => Some(content),
_ => None,
});
match latest_user {
Some(content) => content.push(Content::Text {
text: notes.join("\n"),
}),
None => {
tracing::debug!("extension note dropped: the request has no user turn")
}
}
}
}
Err(halt) => return Err(halt),
}
}
let retry = &config.retry_config;
let mut attempt = 0;
let result = loop {
let stream_config = StreamConfig {
model: config.model.clone(),
system_prompt: context.system_prompt.clone(),
messages: llm_messages.clone(),
tools: tool_defs.clone(),
thinking_level: config.thinking_level,
api_key: config.api_key.clone(),
max_tokens: config.max_tokens,
temperature: config.temperature,
model_config: config.model_config.clone(),
cache_config: config.cache_config.clone(),
output_schema: config.output_schema.clone(),
};
let (stream_tx, mut stream_rx) = mpsc::unbounded_channel();
let provider_cancel = cancel.clone();
let event_tx = tx.clone();
let model_for_events = config.model.clone();
let forward_handle = crate::rt::spawn(async move {
let mut partial_message: Option<AgentMessage> = None;
let mut ended = false;
while let Some(event) = stream_rx.recv().await {
match &event {
StreamEvent::Start => {
let placeholder = AgentMessage::Llm(Message::Assistant {
content: Vec::new(),
stop_reason: StopReason::Stop,
model: model_for_events.clone(),
provider: String::new(),
usage: Usage::default(),
timestamp: now_ms(),
error_message: None,
});
partial_message = Some(placeholder.clone());
event_tx
.send(AgentEvent::MessageStart {
message: placeholder,
})
.ok();
}
StreamEvent::TextDelta { delta, .. } => {
if let Some(ref msg) = partial_message {
event_tx
.send(AgentEvent::MessageUpdate {
message: msg.clone(),
delta: StreamDelta::Text {
delta: delta.clone(),
},
})
.ok();
}
}
StreamEvent::ThinkingDelta { delta, .. } => {
if let Some(ref msg) = partial_message {
event_tx
.send(AgentEvent::MessageUpdate {
message: msg.clone(),
delta: StreamDelta::Thinking {
delta: delta.clone(),
},
})
.ok();
}
}
StreamEvent::ToolCallDelta { delta, .. } => {
if let Some(ref msg) = partial_message {
event_tx
.send(AgentEvent::MessageUpdate {
message: msg.clone(),
delta: StreamDelta::ToolCallDelta {
delta: delta.clone(),
},
})
.ok();
}
}
StreamEvent::Done { message } => {
let am: AgentMessage = message.clone().into();
if partial_message.is_none() {
event_tx
.send(AgentEvent::MessageStart {
message: am.clone(),
})
.ok();
}
partial_message = Some(am.clone());
ended = true;
event_tx.send(AgentEvent::MessageEnd { message: am }).ok();
}
StreamEvent::Error { message } => {
let am: AgentMessage = message.clone().into();
if partial_message.is_none() {
event_tx
.send(AgentEvent::MessageStart {
message: am.clone(),
})
.ok();
}
partial_message = Some(am.clone());
ended = true;
event_tx.send(AgentEvent::MessageEnd { message: am }).ok();
}
_ => {}
}
}
AttemptEvents {
started: partial_message.is_some(),
ended,
}
});
let result = config
.provider
.stream(stream_config, stream_tx, provider_cancel)
.await;
let attempt_events = match forward_handle.await {
Ok(events) => events,
Err(e) => {
warn!("event forwarder failed; this attempt's events may be incomplete: {e}");
AttemptEvents::default()
}
};
match &result {
Err(e) if e.is_retryable() && attempt < retry.max_retries && !cancel.is_cancelled() => {
if attempt_events.needs_end() {
let failed = error_message(
&config.model,
format!(
"attempt {} of {} failed and will be retried: {e}",
attempt + 1,
retry.max_retries + 1
),
);
tx.send(AgentEvent::MessageEnd {
message: failed.into(),
})
.ok();
}
attempt += 1;
let delay = e
.retry_after()
.map(|d| d.min(std::time::Duration::from_millis(retry.max_delay_ms)))
.unwrap_or_else(|| retry.delay_for_attempt(attempt));
crate::retry::log_retry(attempt, retry.max_retries, &delay, e);
tx.send(AgentEvent::provider_retry(
attempt,
retry.max_retries + 1,
e.to_string(),
delay,
))
.ok();
let backoff = std::pin::pin!(crate::rt::sleep(delay));
let cancelled = std::pin::pin!(cancel.cancelled());
if let futures::future::Either::Right(_) =
futures::future::select(backoff, cancelled).await
{
break (Err(ProviderError::Cancelled), AttemptEvents::default());
}
continue;
}
_ => break (result, attempt_events),
}
};
let (result, attempt_events) = result;
match result {
Ok(msg) => {
if attempt_events.is_silent() {
announce(tx, &msg);
} else if attempt_events.needs_end() {
tx.send(AgentEvent::MessageEnd {
message: msg.clone().into(),
})
.ok();
}
Ok(msg)
}
Err(e) => {
warn!("Provider error: {}", e);
let mut failed = error_message(&config.model, e.to_string());
if matches!(e, ProviderError::Cancelled) || cancel.is_cancelled() {
if let Message::Assistant { stop_reason, .. } = &mut failed {
*stop_reason = StopReason::Aborted;
}
}
if attempt_events.is_silent() {
announce(tx, &failed);
} else if attempt_events.needs_end() {
tx.send(AgentEvent::MessageEnd {
message: failed.clone().into(),
})
.ok();
}
Ok(failed)
}
}
}
fn announce(tx: &mpsc::UnboundedSender<AgentEvent>, message: &Message) {
tx.send(AgentEvent::MessageStart {
message: message.clone().into(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: message.clone().into(),
})
.ok();
}
#[derive(Debug, Default, Clone, Copy)]
struct AttemptEvents {
started: bool,
ended: bool,
}
impl AttemptEvents {
fn needs_end(self) -> bool {
self.started && !self.ended
}
fn is_silent(self) -> bool {
!self.started && !self.ended
}
}
fn error_message(model: &str, error: String) -> Message {
Message::Assistant {
content: vec![Content::Text {
text: String::new(),
}],
stop_reason: StopReason::Error,
model: model.to_string(),
provider: "unknown".into(),
usage: Usage::default(),
timestamp: now_ms(),
error_message: Some(error),
}
}
struct ToolExecutionResult {
tool_results: Vec<Message>,
steering_messages: Option<Vec<AgentMessage>>,
sub_agent_stats: Vec<SessionStats>,
}
fn structured_output_is_tool_forced(model_config: Option<&ModelConfig>) -> bool {
!model_config.is_some_and(|mc| {
mc.api == crate::provider::ApiProtocol::AnthropicMessages
&& mc
.anthropic
.as_ref()
.is_some_and(|c| c.native_structured_output)
})
}
fn unwrap_structured_tool_call(
message: Message,
schema: Option<&crate::provider::OutputSchema>,
) -> Message {
let Some(schema) = schema else {
return message;
};
let Message::Assistant {
content,
model,
provider,
usage,
stop_reason,
error_message,
..
} = &message
else {
return message;
};
let Some(payload) = content.iter().find_map(|c| match c {
Content::ToolCall {
name, arguments, ..
} if *name == schema.name => Some(arguments.clone()),
_ => None,
}) else {
return message;
};
let mut new_content: Vec<Content> = content
.iter()
.filter(|c| !matches!(c, Content::ToolCall { name, .. } if *name == schema.name))
.cloned()
.collect();
new_content.push(Content::Text {
text: payload.to_string(),
});
let new_stop = if *stop_reason == StopReason::ToolUse {
StopReason::Stop
} else {
stop_reason.clone()
};
let rebuilt = Message::assistant(
new_content,
new_stop,
model.clone(),
provider.clone(),
usage.clone(),
);
match error_message {
Some(msg) => rebuilt.with_error_message(msg.clone()),
None => rebuilt,
}
}
#[derive(Clone, Copy)]
struct Gate<'a> {
middleware: &'a [Arc<dyn ToolMiddleware>],
history: &'a [AgentMessage],
extensions: &'a crate::extension::ActiveExtensions,
delegation: &'a Delegation,
}
async fn execute_tool_calls(
tools: &[Box<dyn AgentTool>],
tool_calls: &[(String, String, serde_json::Value)],
tx: &mpsc::UnboundedSender<AgentEvent>,
cancel: &tokio_util::sync::CancellationToken,
get_steering: Option<&GetMessagesFn>,
strategy: &ToolExecutionStrategy,
gate: Gate<'_>,
) -> ToolExecutionResult {
match strategy {
ToolExecutionStrategy::Sequential => {
execute_sequential(tools, tool_calls, tx, cancel, get_steering, gate).await
}
ToolExecutionStrategy::Parallel => {
execute_batch(tools, tool_calls, tx, cancel, get_steering, gate).await
}
ToolExecutionStrategy::Batched { size } => {
let mut results: Vec<Message> = Vec::new();
let mut steering_messages: Option<Vec<AgentMessage>> = None;
let mut sub_agent_stats: Vec<SessionStats> = Vec::new();
for (batch_idx, batch) in tool_calls.chunks((*size).max(1)).enumerate() {
let batch_result = execute_batch(tools, batch, tx, cancel, None, gate).await;
results.extend(batch_result.tool_results);
sub_agent_stats.extend(batch_result.sub_agent_stats);
if let Some(get_steering_fn) = get_steering {
let steering = get_steering_fn();
if !steering.is_empty() {
steering_messages = Some(steering);
let executed = (batch_idx + 1) * *size;
if executed < tool_calls.len() {
for (skip_id, skip_name, _) in &tool_calls[executed..] {
results.push(skip_tool_call(skip_id, skip_name, tx));
}
}
break;
}
}
}
ToolExecutionResult {
tool_results: results,
steering_messages,
sub_agent_stats,
}
}
}
}
async fn execute_sequential(
tools: &[Box<dyn AgentTool>],
tool_calls: &[(String, String, serde_json::Value)],
tx: &mpsc::UnboundedSender<AgentEvent>,
cancel: &tokio_util::sync::CancellationToken,
get_steering: Option<&GetMessagesFn>,
gate: Gate<'_>,
) -> ToolExecutionResult {
let mut results: Vec<Message> = Vec::new();
let mut steering_messages: Option<Vec<AgentMessage>> = None;
let mut sub_agent_stats: Vec<SessionStats> = Vec::new();
for (index, (id, name, args)) in tool_calls.iter().enumerate() {
let (result_msg, delegated) =
execute_single_tool(tools, id, name, args, tx, cancel, gate).await;
results.push(result_msg);
sub_agent_stats.extend(delegated);
if let Some(get_steering_fn) = get_steering {
let steering = get_steering_fn();
if !steering.is_empty() {
steering_messages = Some(steering);
for (skip_id, skip_name, _) in &tool_calls[index + 1..] {
results.push(skip_tool_call(skip_id, skip_name, tx));
}
break;
}
}
}
ToolExecutionResult {
tool_results: results,
steering_messages,
sub_agent_stats,
}
}
async fn execute_batch(
tools: &[Box<dyn AgentTool>],
tool_calls: &[(String, String, serde_json::Value)],
tx: &mpsc::UnboundedSender<AgentEvent>,
cancel: &tokio_util::sync::CancellationToken,
get_steering: Option<&GetMessagesFn>,
gate: Gate<'_>,
) -> ToolExecutionResult {
use futures::future::join_all;
let futures: Vec<_> = tool_calls
.iter()
.map(|(id, name, args)| execute_single_tool(tools, id, name, args, tx, cancel, gate))
.collect();
let batch_results = join_all(futures).await;
let mut sub_agent_stats: Vec<SessionStats> = Vec::new();
let results: Vec<Message> = batch_results
.into_iter()
.map(|(msg, delegated)| {
sub_agent_stats.extend(delegated);
msg
})
.collect();
let steering_messages = if let Some(get_steering_fn) = get_steering {
let steering = get_steering_fn();
if steering.is_empty() {
None
} else {
Some(steering)
}
} else {
None
};
ToolExecutionResult {
tool_results: results,
steering_messages,
sub_agent_stats,
}
}
async fn execute_single_tool(
tools: &[Box<dyn AgentTool>],
id: &str,
name: &str,
args: &serde_json::Value,
tx: &mpsc::UnboundedSender<AgentEvent>,
cancel: &tokio_util::sync::CancellationToken,
gate: Gate<'_>,
) -> (Message, Vec<SessionStats>) {
if let Some(raw) = crate::provider::unparsed_tool_arguments(args) {
let (msg, _) = unparsed_arguments_tool_call(id, name, args, raw, tx);
return (msg, Vec::new());
}
if cancel.is_cancelled() {
return (cancelled_tool_call(id, name, args, tx), Vec::new());
}
if gate.extensions.has_failure() {
return (failed_run_tool_call(id, name, args, tx), Vec::new());
}
let mut effective_args = args.clone();
let prompts = if gate.middleware.is_empty() && gate.extensions.is_empty() {
Vec::new()
} else {
run_prompts()
};
for mw in gate.middleware {
let call = ToolCallRequest {
tool_call_id: id,
tool_name: name,
args: &effective_args,
messages: gate.history,
run_prompts: &prompts,
};
let decision = {
use futures::FutureExt;
std::panic::AssertUnwindSafe(mw.before_tool(&call))
.catch_unwind()
.await
.unwrap_or_else(|_| {
tracing::warn!(
tool = name,
tool_call_id = id,
"tool middleware panicked; denying the call"
);
ToolDecision::Deny("tool middleware panicked".into())
})
};
match decision {
ToolDecision::Allow => {}
ToolDecision::Modify(new_args) => effective_args = new_args,
ToolDecision::Deny(reason) => {
let (msg, _) = denied_tool_call(id, name, &effective_args, &reason, tx);
return (msg, Vec::new());
}
}
}
if !gate.extensions.is_empty() {
let call = ToolCallRequest {
tool_call_id: id,
tool_name: name,
args: &effective_args,
messages: gate.history,
run_prompts: &prompts,
};
match gate.extensions.before_tool(call).await {
Ok(final_args) => effective_args = final_args,
Err(reason) => {
let (msg, _) = denied_tool_call(id, name, &effective_args, &reason, tx);
return (msg, Vec::new());
}
}
}
let args = &effective_args;
if cancel.is_cancelled() {
return (cancelled_tool_call(id, name, args, tx), Vec::new());
}
if gate.extensions.has_failure() {
return (failed_run_tool_call(id, name, args, tx), Vec::new());
}
let tool = tools.iter().find(|t| t.name() == name);
tx.send(AgentEvent::ToolExecutionStart {
tool_call_id: id.to_string(),
tool_name: name.to_string(),
args: args.clone(),
})
.ok();
let withhold = gate.extensions.filters_tool_output();
let on_update: Option<ToolUpdateFn> = if withhold {
None
} else {
let tx = tx.clone();
let id = id.to_string();
let name = name.to_string();
Some(Arc::new(move |partial: ToolResult| {
tx.send(AgentEvent::ToolExecutionUpdate {
tool_call_id: id.clone(),
tool_name: name.clone(),
partial_result: partial,
})
.ok();
}))
};
let on_progress: Option<ProgressFn> = if withhold {
None
} else {
let tx = tx.clone();
let id = id.to_string();
let name = name.to_string();
Some(Arc::new(move |text: String| {
tx.send(AgentEvent::ProgressMessage {
tool_call_id: id.clone(),
tool_name: name.clone(),
text,
})
.ok();
}))
};
let sub_agent_report: SubAgentReport = Arc::default();
let ctx = ToolContext {
tool_call_id: id.to_string(),
tool_name: name.to_string(),
cancel: cancel.child_token(),
on_update,
on_progress,
sub_agent_report: Some(sub_agent_report.clone()),
delegation: Delegation {
call_id: Some(id.to_string()),
..gate.delegation.clone()
},
};
let tool_span = tracing::info_span!(
"tool",
tool = %name,
tool_call_id = %id,
is_error = tracing::field::Empty,
);
use tracing::Instrument;
let (result, is_error) = match tool {
Some(tool) => {
let execution = {
use futures::FutureExt;
std::panic::AssertUnwindSafe(async { tool.execute(args.clone(), ctx).await })
.catch_unwind()
.instrument(tool_span.clone())
.await
.unwrap_or_else(|payload| {
let why = crate::tool_source::panic_message(&*payload);
tracing::error!(tool = name, tool_call_id = id, panic = %why, "tool panicked");
Err(ToolError::Failed(format!("tool '{name}' panicked: {why}")))
})
};
match execution {
Ok(r) => (r, false),
Err(e) => (
ToolResult {
content: vec![Content::Text {
text: e.to_string(),
}],
details: serde_json::Value::Null,
},
true,
),
}
}
None => (
ToolResult {
content: vec![Content::Text {
text: format!("Tool {} not found", name),
}],
details: serde_json::Value::Null,
},
true,
),
};
tool_span.record("is_error", is_error);
let delegated =
std::mem::take(&mut *sub_agent_report.lock().unwrap_or_else(|e| e.into_inner()));
let mut result = result;
let mut is_error = is_error;
if tool.is_some() && !gate.extensions.is_empty() {
let call = ToolCallRequest {
tool_call_id: id,
tool_name: name,
args,
messages: gate.history,
run_prompts: &prompts,
};
let mut output = crate::extension::ToolOutput::new(result, is_error);
gate.extensions.after_tool(&call, &mut output).await;
result = output.result;
is_error = output.is_error;
}
if let Some((first, rest)) = delegated.split_first() {
let mut combined = first.clone();
for stats in rest {
combined.merge(stats);
}
attach_sub_agent_stats(&mut result.details, &combined);
}
tx.send(AgentEvent::ToolExecutionEnd {
tool_call_id: id.to_string(),
tool_name: name.to_string(),
result: result.clone(),
is_error,
})
.ok();
let tool_result_msg = Message::ToolResult {
tool_call_id: id.to_string(),
tool_name: name.to_string(),
content: result.content,
is_error,
timestamp: now_ms(),
};
tx.send(AgentEvent::MessageStart {
message: tool_result_msg.clone().into(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: tool_result_msg.clone().into(),
})
.ok();
(tool_result_msg, delegated)
}
fn attach_sub_agent_stats(details: &mut serde_json::Value, stats: &SessionStats) {
if details.is_null() {
*details = serde_json::json!({});
}
if let Some(obj) = details.as_object_mut() {
if !obj.contains_key(SUB_AGENT_STATS_KEY) {
if let Ok(v) = serde_json::to_value(stats) {
obj.insert(SUB_AGENT_STATS_KEY.to_string(), v);
}
}
}
}
fn denied_tool_call(
id: &str,
name: &str,
args: &serde_json::Value,
reason: &str,
tx: &mpsc::UnboundedSender<AgentEvent>,
) -> (Message, bool) {
tracing::warn!(
tool = name,
tool_call_id = id,
reason,
"tool call denied by middleware"
);
unexecuted_tool_call(id, name, args, format!("Tool call denied: {}", reason), tx)
}
fn unparsed_arguments_tool_call(
id: &str,
name: &str,
args: &serde_json::Value,
raw: &str,
tx: &mpsc::UnboundedSender<AgentEvent>,
) -> (Message, bool) {
let text = match serde_json::from_str::<serde_json::Value>(raw) {
Err(e) => {
tracing::warn!(
tool = name,
tool_call_id = id,
len = raw.len(),
parse_error = %e,
"tool call not executed: arguments did not parse as JSON"
);
if e.is_eof() {
format!(
"The arguments for tool `{name}` were cut off before they were complete \
(the response likely hit the output token limit): {e}. The tool was not \
run. Do not resend the same call unchanged — make the arguments smaller \
(for example, split large content across several calls)."
)
} else {
format!(
"The arguments for tool `{name}` are not valid JSON: {e}. The tool was not \
run. Send the arguments as a single valid JSON object."
)
}
}
Ok(v) => {
let kind = match v {
serde_json::Value::String(_) => "a string",
serde_json::Value::Number(_) => "a number",
serde_json::Value::Array(_) => "an array",
serde_json::Value::Bool(_) => "a boolean",
serde_json::Value::Null | serde_json::Value::Object(_) => "",
};
tracing::warn!(
tool = name,
tool_call_id = id,
kind,
"tool call not executed: arguments were not delivered as a JSON object"
);
if kind.is_empty() {
format!(
"The arguments for tool `{name}` were not delivered as a parsed JSON \
object. The tool was not run."
)
} else {
format!(
"The arguments for tool `{name}` were not a JSON object (got {kind}). \
The tool was not run. Send the arguments as a single JSON object — \
not encoded as a string."
)
}
}
};
unexecuted_tool_call(id, name, args, text, tx)
}
fn failed_run_tool_call(
id: &str,
name: &str,
args: &serde_json::Value,
tx: &mpsc::UnboundedSender<AgentEvent>,
) -> Message {
tracing::debug!(
tool = name,
tool_call_id = id,
"tool call not executed: a required extension failed the run"
);
unexecuted_tool_call(
id,
name,
args,
"Tool call not run: a required extension failed the run.".to_string(),
tx,
)
.0
}
fn cancelled_tool_call(
id: &str,
name: &str,
args: &serde_json::Value,
tx: &mpsc::UnboundedSender<AgentEvent>,
) -> Message {
tracing::debug!(
tool = name,
tool_call_id = id,
"tool call not executed: the run was cancelled"
);
unexecuted_tool_call(id, name, args, CANCELLED_TOOL_RESULT_TEXT.to_string(), tx).0
}
fn unexecuted_tool_call(
id: &str,
name: &str,
args: &serde_json::Value,
text: String,
tx: &mpsc::UnboundedSender<AgentEvent>,
) -> (Message, bool) {
tx.send(AgentEvent::ToolExecutionStart {
tool_call_id: id.to_string(),
tool_name: name.to_string(),
args: args.clone(),
})
.ok();
let result = ToolResult {
content: vec![Content::Text { text }],
details: serde_json::Value::Null,
};
tx.send(AgentEvent::ToolExecutionEnd {
tool_call_id: id.to_string(),
tool_name: name.to_string(),
result: result.clone(),
is_error: true,
})
.ok();
let msg = Message::ToolResult {
tool_call_id: id.to_string(),
tool_name: name.to_string(),
content: result.content,
is_error: true,
timestamp: now_ms(),
};
tx.send(AgentEvent::MessageStart {
message: msg.clone().into(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: msg.clone().into(),
})
.ok();
(msg, true)
}
fn skip_tool_call(
tool_call_id: &str,
tool_name: &str,
tx: &mpsc::UnboundedSender<AgentEvent>,
) -> Message {
let result = ToolResult {
content: vec![Content::Text {
text: "Skipped due to queued user message.".into(),
}],
details: serde_json::Value::Null,
};
tx.send(AgentEvent::ToolExecutionStart {
tool_call_id: tool_call_id.into(),
tool_name: tool_name.into(),
args: serde_json::Value::Null,
})
.ok();
tx.send(AgentEvent::ToolExecutionEnd {
tool_call_id: tool_call_id.into(),
tool_name: tool_name.into(),
result: result.clone(),
is_error: true,
})
.ok();
let msg = Message::ToolResult {
tool_call_id: tool_call_id.into(),
tool_name: tool_name.into(),
content: result.content,
is_error: true,
timestamp: now_ms(),
};
tx.send(AgentEvent::MessageStart {
message: msg.clone().into(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: msg.clone().into(),
})
.ok();
msg
}