use af_context::ToolCallId;
use std::time::{Duration, Instant};
use af_agent::{HookDecision, ToolConcurrency, ToolExecutionContext, ToolExecutionFuture};
use af_agent_session::{Event, InteractionResolution};
use futures::future::join_all;
use serde_json::Value;
use crate::{
extension::await_extension, AgentRuntime, CancellationToken, EventWriter, RuntimeError,
TurnRequest,
};
#[derive(Clone)]
pub(super) struct PlannedCall {
pub(super) transcript_id: ToolCallId,
pub(super) id: ToolCallId,
pub(super) name: String,
pub(super) arguments: Value,
pub(super) step: u32,
pub(super) source_event_seq: u64,
pub(super) preflight_error: Option<String>,
}
impl AgentRuntime {
pub(super) async fn authorize_call(
&self,
request: &TurnRequest,
call: &PlannedCall,
interaction_resolution: Option<InteractionResolution>,
cancellation: CancellationToken,
) -> HookDecision {
let timeout = self
.tools
.get(&call.name)
.map_or(Duration::from_secs(15), |tool| {
Duration::from_secs(tool.meta().timeout_secs.max(1))
});
let extension_cancellation = cancellation.child();
let deadline = Instant::now() + timeout;
let context = ToolExecutionContext {
request: request.context.clone(),
session_id: request.session_id.clone(),
run_id: request.run_id.clone(),
step: call.step,
call_id: call.id.clone(),
source_event_seq: call.source_event_seq,
interaction_resolution,
cancellation: extension_cancellation.clone(),
deadline,
};
for hook in &self.hooks {
let decision = match await_extension(
"before_tool hook",
extension_cancellation.clone(),
deadline,
hook.before_tool(&context, &call.name, &call.arguments),
)
.await
{
Ok(decision) => decision,
Err(error) => {
return HookDecision::Deny {
reason: error.to_string(),
}
}
};
if decision != HookDecision::Continue {
return decision;
}
}
HookDecision::Continue
}
pub(super) async fn execute_tools(
&self,
request: &TurnRequest,
calls: &[PlannedCall],
writer: &dyn EventWriter,
cancellation: CancellationToken,
interaction_resolution: Option<InteractionResolution>,
) -> Result<Vec<Result<Value, String>>, RuntimeError> {
let mut results = Vec::with_capacity(calls.len());
let mut index = 0;
while index < calls.len() {
let call = &calls[index];
let exclusive = self
.tools
.get(&call.name)
.is_some_and(|tool| tool.meta().concurrency == ToolConcurrency::Exclusive);
let batch_len = if exclusive {
1
} else {
calls[index..]
.iter()
.take_while(|candidate| {
self.tools.get(&candidate.name).is_some_and(|tool| {
tool.meta().concurrency == ToolConcurrency::Concurrent
})
})
.count()
.min(self.limits.max_parallel_tools.max(1))
};
if batch_len == 0 {
results.push(Err(format!("tool '{}' is unavailable", call.name)));
index += 1;
continue;
}
writer
.append(
calls[index..index + batch_len]
.iter()
.map(|call| Event::ToolExecutionStarted {
run_id: request.run_id.clone(),
step: call.step,
call_id: call.id.clone(),
})
.collect(),
)
.await?;
let futures = calls[index..index + batch_len].iter().map(|call| {
let tools = self.tools.clone();
let hooks = self.hooks.clone();
let name = call.name.clone();
let args = call.arguments.clone();
let run_cancellation = cancellation.clone();
let call_cancellation = run_cancellation.child();
let timeout = tools.get(&name).map_or(Duration::from_secs(15), |tool| {
Duration::from_secs(tool.meta().timeout_secs.max(1))
});
let context = ToolExecutionContext {
request: request.context.clone(),
session_id: request.session_id.clone(),
run_id: request.run_id.clone(),
step: call.step,
call_id: call.id.clone(),
source_event_seq: call.source_event_seq,
interaction_resolution,
cancellation: call_cancellation.clone(),
deadline: Instant::now() + timeout,
};
async move {
if run_cancellation.is_cancelled() {
return Err("cancelled".into());
}
let mut execution: ToolExecutionFuture<'_> =
Box::pin(tools.execute_with_context(&name, &context, args));
for hook in hooks.iter().rev() {
execution = hook.around_tool(&context, &name, &call.arguments, execution);
}
tokio::pin!(execution);
let deadline = tokio::time::sleep(timeout);
tokio::pin!(deadline);
let interrupted = tokio::select! {
result = &mut execution => {
let value = result?;
tools.validate_output(&name, &value)?;
return Ok(value);
},
_ = &mut deadline => format!("tool timed out after {}s", timeout.as_secs()),
_ = run_cancellation.cancelled() => "cancelled".to_string(),
};
call_cancellation.cancel();
match tokio::time::timeout(Duration::from_millis(250), &mut execution).await {
Ok(result) => {
let value = result?;
tools.validate_output(&name, &value)?;
Ok(value)
}
Err(_) => Err(format!("{interrupted}; tool_outcome_unknown")),
}
}
});
let batch = join_all(futures).await;
writer
.append(
calls[index..index + batch_len]
.iter()
.zip(&batch)
.flat_map(|(call, result)| {
let (result, is_error) = match result {
Ok(value) => (value.clone(), false),
Err(error) => (serde_json::json!({"error":error}), true),
};
let cost_units = self
.tools
.get(&call.name)
.map_or(0, |tool| tool.meta().cost_units);
[
Event::ToolResult {
run_id: request.run_id.clone(),
step: call.step,
call_id: call.id.clone(),
result,
is_error,
},
Event::UsageRecorded {
run_id: request.run_id.clone(),
operation_id: format!("tool:{}", call.id),
prompt_tokens: 0,
completion_tokens: 0,
cost_units,
},
]
})
.collect(),
)
.await?;
results.extend(batch);
index += batch_len;
}
Ok(results)
}
pub(super) async fn run_after_hooks(
&self,
request: &TurnRequest,
call: &PlannedCall,
result: &Value,
interaction_resolution: Option<InteractionResolution>,
cancellation: CancellationToken,
) -> Result<(), String> {
let timeout = self
.tools
.get(&call.name)
.map_or(Duration::from_secs(15), |tool| {
Duration::from_secs(tool.meta().timeout_secs.max(1))
});
let extension_cancellation = cancellation.child();
let deadline = Instant::now() + timeout;
let context = ToolExecutionContext {
request: request.context.clone(),
session_id: request.session_id.clone(),
run_id: request.run_id.clone(),
step: call.step,
call_id: call.id.clone(),
source_event_seq: call.source_event_seq,
interaction_resolution,
cancellation: extension_cancellation.clone(),
deadline,
};
for hook in &self.hooks {
await_extension(
"after_tool hook",
extension_cancellation.clone(),
deadline,
hook.after_tool(&context, &call.name, &call.arguments, result),
)
.await
.map_err(|error| error.to_string())?
.map_err(|error| error.to_string())?;
}
Ok(())
}
}