Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

feat: add opt-in early tool call dispatch

Browse Source
This commit is contained in:
Hare
2026-09-29 11:07:59 +09:00
parent d1303ab638
commit cc2f8b3422
15 changed files with 1433 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/src/engine.rs
+727 -24
View File
@@ -1,7 +1,8 @@
use std::collections::{HashMap, HashSet};
use std::collections::{HashMap, HashSet, VecDeque};
use std::{future::Future, marker::PhantomData, pin::Pin, sync::Arc, time::Instant};
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use serde_json::{Value, json};
use tokio::sync::mpsc;
use tokio::time::Instant as TokioInstant;
@@ -25,7 +26,7 @@ use crate::{
llm_client::{
ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream,
ToolDefinition, error::is_retryable, event::Event, retry::RetryPolicy,
transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT, types::parse_tool_arguments,
transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT,
},
state::{EngineState, Locked, Mutable},
timeline::event::{ErrorEvent, StatusEvent, UsageEvent},
@@ -78,11 +79,20 @@ pub enum ToolRegistryError {
DuplicateName(String),
}
/// When provider-confirmed tool calls may begin execution.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum ToolCallDispatchMode {
/// Wait for the complete assistant response before starting any tools.
#[default]
AfterResponse,
/// Start calls at trustworthy provider per-call completion boundaries.
OnToolCallComplete,
}
/// Engine configuration
#[derive(Debug, Clone, Default)]
pub struct EngineConfig {
// Reserved for future extensions (currently empty)
_private: (),
pub tool_call_dispatch: ToolCallDispatchMode,
}
/// Project terminal tool outputs into the assistant's original ToolCall order.
@@ -196,6 +206,7 @@ pub struct EngineRunOutput<C: LlmClient, A: Send + Sync = ()> {
}
/// Internal: tool execution result
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ToolExecutionResult {
Completed,
Paused,
@@ -207,6 +218,56 @@ struct ToolExecutionAttempt {
terminal: bool,
}
type ToolTerminalTask = Pin<Box<dyn Future<Output = (String, ToolResult)> + Send>>;
type ToolCallExecutionInfo = (
ToolCall,
crate::tool::ToolMeta,
Arc<dyn crate::tool::Tool>,
ToolExecutionContext,
);
/// One response-scoped execution batch whose calls are admitted as their
/// provider-confirmed blocks complete.
struct EarlyToolExecutionBatch {
batch_id: String,
next_call_index: usize,
calls: Vec<ToolCall>,
seen_call_ids: HashSet<String>,
call_info_map: HashMap<String, ToolCallExecutionInfo>,
started_calls: Vec<(String, String)>,
attempt_fence: ToolExecutionAttemptFence,
futures: FuturesUnordered<ToolTerminalTask>,
execution_handles: HashMap<String, ToolExecutionHandle>,
synthetic_results: VecDeque<ToolResult>,
terminal_call_ids: HashSet<String>,
dispatched_any: bool,
}
impl EarlyToolExecutionBatch {
fn new(batch_id: String) -> Self {
Self {
batch_id,
next_call_index: 0,
calls: Vec::new(),
seen_call_ids: HashSet::new(),
call_info_map: HashMap::new(),
started_calls: Vec::new(),
attempt_fence: ToolExecutionAttemptFence::default(),
futures: FuturesUnordered::new(),
execution_handles: HashMap::new(),
synthetic_results: VecDeque::new(),
terminal_call_ids: HashSet::new(),
dispatched_any: false,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum EarlyToolAdmission {
Continue,
Pause,
}
/// Per-batch compare-and-set fence for terminal ToolResult commits.
///
/// A completion may commit only when its attempt id still matches the active
@@ -312,6 +373,11 @@ enum StreamCompletion {
Interrupted { reason: String },
}
struct StreamResponseOutput {
completion: StreamCompletion,
early_tools: Option<EarlyToolExecutionBatch>,
}
pub struct Engine<C: LlmClient, S: EngineState = Mutable, A: Send + Sync = ()> {
/// LLM client
client: C,
@@ -362,6 +428,8 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable, A: Send + Sync = ()> {
tool_execution_batch_count: usize,
/// Maximum number of AgentTurns (None = unlimited)
max_turns: Option<u32>,
/// Selected tool-call execution start boundary.
tool_call_dispatch_mode: ToolCallDispatchMode,
/// Caller-selected policy for interrupting started provider operations.
tool_execution_policy: ToolExecutionPolicy,
/// AgentTurn-start callbacks (1:1 with LlmCall today)
@@ -1184,7 +1252,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
} = item
{
if !answered_call_ids.contains(call_id) {
let input = parse_tool_arguments(arguments);
let input = serde_json::from_str::<Value>(arguments)
.ok()
.filter(Value::is_object)
.unwrap_or(Value::Null);
pending_calls.push(ToolCall {
id: call_id.clone(),
name: name.clone(),
@@ -1201,6 +1272,342 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}
}
async fn admit_early_tool_calls(
&mut self,
history: &History<A>,
batch: &mut EarlyToolExecutionBatch,
tool_calls: Vec<ToolCall>,
) -> Result<EarlyToolAdmission, EngineError> {
for mut tool_call in tool_calls {
if !batch.seen_call_ids.insert(tool_call.id.clone()) {
warn!(call_id = %tool_call.id, "Ignoring duplicate completed tool-call event");
continue;
}
let call_index = batch.next_call_index;
batch.next_call_index += 1;
batch.calls.push(tool_call.clone());
let expected_tool_use_id = tool_call.id.clone();
let context = ToolExecutionContext::new(&tool_call.id, &batch.batch_id, call_index);
if !tool_call.input.is_object() {
batch.synthetic_results.push_back(ToolResult::error(
&tool_call.id,
"Tool arguments must be one complete JSON object",
));
batch.dispatched_any = true;
continue;
}
let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) else {
let attempt_id = context.batch_id.clone();
batch
.attempt_fence
.register(tool_call.id.clone(), attempt_id.clone());
batch
.started_calls
.push((tool_call.id.clone(), attempt_id.clone()));
let call_id = tool_call.id.clone();
let name = tool_call.name.clone();
batch.futures.push(Box::pin(async move {
(
attempt_id,
ToolResult::error(&call_id, format!("Tool not found: {name}")),
)
}));
batch.dispatched_any = true;
continue;
};
let invocation = self.interceptor_invocation(
InterceptorPhase::PreToolCall,
Some(self.turn_count),
Some(InterceptorCallId::Tool(expected_tool_use_id.clone())),
call_index,
);
let mut info = ToolCallInfo {
invocation,
history: history.entries(),
call: tool_call.clone(),
meta,
tool,
context,
};
let action = self
.interceptor
.pre_tool_call(&mut info)
.await
.map_err(|error| {
EngineError::from(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
error,
))
})?;
if info.call.id != expected_tool_use_id {
return Err(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
InterceptorError::new(
InterceptorErrorCategory::ContractViolation,
"pre-tool interceptor changed immutable tool call identity",
),
)
.into());
}
tool_call = info.call;
let mut context = info.context;
context.call_id = tool_call.id.clone();
batch.call_info_map.insert(
tool_call.id.clone(),
(
tool_call.clone(),
info.meta.clone(),
info.tool.clone(),
context.clone(),
),
);
match action {
PreToolAction::Continue => {
let attempt_id = context.batch_id.clone();
batch
.attempt_fence
.register(tool_call.id.clone(), attempt_id.clone());
batch
.started_calls
.push((tool_call.id.clone(), attempt_id.clone()));
let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default();
let call_id = tool_call.id.clone();
let (handle, terminal) =
ToolExecutionHandle::start(info.tool, input_json, context);
batch.execution_handles.insert(call_id.clone(), handle);
batch.futures.push(Box::pin(async move {
let result = match terminal.await {
ToolExecutionTerminal::Confirmed(Ok(output)) => {
ToolResult::from_output(&call_id, output)
}
ToolExecutionTerminal::Confirmed(Err(ToolError::Cancelled(output))) => {
ToolResult::from_output_with_disposition(
&call_id,
output,
ToolResultDisposition::Cancelled,
)
}
ToolExecutionTerminal::Confirmed(Err(ToolError::Interrupted(
output,
))) => ToolResult::from_output_with_disposition(
&call_id,
output,
ToolResultDisposition::Interrupted,
),
ToolExecutionTerminal::Confirmed(Err(error)) => {
ToolResult::error(&call_id, error.to_string())
}
ToolExecutionTerminal::OutcomeUnknown => {
ToolResult::outcome_unknown(&call_id)
}
};
(attempt_id, result)
}));
batch.dispatched_any = true;
}
PreToolAction::SyntheticResult(result) => {
if result.tool_use_id != expected_tool_use_id {
return Err(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
InterceptorError::new(
InterceptorErrorCategory::ContractViolation,
"synthetic tool result changed immutable tool call identity",
),
)
.into());
}
batch.synthetic_results.push_back(result);
batch.dispatched_any = true;
}
PreToolAction::Skip => {}
PreToolAction::Abort(reason) => return Err(EngineError::Aborted(reason)),
PreToolAction::Pause => return Ok(EarlyToolAdmission::Pause),
}
}
Ok(EarlyToolAdmission::Continue)
}
async fn cancel_and_terminalize_early_batch(
&mut self,
history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
batch: &mut EarlyToolExecutionBatch,
) -> Option<EngineError> {
use futures::StreamExt as _;
let cancellation_request_deadline =
TokioInstant::now() + self.tool_execution_policy.cancellation_request_timeout;
let cancellation_requests = batch
.execution_handles
.iter()
.filter(|(call_id, _)| !batch.terminal_call_ids.contains(*call_id))
.map(|(call_id, handle)| {
let call_id = call_id.clone();
let handle = handle.clone();
async move {
(
call_id,
handle.cancel_before(cancellation_request_deadline).await,
)
}
});
let cancellation_requests: FuturesUnordered<_> = cancellation_requests.collect();
for (call_id, result) in cancellation_requests.collect::<Vec<_>>().await {
if let Err(error) = result {
warn!(%call_id, error = %error, "Tool cooperative cancellation request failed");
}
}
let mut batch_error = None;
let deadline =
TokioInstant::now() + self.tool_execution_policy.terminal_confirmation_timeout;
while !batch.futures.is_empty() {
tokio::select! {
biased;
result = batch.futures.next() => {
let (attempt_id, result) = result.expect("non-empty FuturesUnordered");
if let Err(error) = self.finalize_and_commit_tool_result(
history,
annotate,
result,
Some(&attempt_id),
&batch.call_info_map,
&mut batch.attempt_fence,
&mut batch.terminal_call_ids,
).await && batch_error.is_none() {
batch_error = Some(error);
}
}
_ = tokio::time::sleep_until(deadline) => break,
}
}
for (call_id, attempt_id) in &batch.started_calls {
if !batch.attempt_fence.is_terminal(call_id) {
if let Some(handle) = batch.execution_handles.get(call_id) {
handle.force_close();
}
if let Err(error) = self
.finalize_and_commit_tool_result(
history,
annotate,
ToolResult::outcome_unknown(call_id),
Some(attempt_id),
&batch.call_info_map,
&mut batch.attempt_fence,
&mut batch.terminal_call_ids,
)
.await
&& batch_error.is_none()
{
batch_error = Some(error);
}
}
}
batch_error
}
async fn drain_early_tool_batch(
&mut self,
history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
mut batch: EarlyToolExecutionBatch,
force_cancel: bool,
initial_pause_requested: bool,
) -> Result<ToolExecutionResult, EngineError> {
while let Some(result) = batch.synthetic_results.pop_front() {
if let Err(error) = self
.finalize_and_commit_tool_result(
history,
annotate,
result,
None,
&batch.call_info_map,
&mut batch.attempt_fence,
&mut batch.terminal_call_ids,
)
.await
{
let cleanup_error = self
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
.await;
return Err(cleanup_error.unwrap_or(error));
}
}
if force_cancel {
if let Some(error) = self
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
.await
{
return Err(error);
}
return Ok(ToolExecutionResult::Completed);
}
let mut pause_requested = initial_pause_requested;
let mut pause_deadline = initial_pause_requested
.then(|| TokioInstant::now() + self.tool_execution_policy.pause_safe_boundary_timeout);
while !batch.futures.is_empty() {
tokio::select! {
biased;
result = batch.futures.next() => {
let (attempt_id, result) = result.expect("non-empty FuturesUnordered");
if let Err(error) = self.finalize_and_commit_tool_result(
history,
annotate,
result,
Some(&attempt_id),
&batch.call_info_map,
&mut batch.attempt_fence,
&mut batch.terminal_call_ids,
).await {
let cleanup_error = self
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
.await;
return Err(cleanup_error.unwrap_or(error));
}
}
pause = self.pause_rx.recv(), if !pause_requested => {
if pause.is_some() {
pause_requested = true;
pause_deadline = Some(
TokioInstant::now()
+ self.tool_execution_policy.pause_safe_boundary_timeout,
);
}
}
_ = tokio::time::sleep_until(pause_deadline.unwrap_or_else(TokioInstant::now)), if pause_deadline.is_some() => {
if let Some(error) = self
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
.await {
return Err(error);
}
return Ok(ToolExecutionResult::Paused);
}
cancel = self.cancel_rx.recv() => {
if cancel.is_some() {
info!("Tool execution cancellation requested");
}
if let Some(error) = self
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
.await {
return Err(error);
}
return Err(EngineError::Cancelled);
}
}
}
Ok(if pause_requested {
ToolExecutionResult::Paused
} else {
ToolExecutionResult::Completed
})
}
/// Execute tools in parallel
///
/// After running pre_tool_call hooks on all tools,
@@ -1231,6 +1638,13 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
for (call_index, mut tool_call) in tool_calls.into_iter().enumerate() {
let expected_tool_use_id = tool_call.id.clone();
let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index);
if !tool_call.input.is_object() {
synthetic_results.push(ToolResult::error(
&tool_call.id,
"Tool arguments must be one complete JSON object",
));
continue;
}
if let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) {
let invocation = self.interceptor_invocation(
InterceptorPhase::PreToolCall,
@@ -1934,8 +2348,30 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
self.request_trace_payload(&request),
);
let request = self.attach_transport_trace(request, current_turn, current_llm_call);
let stream_outcome = self
.stream_response(request, current_turn, current_llm_call)
let early_dispatch = self.tool_call_dispatch_mode
== ToolCallDispatchMode::OnToolCallComplete
&& self.client.tool_call_completion_support()
== crate::llm_client::ToolCallCompletionSupport::PerBlock;
if self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete
&& !early_dispatch
{
self.emit_warning(
"tool-call early dispatch deferred until response completion: provider does not expose a trustworthy per-call completion boundary",
);
}
let assistant_start = history.len();
let StreamResponseOutput {
completion: stream_outcome,
mut early_tools,
} = self
.stream_response(
request,
current_turn,
current_llm_call,
history,
annotate,
early_dispatch,
)
.await?;
for cb in &self.llm_call_end_cbs {
@@ -1944,6 +2380,47 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
self.llm_call_count += 1;
if let StreamCompletion::Interrupted { reason } = stream_outcome {
if early_tools
.as_ref()
.is_some_and(|batch| batch.dispatched_any)
{
// Once a side effect may have started, automatic stream
// continuation would allow the model to regenerate the same
// operation under a fresh call id. Close every owned attempt
// and stop this run instead.
self.timeline.abort_current_block();
self.timeline.flush_usage();
let reasoning_items = self.thinking_block_collector.take_collected();
let text_blocks = self.text_block_collector.take_collected();
let _partial_tool_calls = self.tool_call_collector.take_collected();
let assistant_items =
self.build_assistant_items(&reasoning_items, &text_blocks, &[]);
if !assistant_items.is_empty()
&& let Err(error) =
self.append_history_items(history, assistant_items, annotate)
{
if let Some(batch) = early_tools.take()
&& let Err(cleanup_error) = self
.drain_early_tool_batch(history, annotate, batch, true, false)
.await
{
return Err(cleanup_error);
}
return Err(error);
}
if let Some(batch) = early_tools.take() {
self.drain_early_tool_batch(history, annotate, batch, true, false)
.await?;
}
return Err(EngineError::Client(ClientError::Api {
status: None,
code: Some("early_tool_stream_interrupted".to_string()),
message: format!(
"LLM stream interrupted after early tool execution started: {reason}"
),
retry_after: None,
}));
}
stream_continuations += 1;
if stream_continuations > MAX_STREAM_CONTINUATIONS {
return Err(EngineError::Client(ClientError::Api {
@@ -1989,11 +2466,61 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
// `append_history_items` so observers see each item as it lands.
let reasoning_items = self.thinking_block_collector.take_collected();
let text_blocks = self.text_block_collector.take_collected();
let tool_calls = self.tool_call_collector.take_collected();
let remaining_tool_calls = self.tool_call_collector.take_collected();
let assistant_items =
self.build_assistant_items(&reasoning_items, &text_blocks, &tool_calls);
let assistant_start = history.len();
self.append_history_items(history, assistant_items, annotate)?;
self.build_assistant_items(&reasoning_items, &text_blocks, &remaining_tool_calls);
if let Err(error) = self.append_history_items(history, assistant_items, annotate) {
if let Some(batch) = early_tools.take()
&& let Err(cleanup_error) = self
.drain_early_tool_batch(history, annotate, batch, true, false)
.await
{
return Err(cleanup_error);
}
return Err(error);
}
if early_dispatch && !remaining_tool_calls.is_empty() {
if early_tools.is_none() {
let batch_id = format!("tool-batch-{}", self.tool_execution_batch_count);
self.tool_execution_batch_count += 1;
early_tools = Some(EarlyToolExecutionBatch::new(batch_id));
}
let admission = self
.admit_early_tool_calls(
history,
early_tools.as_mut().expect("batch initialized"),
remaining_tool_calls.clone(),
)
.await;
let admission = match admission {
Ok(admission) => admission,
Err(error) => {
if let Some(batch) = early_tools.take()
&& let Err(cleanup_error) = self
.drain_early_tool_batch(history, annotate, batch, true, false)
.await
{
return Err(cleanup_error);
}
return Err(error);
}
};
if admission == EarlyToolAdmission::Pause {
if let Some(batch) = early_tools.take() {
self.drain_early_tool_batch(history, annotate, batch, false, true)
.await?;
}
return Ok(EngineResult::Paused);
}
}
let mut tool_calls = early_tools
.as_ref()
.map(|batch| batch.calls.clone())
.unwrap_or_default();
if !early_dispatch {
tool_calls = remaining_tool_calls;
}
let assistant_invocation = self.interceptor_invocation(
InterceptorPhase::AssistantTurnEnd,
@@ -2011,29 +2538,64 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
history: history.entries(),
tool_calls: &tool_calls,
})
.await
.map_err(|error| {
EngineError::from(InterceptorFailure::new(
.await;
let assistant_turn_action = match assistant_turn_action {
Ok(action) => action,
Err(error) => {
let error = EngineError::from(InterceptorFailure::new(
InterceptorPhase::AssistantTurnEnd,
error,
))
})?;
));
if let Some(batch) = early_tools.take()
&& let Err(cleanup_error) = self
.drain_early_tool_batch(history, annotate, batch, true, false)
.await
{
return Err(cleanup_error);
}
return Err(error);
}
};
match assistant_turn_action {
TurnEndAction::Finish if tool_calls.is_empty() => {
return Ok(EngineResult::Finished);
}
TurnEndAction::Finish => {}
TurnEndAction::ContinueWithMessages(additional) => {
self.append_history_items(history, additional, annotate)?;
if let Err(error) = self.append_history_items(history, additional, annotate) {
if let Some(batch) = early_tools.take()
&& let Err(cleanup_error) = self
.drain_early_tool_batch(history, annotate, batch, true, false)
.await
{
return Err(cleanup_error);
}
return Err(error);
}
if tool_calls.is_empty() {
continue;
}
}
TurnEndAction::Pause => {
if let Some(batch) = early_tools.take() {
self.drain_early_tool_batch(history, annotate, batch, false, true)
.await?;
}
return Ok(EngineResult::Paused);
}
}
if let Some(batch) = early_tools.take() {
if self
.drain_early_tool_batch(history, annotate, batch, false, false)
.await?
== ToolExecutionResult::Paused
{
return Ok(EngineResult::Paused);
}
continue;
}
if let Some(result) = self
.execute_and_commit_tools(history, annotate, tool_calls)
.await?
@@ -2224,7 +2786,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
request: Request,
turn: usize,
llm_call: usize,
) -> Result<StreamCompletion, EngineError> {
history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
early_dispatch: bool,
) -> Result<StreamResponseOutput, EngineError> {
debug!(
item_count = request.items.len(),
tool_count = request.tools.len(),
@@ -2233,6 +2798,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
);
let mut stream = self.open_stream_with_retry(request, turn, llm_call).await?;
let mut early_tools: Option<EarlyToolExecutionBatch> = None;
let mut event_count: usize = 0;
loop {
@@ -2254,8 +2820,11 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
Err(err) => {
// 部分情報でも発火しておく（料金会計用）
self.timeline.flush_usage();
return Ok(StreamCompletion::Interrupted {
reason: err.to_string(),
return Ok(StreamResponseOutput {
completion: StreamCompletion::Interrupted {
reason: err.to_string(),
},
early_tools,
});
}
};
@@ -2269,9 +2838,116 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}
self.emit_stream_event(turn, llm_call, &event);
self.timeline.dispatch(&event);
if early_dispatch {
let collected = self.tool_call_collector.take_collected();
if !collected.is_empty() {
if early_tools.is_none() {
let batch_id = format!(
"tool-batch-{}",
self.tool_execution_batch_count
);
self.tool_execution_batch_count += 1;
early_tools =
Some(EarlyToolExecutionBatch::new(batch_id));
}
let fresh_calls = {
let batch = early_tools.as_ref().expect("batch initialized");
collected
.into_iter()
.filter(|call| {
!batch.seen_call_ids.contains(&call.id)
})
.collect::<Vec<_>>()
};
if !fresh_calls.is_empty() {
// Commit every completed assistant block observed so far,
// including the ToolCall itself, before provider ownership
// can start. This is the durable fail-closed gate.
let reasoning =
self.thinking_block_collector.take_collected();
let text = self.text_block_collector.take_collected();
let assistant_items = self.build_assistant_items(
&reasoning,
&text,
&fresh_calls,
);
if let Err(error) = self.append_history_items(
history,
assistant_items,
annotate,
) {
if let Some(mut batch) = early_tools.take() {
let _ = self
.cancel_and_terminalize_early_batch(
history,
annotate,
&mut batch,
)
.await;
}
return Err(error);
}
let admission = {
let batch =
early_tools.as_mut().expect("batch initialized");
self.admit_early_tool_calls(
history,
batch,
fresh_calls,
)
.await
};
match admission {
Ok(EarlyToolAdmission::Continue) => {}
Ok(EarlyToolAdmission::Pause) => {
self.timeline.abort_current_block();
self.timeline.flush_usage();
if let Some(batch) = early_tools.take() {
self.drain_early_tool_batch(
history,
annotate,
batch,
false,
true,
)
.await?;
}
return Err(EngineError::PauseRequested);
}
Err(error) => {
self.timeline.abort_current_block();
self.timeline.flush_usage();
if let Some(mut batch) = early_tools.take() {
let cleanup_error = self
.cancel_and_terminalize_early_batch(
history,
annotate,
&mut batch,
)
.await;
return Err(cleanup_error.unwrap_or(error));
}
return Err(error);
}
}
}
}
}
if let Event::Error(err) = &event {
self.timeline.abort_current_block();
self.timeline.flush_usage();
if let Some(mut batch) = early_tools.take() {
if let Some(error) = self
.cancel_and_terminalize_early_batch(
history,
annotate,
&mut batch,
)
.await
{
return Err(error);
}
}
return Err(EngineError::Client(ClientError::Api {
status: None,
code: err.code.clone(),
@@ -2291,6 +2967,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
// Preserve any UsageEvent already received for billing, but
// do not treat this paused request as a normal completion.
self.timeline.flush_usage();
if let Some(batch) = early_tools.take() {
self.drain_early_tool_batch(history, annotate, batch, false, true)
.await?;
}
return Err(EngineError::PauseRequested);
}
cancel = self.cancel_rx.recv() => {
@@ -2299,6 +2979,14 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}
self.timeline.abort_current_block();
self.timeline.flush_usage();
if let Some(mut batch) = early_tools.take() {
if let Some(error) = self
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
.await
{
return Err(error);
}
}
return Err(EngineError::Cancelled);
}
}
@@ -2306,7 +2994,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
// ストリーム完了時に集約済み Usage を 1 度だけ発火
self.timeline.flush_usage();
debug!(event_count = event_count, "Stream completed");
Ok(StreamCompletion::Complete)
Ok(StreamResponseOutput {
completion: StreamCompletion::Complete,
early_tools,
})
}
/// Execute tools and push results to history.
@@ -2361,6 +3052,7 @@ impl<C: LlmClient, A: Send + Sync> Engine<C, Mutable, A> {
llm_call_count: 0,
tool_execution_batch_count: 0,
max_turns: None,
tool_call_dispatch_mode: ToolCallDispatchMode::default(),
tool_execution_policy: ToolExecutionPolicy::default(),
turn_start_cbs: Vec::new(),
turn_end_cbs: Vec::new(),
@@ -2561,9 +3253,18 @@ impl<C: LlmClient, A: Send + Sync> Engine<C, Mutable, A> {
self.max_turns = max_turns;
}
/// Apply configuration (reserved for future extensions)
#[allow(dead_code)]
pub fn config(self, _config: EngineConfig) -> Self {
/// Select when provider-confirmed tool calls may begin execution.
pub fn set_tool_call_dispatch_mode(&mut self, mode: ToolCallDispatchMode) {
self.tool_call_dispatch_mode = mode;
}
pub fn tool_call_dispatch_mode(&self) -> ToolCallDispatchMode {
self.tool_call_dispatch_mode
}
/// Apply configuration.
pub fn config(mut self, config: EngineConfig) -> Self {
self.tool_call_dispatch_mode = config.tool_call_dispatch;
self
}
@@ -2640,6 +3341,7 @@ impl<C: LlmClient, A: Send + Sync> Engine<C, Mutable, A> {
llm_call_count: self.llm_call_count,
tool_execution_batch_count: self.tool_execution_batch_count,
max_turns: self.max_turns,
tool_call_dispatch_mode: self.tool_call_dispatch_mode,
tool_execution_policy: self.tool_execution_policy,
turn_start_cbs: self.turn_start_cbs,
turn_end_cbs: self.turn_end_cbs,
@@ -2823,6 +3525,7 @@ impl<C: LlmClient, A: Send + Sync> Engine<C, Locked, A> {
llm_call_count: self.llm_call_count,
tool_execution_batch_count: self.tool_execution_batch_count,
max_turns: self.max_turns,
tool_call_dispatch_mode: self.tool_call_dispatch_mode,
tool_execution_policy: self.tool_execution_policy,
turn_start_cbs: self.turn_start_cbs,
turn_end_cbs: self.turn_end_cbs,
crates/agen/src/lib.rs
+1 -1
View File
@@ -22,7 +22,7 @@ pub use agen_macros::{description, tool, tool_registry};
pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope};
pub use engine::{
Engine, EngineConfig, EngineError, EngineResult, EngineRunExit, EngineRunOutput,
LlmRetryNotice, RunInterruptionReason, ToolRegistryError,
LlmRetryNotice, RunInterruptionReason, ToolCallDispatchMode, ToolRegistryError,
};
pub use handler::ToolUseBlockStart;
pub use history::{History, HistoryEntry};
crates/agen/src/llm_client/client.rs
+19
View File
@@ -38,6 +38,16 @@ impl std::fmt::Display for ConfigWarning {
pub type ResponseStream = Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>;
/// Whether normalized `BlockStop(ToolUse)` events prove one call is complete.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum ToolCallCompletionSupport {
/// Individual tool-call completion events are provider-authoritative.
PerBlock,
/// Tool calls are complete only once the response stream itself completes.
#[default]
ResponseComplete,
}
/// LLMクライアントのtrait
///
/// 各プロバイダはこのtraitを実装し、統一されたインターフェースを提供する。
@@ -59,6 +69,11 @@ pub trait LlmClient: Send + Sync {
/// compaction) without access to the original construction parameters.
fn clone_boxed(&self) -> Box<dyn LlmClient>;
/// Describe the strongest safe tool-call completion boundary.
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
ToolCallCompletionSupport::ResponseComplete
}
/// 設定をバリデーションし、未サポートの設定があれば警告を返す
///
/// # Arguments
@@ -92,6 +107,10 @@ impl LlmClient for Box<dyn LlmClient> {
(**self).clone_boxed()
}
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
(**self).tool_call_completion_support()
}
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
(**self).validate_config(config)
}
crates/agen/src/llm_client/scheme/anthropic/scheme_impl.rs
+5
View File
@@ -9,6 +9,7 @@ use crate::llm_client::{
ClientError,
auth::AuthRequirement,
capability::ModelCapability,
client::ToolCallCompletionSupport,
event::{BlockType, Event, ReasoningBlockData},
scheme::Scheme,
types::Request,
@@ -104,4 +105,8 @@ impl Scheme for AnthropicScheme {
fn default_capability(&self) -> ModelCapability {
super::capability::default_capability()
}
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
ToolCallCompletionSupport::PerBlock
}
}
crates/agen/src/llm_client/scheme/mod.rs
+32 -1
View File
@@ -16,7 +16,7 @@ use serde_json::Value;
use super::auth::AuthRequirement;
use super::capability::ModelCapability;
use super::client::ConfigWarning;
use super::client::{ConfigWarning, ToolCallCompletionSupport};
use super::error::ClientError;
use super::event::Event;
use super::types::{Request, RequestConfig};
@@ -81,6 +81,11 @@ pub trait Scheme: Clone + Send + Sync + 'static {
/// 高レベルの client 構築層の責務で、scheme はここには関与しない。
fn default_capability(&self) -> ModelCapability;
/// Strongest safe completion boundary for streamed tool calls.
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
ToolCallCompletionSupport::ResponseComplete
}
/// scheme 側でサポートしていない `RequestConfig` フィールドを
/// 警告として返す（例: OpenAI Chat は `top_k` 非対応）。
/// デフォルトは空 Vec。
@@ -89,3 +94,29 @@ pub trait Scheme: Clone + Send + Sync + 'static {
Vec::new()
}
}
#[cfg(test)]
mod completion_support_tests {
use super::*;
use crate::llm_client::ToolCallCompletionSupport;
#[test]
fn only_schemes_with_authoritative_per_call_done_events_enable_early_dispatch() {
assert_eq!(
anthropic::AnthropicScheme::new().tool_call_completion_support(),
ToolCallCompletionSupport::PerBlock
);
assert_eq!(
openai_responses::OpenAIResponsesScheme::new().tool_call_completion_support(),
ToolCallCompletionSupport::PerBlock
);
assert_eq!(
openai_chat::OpenAIScheme::new().tool_call_completion_support(),
ToolCallCompletionSupport::ResponseComplete
);
assert_eq!(
gemini::GeminiScheme::new().tool_call_completion_support(),
ToolCallCompletionSupport::ResponseComplete
);
}
}
crates/agen/src/llm_client/scheme/openai_responses/scheme_impl.rs
+5 -1
View File
@@ -6,7 +6,7 @@ use crate::llm_client::{
ClientError,
auth::AuthRequirement,
capability::ModelCapability,
client::ConfigWarning,
client::{ConfigWarning, ToolCallCompletionSupport},
event::Event,
scheme::Scheme,
types::{Request, RequestConfig},
@@ -56,6 +56,10 @@ impl Scheme for OpenAIResponsesScheme {
super::capability::default_capability()
}
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
ToolCallCompletionSupport::PerBlock
}
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
let mut warnings = Vec::new();
// Some compatible backends reject `max_output_tokens` with HTTP 400.
crates/agen/src/llm_client/transport.rs
+4
View File
@@ -510,6 +510,10 @@ impl<S: Scheme + Clone + 'static> LlmClient for HttpTransport<S> {
self.scheme.validate_config(config)
}
fn tool_call_completion_support(&self) -> super::ToolCallCompletionSupport {
self.scheme.tool_call_completion_support()
}
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError> {
let total_started = Instant::now();
let path = self.scheme.path(&self.model_id);
crates/agen/src/timeline/timeline.rs
+37 -42
View File
@@ -256,8 +256,8 @@ where
H: Handler<ToolUseBlockKind>,
{
handler: H,
scope: Option<H::Scope>,
current_tool: Option<(String, String)>, // (id, name)
scopes: HashMap<usize, H::Scope>,
tools: HashMap<usize, (String, String)>, // index -> (id, name)
}
impl<H> ToolUseBlockHandlerWrapper<H>
@@ -267,8 +267,8 @@ where
fn new(handler: H) -> Self {
Self {
handler,
scope: None,
current_tool: None,
scopes: HashMap::new(),
tools: HashMap::new(),
}
}
}
@@ -279,60 +279,55 @@ where
H::Scope: Send + Sync,
{
fn dispatch_start(&mut self, start: &BlockStart) {
if let Some(scope) = &mut self.scope {
if let BlockMetadata::ToolUse { id, name } = &start.metadata {
self.current_tool = Some((id.clone(), name.clone()));
self.handler.on_event(
scope,
&ToolUseBlockEvent::Start(ToolUseBlockStart {
index: start.index,
id: id.clone(),
name: name.clone(),
}),
);
}
if let BlockMetadata::ToolUse { id, name } = &start.metadata {
let scope = self.scopes.entry(start.index).or_default();
self.tools.insert(start.index, (id.clone(), name.clone()));
self.handler.on_event(
scope,
&ToolUseBlockEvent::Start(ToolUseBlockStart {
index: start.index,
id: id.clone(),
name: name.clone(),
}),
);
}
}
fn dispatch_delta(&mut self, delta: &BlockDelta) {
if let Some(scope) = &mut self.scope {
if let DeltaContent::InputJson(json) = &delta.delta {
self.handler
.on_event(scope, &ToolUseBlockEvent::InputJsonDelta(json.clone()));
}
if let DeltaContent::InputJson(json) = &delta.delta {
let scope = self.scopes.entry(delta.index).or_default();
self.handler
.on_event(scope, &ToolUseBlockEvent::InputJsonDelta(json.clone()));
}
}
fn dispatch_stop(&mut self, stop: &BlockStop) {
if let Some(scope) = &mut self.scope {
if let Some((id, name)) = self.current_tool.take() {
self.handler.on_event(
scope,
&ToolUseBlockEvent::Stop(ToolUseBlockStop {
index: stop.index,
id,
name,
}),
);
}
if let (Some(mut scope), Some((id, name))) = (
self.scopes.remove(&stop.index),
self.tools.remove(&stop.index),
) {
self.handler.on_event(
&mut scope,
&ToolUseBlockEvent::Stop(ToolUseBlockStop {
index: stop.index,
id,
name,
}),
);
}
}
fn dispatch_abort(&mut self, _abort: &BlockAbort) {
self.current_tool = None;
fn dispatch_abort(&mut self, abort: &BlockAbort) {
self.scopes.remove(&abort.index);
self.tools.remove(&abort.index);
}
fn start_scope(&mut self) {
self.scope = Some(H::Scope::default());
}
fn start_scope(&mut self) {}
fn end_scope(&mut self) {
self.scope = None;
self.current_tool = None;
}
fn end_scope(&mut self) {}
fn has_scope(&self) -> bool {
self.scope.is_some()
!self.scopes.is_empty()
}
}
crates/agen/src/timeline/tool_call_collector.rs
+49 -2
View File
@@ -5,7 +5,6 @@
use crate::{
handler::{Handler, ToolUseBlockEvent, ToolUseBlockKind},
llm_client::types::parse_tool_arguments,
tool::ToolCall,
};
use std::sync::{Arc, Mutex};
@@ -85,7 +84,18 @@ impl Handler<ToolUseBlockKind> for ToolCallCollector {
// ブロック完了時にToolCallを確定
if let (Some(id), Some(name)) = (scope.current_id.take(), scope.current_name.take())
{
let input = parse_tool_arguments(&scope.input_json_buffer);
// A Stop event proves the provider finished the block, but
// malformed/non-object arguments are still not executable.
// Preserve an explicit null sentinel so the engine can commit
// a terminal validation error without guessing `{}`.
let input = if scope.input_json_buffer.trim().is_empty() {
serde_json::Value::Object(serde_json::Map::new())
} else {
serde_json::from_str::<serde_json::Value>(&scope.input_json_buffer)
.ok()
.filter(serde_json::Value::is_object)
.unwrap_or(serde_json::Value::Null)
};
let tool_call = ToolCall { id, name, input };
@@ -144,6 +154,43 @@ mod tests {
);
}
#[test]
fn test_collect_interleaved_tool_call_indexes_and_ignore_duplicate_stop() {
let collector = ToolCallCollector::new();
let mut timeline = Timeline::new();
timeline.on_tool_use_block(collector.clone());
timeline.dispatch(&Event::tool_use_start(0, "call_0", "tool_0"));
timeline.dispatch(&Event::tool_input_delta(0, r#"{"a":"#));
timeline.dispatch(&Event::tool_use_start(1, "call_1", "tool_1"));
timeline.dispatch(&Event::tool_input_delta(1, r#"{"b":2}"#));
timeline.dispatch(&Event::tool_input_delta(0, r#"1}"#));
timeline.dispatch(&Event::tool_use_stop(1));
timeline.dispatch(&Event::tool_use_stop(1));
timeline.dispatch(&Event::tool_use_stop(0));
let calls = collector.take_collected();
assert_eq!(calls.len(), 2);
assert_eq!(calls[0].id, "call_1");
assert_eq!(calls[0].input, serde_json::json!({"b": 2}));
assert_eq!(calls[1].id, "call_0");
assert_eq!(calls[1].input, serde_json::json!({"a": 1}));
}
#[test]
fn malformed_completed_arguments_are_not_normalized_to_executable_object() {
let collector = ToolCallCollector::new();
let mut timeline = Timeline::new();
timeline.on_tool_use_block(collector.clone());
timeline.dispatch(&Event::tool_use_start(0, "bad", "mutate"));
timeline.dispatch(&Event::tool_input_delta(0, r#"{"unterminated":"#));
timeline.dispatch(&Event::tool_use_stop(0));
let calls = collector.take_collected();
assert_eq!(calls.len(), 1);
assert!(calls[0].input.is_null());
}
#[test]
fn test_collect_multiple_tool_calls() {
let collector = ToolCallCollector::new();
crates/agen/tests/parallel_execution_test.rs
+471 -1
View File
@@ -11,12 +11,16 @@ use agen::interceptor::{
PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo,
};
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::llm_client::{
ClientError, LlmClient, Request, ResponseStream, ToolCallCompletionSupport,
};
use agen::tool::{
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
ToolResultDisposition,
};
use agen::{
Engine, EngineError, EngineRunExit, History, Item, RunInterruptionReason, ToolExecutionPolicy,
Engine, EngineError, EngineRunExit, History, Item, RunInterruptionReason, ToolCallDispatchMode,
ToolExecutionPolicy,
};
use async_trait::async_trait;
@@ -251,10 +255,476 @@ impl Tool for ContextRecordingTool {
}
}
#[derive(Clone)]
struct ControlledStreamClient {
receiver: Arc<Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<Result<Event, ClientError>>>>>,
completion_support: ToolCallCompletionSupport,
}
impl ControlledStreamClient {
fn new(
completion_support: ToolCallCompletionSupport,
) -> (
Self,
tokio::sync::mpsc::UnboundedSender<Result<Event, ClientError>>,
) {
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
(
Self {
receiver: Arc::new(Mutex::new(Some(rx))),
completion_support,
},
tx,
)
}
}
#[async_trait]
impl LlmClient for ControlledStreamClient {
async fn stream(&self, _request: Request) -> Result<ResponseStream, ClientError> {
let mut rx = self
.receiver
.lock()
.unwrap()
.take()
.ok_or_else(|| ClientError::Api {
status: Some(500),
code: Some("controlled_stream_exhausted".to_string()),
message: "controlled stream already consumed".to_string(),
retry_after: None,
})?;
Ok(Box::pin(futures::stream::poll_fn(move |cx| {
rx.poll_recv(cx)
})))
}
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
self.completion_support
}
}
#[derive(Clone)]
struct BarrierTool {
name: &'static str,
starts: Arc<AtomicUsize>,
started: Arc<tokio::sync::Notify>,
release: Arc<tokio::sync::Notify>,
}
impl BarrierTool {
fn new(name: &'static str) -> Self {
Self {
name,
starts: Arc::new(AtomicUsize::new(0)),
started: Arc::new(tokio::sync::Notify::new()),
release: Arc::new(tokio::sync::Notify::new()),
}
}
fn definition(&self) -> ToolDefinition {
let tool = self.clone();
Arc::new(move || {
(
ToolMeta::new(tool.name).input_schema(serde_json::json!({"type":"object"})),
Arc::new(tool.clone()) as Arc<dyn Tool>,
)
})
}
}
#[async_trait]
impl Tool for BarrierTool {
async fn execute(
&self,
_input_json: &str,
_ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
self.starts.fetch_add(1, Ordering::SeqCst);
self.started.notify_one();
self.release.notified().await;
Ok(format!("{} complete", self.name).into())
}
}
// =============================================================================
// Tests
// =============================================================================
#[tokio::test]
async fn default_dispatch_waits_for_response_completion() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let tool = BarrierTool::new("barrier_default");
let probe = tool.clone();
let mut engine = Engine::new(client);
engine.set_max_turns(Some(1));
engine.register_tool(tool.definition());
let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel();
engine.on_stream_event(move |_, _, event| {
if matches!(event, Event::Ping(_)) {
let _ = observed_tx.send(());
}
});
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
tx.send(Ok(Event::tool_use_start(
0,
"call_default",
"barrier_default",
)))
.unwrap();
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
tx.send(Ok(Event::ping())).unwrap();
observed_rx.recv().await.unwrap();
assert_eq!(probe.starts.load(Ordering::SeqCst), 0);
tx.send(Ok(Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})))
.unwrap();
drop(tx);
tokio::time::timeout(Duration::from_secs(1), probe.started.notified())
.await
.expect("tool starts after response completion");
probe.release.notify_one();
let _ = run.await.unwrap();
assert_eq!(probe.starts.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_result() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let first = BarrierTool::new("barrier_a");
let second = BarrierTool::new("barrier_b");
let first_probe = first.clone();
let second_probe = second.clone();
let mut engine = Engine::new(client);
engine.set_max_turns(Some(1));
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
engine.register_tool(first.definition());
engine.register_tool(second.definition());
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run both").await;
(output, history)
});
tx.send(Ok(Event::tool_use_start(0, "call_a", "barrier_a")))
.unwrap();
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
tokio::time::timeout(Duration::from_secs(1), first_probe.started.notified())
.await
.expect("first tool starts while stream is open");
tx.send(Ok(Event::tool_use_start(1, "call_b", "barrier_b")))
.unwrap();
tx.send(Ok(Event::tool_input_delta(1, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(1))).unwrap();
tokio::time::timeout(Duration::from_secs(1), second_probe.started.notified())
.await
.expect("second tool starts before first completes");
assert_eq!(first_probe.starts.load(Ordering::SeqCst), 1);
assert_eq!(second_probe.starts.load(Ordering::SeqCst), 1);
tx.send(Ok(Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})))
.unwrap();
drop(tx);
first_probe.release.notify_one();
second_probe.release.notify_one();
let (_output, history) = run.await.unwrap();
let calls = history
.iter()
.filter(|entry| matches!(entry.item, Item::ToolCall { .. }))
.count();
let results = history
.iter()
.filter(|entry| matches!(entry.item, Item::ToolResult { .. }))
.count();
assert_eq!((calls, results), (2, 2));
}
#[tokio::test]
async fn later_history_failure_terminalizes_already_started_early_sibling() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let first = BarrierTool::new("history_first");
let second = BarrierTool::new("history_second");
let first_probe = first.clone();
let second_probe = second.clone();
let mut engine = Engine::new(client);
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
engine.set_tool_execution_policy(ToolExecutionPolicy {
pause_safe_boundary_timeout: Duration::from_millis(1),
cancellation_request_timeout: Duration::from_millis(1),
terminal_confirmation_timeout: Duration::from_millis(1),
});
engine.register_tool(first.definition());
engine.register_tool(second.definition());
engine.on_history_append(|item| match item {
Item::ToolCall { call_id, .. } if call_id == "call_history_second" => {
Err("durable append rejected".to_string())
}
_ => Ok(()),
});
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
tx.send(Ok(Event::tool_use_start(
0,
"call_history_first",
"history_first",
)))
.unwrap();
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
tokio::time::timeout(Duration::from_secs(1), first_probe.started.notified())
.await
.expect("first side effect starts");
tx.send(Ok(Event::tool_use_start(
1,
"call_history_second",
"history_second",
)))
.unwrap();
tx.send(Ok(Event::tool_input_delta(1, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(1))).unwrap();
drop(tx);
let (output, history) = run.await.unwrap();
assert!(matches!(
output.result,
EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(
EngineError::HistoryAppend(ref message)
)) if message == "durable append rejected"
));
assert_eq!(first_probe.starts.load(Ordering::SeqCst), 1);
assert_eq!(second_probe.starts.load(Ordering::SeqCst), 0);
assert!(history.iter().any(|entry| matches!(
&entry.item,
Item::ToolResult {
call_id,
disposition: ToolResultDisposition::OutcomeUnknown,
..
} if call_id == "call_history_first"
)));
}
#[tokio::test]
async fn malformed_completed_arguments_commit_error_without_executing_tool() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let tool = SlowTool::new("must_not_run", 0);
let probe = tool.clone();
let mut engine = Engine::new(client);
engine.set_max_turns(Some(1));
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
engine.register_tool(tool.definition());
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
tx.send(Ok(Event::tool_use_start(
0,
"call_malformed",
"must_not_run",
)))
.unwrap();
tx.send(Ok(Event::tool_input_delta(0, r#"{"broken":"#)))
.unwrap();
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
tx.send(Ok(Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})))
.unwrap();
drop(tx);
let (_output, history) = run.await.unwrap();
assert_eq!(probe.call_count(), 0);
assert!(history.iter().any(|entry| matches!(
&entry.item,
Item::ToolResult {
call_id,
summary,
..
} if call_id == "call_malformed" && summary.contains("complete JSON object")
)));
}
#[tokio::test]
async fn pause_during_early_execution_drains_result_and_resume_does_not_rerun() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let tool = BarrierTool::new("barrier_pause");
let probe = tool.clone();
let mut engine = Engine::new(client);
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
engine.register_tool(tool.definition());
let pause = engine.pause_sender();
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
tx.send(Ok(Event::tool_use_start(0, "call_pause", "barrier_pause")))
.unwrap();
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
tokio::time::timeout(Duration::from_secs(1), probe.started.notified())
.await
.expect("tool starts before pause");
pause.send(()).await.unwrap();
probe.release.notify_one();
drop(tx);
let (output, mut history) = run.await.unwrap();
assert!(matches!(output.result, EngineRunExit::Paused));
assert!(history.iter().any(|entry| matches!(
&entry.item,
Item::ToolResult { call_id, .. } if call_id == "call_pause"
)));
let mut engine = output.engine;
let _ = engine.resume(&mut history).await;
assert_eq!(probe.starts.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn interrupted_stream_after_early_start_terminalizes_without_continuation_or_reexecution() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let tool = BarrierTool::new("barrier_interrupted");
let probe = tool.clone();
let mut engine = Engine::new(client);
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
engine.register_tool(tool.definition());
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
tx.send(Ok(Event::tool_use_start(
0,
"call_interrupted",
"barrier_interrupted",
)))
.unwrap();
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
tokio::time::timeout(Duration::from_secs(1), probe.started.notified())
.await
.expect("tool starts before stream interruption");
tx.send(Err(ClientError::Api {
status: None,
code: Some("connection_lost".to_string()),
message: "connection lost".to_string(),
retry_after: None,
}))
.unwrap();
probe.release.notify_one();
drop(tx);
let (output, mut history) = run.await.unwrap();
assert!(matches!(
output.result,
EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(EngineError::Client(
ClientError::Api {
code: Some(ref code),
..
}
))) if code == "early_tool_stream_interrupted"
));
assert_eq!(probe.starts.load(Ordering::SeqCst), 1);
assert_eq!(
history
.iter()
.filter(|entry| matches!(&entry.item, Item::ToolCall { call_id, .. } if call_id == "call_interrupted"))
.count(),
1
);
assert_eq!(
history
.iter()
.filter(|entry| matches!(&entry.item, Item::ToolResult { call_id, .. } if call_id == "call_interrupted"))
.count(),
1
);
let mut engine = output.engine;
let _ = engine.resume(&mut history).await;
assert_eq!(
probe.starts.load(Ordering::SeqCst),
1,
"terminalized early calls must not be resumed"
);
}
#[tokio::test]
async fn unsupported_provider_visibly_defers_opt_in_dispatch() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::ResponseComplete);
let tool = BarrierTool::new("barrier_fallback");
let probe = tool.clone();
let mut engine = Engine::new(client);
engine.set_max_turns(Some(1));
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
engine.register_tool(tool.definition());
let warnings = Arc::new(Mutex::new(Vec::new()));
let warning_probe = warnings.clone();
engine.on_warning(move |warning| warning_probe.lock().unwrap().push(warning.to_string()));
let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel();
engine.on_stream_event(move |_, _, event| {
if matches!(event, Event::Ping(_)) {
let _ = observed_tx.send(());
}
});
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
tx.send(Ok(Event::tool_use_start(
0,
"call_fallback",
"barrier_fallback",
)))
.unwrap();
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
tx.send(Ok(Event::ping())).unwrap();
observed_rx.recv().await.unwrap();
assert_eq!(probe.starts.load(Ordering::SeqCst), 0);
assert!(warnings.lock().unwrap().iter().any(|warning| {
warning.contains("provider does not expose a trustworthy per-call completion boundary")
}));
tx.send(Ok(Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})))
.unwrap();
drop(tx);
tokio::time::timeout(Duration::from_secs(1), probe.started.notified())
.await
.expect("fallback starts tool only after response completion");
probe.release.notify_one();
let _ = run.await.unwrap();
}
/// Verify that multiple tools are executed in parallel
///
/// If each tool takes 100ms, sequential execution would take 300ms+,
crates/manifest/src/config.rs
+37
View File
@@ -542,6 +542,8 @@ pub struct EngineManifestConfig {
#[serde(default)]
pub reasoning: Option<ReasoningControl>,
#[serde(default)]
pub tool_call_dispatch: Option<crate::ToolCallDispatchMode>,
#[serde(default)]
pub tool_output: ToolOutputLimitsPartial,
#[serde(default)]
pub file_upload: FileUploadLimitsPartial,
@@ -868,6 +870,7 @@ impl EngineManifestConfig {
top_k: upper.top_k.or(self.top_k),
stop_sequences: upper.stop_sequences.or(self.stop_sequences),
reasoning: upper.reasoning.or(self.reasoning),
tool_call_dispatch: upper.tool_call_dispatch.or(self.tool_call_dispatch),
tool_output: self.tool_output.merge(upper.tool_output),
file_upload: self.file_upload.merge(upper.file_upload),
}
@@ -1171,6 +1174,7 @@ impl TryFrom<WorkerManifestConfig> for WorkerManifest {
top_k: cfg.engine.top_k,
stop_sequences: cfg.engine.stop_sequences.unwrap_or_default(),
reasoning: cfg.engine.reasoning,
tool_call_dispatch: cfg.engine.tool_call_dispatch.unwrap_or_default(),
tool_output: ToolOutputLimits {
default_max_bytes: cfg
.engine
@@ -1343,6 +1347,39 @@ mod tests {
}
}
#[test]
fn tool_call_dispatch_defaults_merges_and_round_trips() {
let manifest: WorkerManifest = minimal_valid().try_into().unwrap();
assert_eq!(
manifest.engine.tool_call_dispatch,
crate::ToolCallDispatchMode::AfterResponse
);
let lower = WorkerManifestConfig {
engine: EngineManifestConfig {
tool_call_dispatch: Some(crate::ToolCallDispatchMode::OnToolCallComplete),
..Default::default()
},
..Default::default()
};
let upper = WorkerManifestConfig::default();
assert_eq!(
lower.merge(upper).engine.tool_call_dispatch,
Some(crate::ToolCallDispatchMode::OnToolCallComplete)
);
let mut configured = minimal_valid();
configured.engine.tool_call_dispatch =
Some(crate::ToolCallDispatchMode::OnToolCallComplete);
let manifest: WorkerManifest = configured.try_into().unwrap();
let snapshot = serde_json::to_value(&manifest).unwrap();
let restored: WorkerManifest = serde_json::from_value(snapshot).unwrap();
assert_eq!(
restored.engine.tool_call_dispatch,
crate::ToolCallDispatchMode::OnToolCallComplete
);
}
#[test]
fn resolve_minimal_succeeds() {
let manifest: WorkerManifest = minimal_valid().try_into().unwrap();
crates/manifest/src/lib.rs
+17
View File
@@ -587,6 +587,20 @@ pub struct WorkerMeta {
pub name: String,
}
/// When the engine may begin executing provider-confirmed tool calls.
///
/// `OnToolCallComplete` advances only the start time: calls may still overlap,
/// so completion order and externally allocated identifiers are not ordered.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ToolCallDispatchMode {
/// Preserve the historical behavior: wait for the whole response stream.
#[default]
AfterResponse,
/// Start each call once the provider emits a trustworthy per-call completion.
OnToolCallComplete,
}
/// Worker-level configuration embedded in the manifest.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EngineManifest {
@@ -614,6 +628,9 @@ pub struct EngineManifest {
pub stop_sequences: Vec<String>,
#[serde(default)]
pub reasoning: Option<ReasoningControl>,
/// Tool-call start timing. Omitted manifests preserve `after_response`.
#[serde(default)]
pub tool_call_dispatch: ToolCallDispatchMode,
/// Byte-size caps applied to tool `content` before it reaches the
/// conversation history. The section is optional in TOML — when
/// omitted, `ToolOutputLimits::default()` (64 KiB default cap, no
crates/worker/src/spawn/tool.rs
+7 -1
View File
@@ -965,6 +965,7 @@ fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfi
stop_sequences: (!manifest.engine.stop_sequences.is_empty())
.then_some(manifest.engine.stop_sequences.clone()),
reasoning: manifest.engine.reasoning.clone(),
tool_call_dispatch: Some(manifest.engine.tool_call_dispatch),
tool_output: ToolOutputLimitsPartial {
default_max_bytes: Some(manifest.engine.tool_output.default_max_bytes),
per_tool: manifest.engine.tool_output.per_tool.clone(),
@@ -2226,7 +2227,8 @@ max_tokens = 3333
Some("reviewer"),
&[("reviewer", "reviewer.toml", REVIEWER_PROFILE)],
);
let parent = parent_manifest(&parent_root, Some(&parent_root.join("deny")));
let mut parent = parent_manifest(&parent_root, Some(&parent_root.join("deny")));
parent.engine.tool_call_dispatch = manifest::ToolCallDispatchMode::OnToolCallComplete;
let scope = vec![abs_rule(&delegated, Permission::Read)];
let profile_config = child_config_from_profile(
@@ -2248,6 +2250,10 @@ max_tokens = 3333
SpawnProfileSelector::Inherit,
);
assert_eq!(
inherit_config.engine.tool_call_dispatch,
Some(manifest::ToolCallDispatchMode::OnToolCallComplete)
);
for config in [profile_config, inherit_config] {
assert_eq!(config.scope.allow, scope);
assert!(config.scope.deny.is_empty());
crates/worker/src/worker.rs
+14
View File
@@ -7138,6 +7138,12 @@ pub fn apply_worker_manifest<C: LlmClient + 'static, A: Send + Sync>(
) {
worker.set_request_config(request_config_from_engine_manifest(wm));
worker.set_max_turns(wm.max_turns.map(|n| n.get()));
worker.set_tool_call_dispatch_mode(match wm.tool_call_dispatch {
manifest::ToolCallDispatchMode::AfterResponse => agen::ToolCallDispatchMode::AfterResponse,
manifest::ToolCallDispatchMode::OnToolCallComplete => {
agen::ToolCallDispatchMode::OnToolCallComplete
}
});
// Worker owns the lifecycle strategy for already-started tool operations.
// The provider must first accept cooperative cancellation, then confirm a
// terminal result before this bounded deadline; Agen handles only the
@@ -8721,6 +8727,7 @@ mod build_summary_prompt_tests {
top_k: Some(40),
stop_sequences: vec!["\n\n".into(), "</stop>".into()],
reasoning: None,
tool_call_dispatch: manifest::ToolCallDispatchMode::OnToolCallComplete,
tool_output: manifest::ToolOutputLimits::default(),
file_upload: manifest::FileUploadLimits::default(),
};
@@ -8732,6 +8739,13 @@ mod build_summary_prompt_tests {
assert_eq!(config.top_p, Some(0.9));
assert_eq!(config.top_k, Some(40));
assert_eq!(config.stop_sequences, vec!["\n\n", "</stop>"]);
let mut engine = Engine::new(NoopClient);
apply_worker_manifest(&mut engine, &manifest);
assert_eq!(
engine.tool_call_dispatch_mode(),
agen::ToolCallDispatchMode::OnToolCallComplete
);
}
#[test]
docs/manifest.toml
+8
View File
@@ -121,6 +121,14 @@ ref = "anthropic/claude-sonnet-4-6"
# reasoning = "medium"
# reasoning = -1
# 任意。デフォルト: "after_response"。
# "on_tool_call_complete" は provider が個別 ToolCall の完成を保証する場合だけ、
# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider では
# 応答完了まで明示的に遅延する。これは直列実行モードではなく、先行 call の完了を
# 待たず後続 call も開始するため、完了順・DB 反映順・採番順は保証しない。
# 早期開始後に stream が切断された場合は自動 continuation せず、安全に terminalize して停止する。
# tool_call_dispatch = "on_tool_call_complete"
# 任意。tool 実行 content の byte 長キャップ。
# セクション省略時は default_max_bytes = 64 * 1024、per_tool 空。
# [worker.tool_output]