diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 1fa4fafd0..2ee42feaa 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -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 { } /// 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 + Send>>; +type ToolCallExecutionInfo = ( + ToolCall, + crate::tool::ToolMeta, + Arc, + 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, + seen_call_ids: HashSet, + call_info_map: HashMap, + started_calls: Vec<(String, String)>, + attempt_fence: ToolExecutionAttemptFence, + futures: FuturesUnordered, + execution_handles: HashMap, + synthetic_results: VecDeque, + terminal_call_ids: HashSet, + 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, +} + pub struct Engine { /// LLM client client: C, @@ -362,6 +428,8 @@ pub struct Engine { tool_execution_batch_count: usize, /// Maximum number of AgentTurns (None = unlimited) max_turns: Option, + /// 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 Engine { } = item { if !answered_call_ids.contains(call_id) { - let input = parse_tool_arguments(arguments); + let input = serde_json::from_str::(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 Engine { } } + async fn admit_early_tool_calls( + &mut self, + history: &History, + batch: &mut EarlyToolExecutionBatch, + tool_calls: Vec, + ) -> Result { + 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, + annotate: &mut impl FnMut(&Item) -> Result, + batch: &mut EarlyToolExecutionBatch, + ) -> Option { + 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::>().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, + annotate: &mut impl FnMut(&Item) -> Result, + mut batch: EarlyToolExecutionBatch, + force_cancel: bool, + initial_pause_requested: bool, + ) -> Result { + 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 Engine { 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 Engine { 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 Engine { 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 Engine { // `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 Engine { 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 Engine { request: Request, turn: usize, llm_call: usize, - ) -> Result { + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + early_dispatch: bool, + ) -> Result { debug!( item_count = request.items.len(), tool_count = request.tools.len(), @@ -2233,6 +2798,7 @@ impl Engine { ); let mut stream = self.open_stream_with_retry(request, turn, llm_call).await?; + let mut early_tools: Option = None; let mut event_count: usize = 0; loop { @@ -2254,8 +2820,11 @@ impl Engine { 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 Engine { } 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::>() + }; + 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 Engine { // 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 Engine { } 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 Engine { // ストリーム完了時に集約済み 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 Engine { 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 Engine { 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 Engine { 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 Engine { 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, diff --git a/crates/agen/src/lib.rs b/crates/agen/src/lib.rs index 89bdaee91..869d5c322 100644 --- a/crates/agen/src/lib.rs +++ b/crates/agen/src/lib.rs @@ -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}; diff --git a/crates/agen/src/llm_client/client.rs b/crates/agen/src/llm_client/client.rs index d21acac85..b112d4291 100644 --- a/crates/agen/src/llm_client/client.rs +++ b/crates/agen/src/llm_client/client.rs @@ -38,6 +38,16 @@ impl std::fmt::Display for ConfigWarning { pub type ResponseStream = Pin> + 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; + /// 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 { (**self).clone_boxed() } + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + (**self).tool_call_completion_support() + } + fn validate_config(&self, config: &RequestConfig) -> Vec { (**self).validate_config(config) } diff --git a/crates/agen/src/llm_client/scheme/anthropic/scheme_impl.rs b/crates/agen/src/llm_client/scheme/anthropic/scheme_impl.rs index ce15baa98..f4b7a0098 100644 --- a/crates/agen/src/llm_client/scheme/anthropic/scheme_impl.rs +++ b/crates/agen/src/llm_client/scheme/anthropic/scheme_impl.rs @@ -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 + } } diff --git a/crates/agen/src/llm_client/scheme/mod.rs b/crates/agen/src/llm_client/scheme/mod.rs index 8306b32db..0ea78b819 100644 --- a/crates/agen/src/llm_client/scheme/mod.rs +++ b/crates/agen/src/llm_client/scheme/mod.rs @@ -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 + ); + } +} diff --git a/crates/agen/src/llm_client/scheme/openai_responses/scheme_impl.rs b/crates/agen/src/llm_client/scheme/openai_responses/scheme_impl.rs index 8a6e4e0a2..0c3679c22 100644 --- a/crates/agen/src/llm_client/scheme/openai_responses/scheme_impl.rs +++ b/crates/agen/src/llm_client/scheme/openai_responses/scheme_impl.rs @@ -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 { let mut warnings = Vec::new(); // Some compatible backends reject `max_output_tokens` with HTTP 400. diff --git a/crates/agen/src/llm_client/transport.rs b/crates/agen/src/llm_client/transport.rs index 1198a9e2a..8eca250f1 100644 --- a/crates/agen/src/llm_client/transport.rs +++ b/crates/agen/src/llm_client/transport.rs @@ -510,6 +510,10 @@ impl LlmClient for HttpTransport { 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 { let total_started = Instant::now(); let path = self.scheme.path(&self.model_id); diff --git a/crates/agen/src/timeline/timeline.rs b/crates/agen/src/timeline/timeline.rs index f972ec31a..fcdac1fc3 100644 --- a/crates/agen/src/timeline/timeline.rs +++ b/crates/agen/src/timeline/timeline.rs @@ -256,8 +256,8 @@ where H: Handler, { handler: H, - scope: Option, - current_tool: Option<(String, String)>, // (id, name) + scopes: HashMap, + tools: HashMap, // index -> (id, name) } impl ToolUseBlockHandlerWrapper @@ -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() } } diff --git a/crates/agen/src/timeline/tool_call_collector.rs b/crates/agen/src/timeline/tool_call_collector.rs index fa04b0ede..3c2ed0ca2 100644 --- a/crates/agen/src/timeline/tool_call_collector.rs +++ b/crates/agen/src/timeline/tool_call_collector.rs @@ -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 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::(&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(); diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index 5df9f386e..f38f9f7cc 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -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>>>>, + completion_support: ToolCallCompletionSupport, +} + +impl ControlledStreamClient { + fn new( + completion_support: ToolCallCompletionSupport, + ) -> ( + Self, + tokio::sync::mpsc::UnboundedSender>, + ) { + 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 { + 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 { + Box::new(self.clone()) + } + + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + self.completion_support + } +} + +#[derive(Clone)] +struct BarrierTool { + name: &'static str, + starts: Arc, + started: Arc, + release: Arc, +} + +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, + ) + }) + } +} + +#[async_trait] +impl Tool for BarrierTool { + async fn execute( + &self, + _input_json: &str, + _ctx: ToolExecutionContext, + ) -> Result { + 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+, diff --git a/crates/manifest/src/config.rs b/crates/manifest/src/config.rs index 9dc0e982e..d291883ef 100644 --- a/crates/manifest/src/config.rs +++ b/crates/manifest/src/config.rs @@ -542,6 +542,8 @@ pub struct EngineManifestConfig { #[serde(default)] pub reasoning: Option, #[serde(default)] + pub tool_call_dispatch: Option, + #[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 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(); diff --git a/crates/manifest/src/lib.rs b/crates/manifest/src/lib.rs index e91c5b913..4cb678470 100644 --- a/crates/manifest/src/lib.rs +++ b/crates/manifest/src/lib.rs @@ -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, #[serde(default)] pub reasoning: Option, + /// 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 diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index f7bd11149..74a17ac47 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -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()); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index ef8bd428c..6ab90b0d0 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -7138,6 +7138,12 @@ pub fn apply_worker_manifest( ) { 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(), "".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", ""]); + + let mut engine = Engine::new(NoopClient); + apply_worker_manifest(&mut engine, &manifest); + assert_eq!( + engine.tool_call_dispatch_mode(), + agen::ToolCallDispatchMode::OnToolCallComplete + ); } #[test] diff --git a/docs/manifest.toml b/docs/manifest.toml index 92b671e6e..274a792e8 100644 --- a/docs/manifest.toml +++ b/docs/manifest.toml @@ -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]