From dcf14ef565161d54a07a8d1b1cf35a373e105d44 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 12:00:36 +0900 Subject: [PATCH] fix: harden early tool call dispatch --- crates/agen/src/engine.rs | 537 ++++++++++++++---- crates/agen/src/interceptor.rs | 14 +- .../scheme/openai_responses/events.rs | 528 +++++++++++++++-- crates/agen/src/llm_client/types.rs | 30 +- crates/agen/src/timeline/mod.rs | 1 + .../agen/src/timeline/tool_call_collector.rs | 48 +- crates/agen/tests/common/mod.rs | 13 +- crates/agen/tests/parallel_execution_test.rs | 422 +++++++++++++- crates/session-store/src/logged_item.rs | 36 +- crates/session-store/src/public_snapshot.rs | 1 + crates/worker/src/ipc/interceptor.rs | 4 + crates/worker/src/spawn/registry.rs | 3 + docs/manifest.toml | 5 +- 13 files changed, 1481 insertions(+), 161 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 2ee42feaa..fa4ab25a6 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -24,7 +24,7 @@ use crate::{ ToolCallInfo, ToolResultInfo, TurnEndAction, }, llm_client::{ - ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream, + ClientError, ConfigWarning, ItemStatus, LlmClient, Request, RequestConfig, ResponseStream, ToolDefinition, error::is_retryable, event::Event, retry::RetryPolicy, transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT, }, @@ -110,33 +110,56 @@ fn materialize_provider_history(items: &[Item]) -> ProviderHistoryProjection { let mut materialized: Vec<_> = items.iter().cloned().enumerate().collect(); let mut call_order = HashMap::::new(); let mut next_call_order = 0usize; - let mut index = 0usize; - while index < materialized.len() { - match &materialized[index].1 { - Item::ToolCall { call_id, .. } => { - call_order.insert(call_id.clone(), next_call_order); - next_call_order += 1; - index += 1; - } - Item::ToolResult { .. } => { - let start = index; - while index < materialized.len() - && matches!(materialized[index].1, Item::ToolResult { .. }) - { - index += 1; - } - materialized[start..index].sort_by_key(|(_, item)| match item { - Item::ToolResult { call_id, .. } => { - call_order.get(call_id).copied().unwrap_or(usize::MAX) - } - _ => unreachable!("tool-result run contains only ToolResult items"), - }); - } - _ => index += 1, + for (_, item) in &materialized { + if let Item::ToolCall { + call_id, + call_index, + .. + } = item + { + call_order.insert(call_id.clone(), call_index.unwrap_or(next_call_order)); + next_call_order += 1; } } + let mut index = 0usize; + while index < materialized.len() { + if !matches!( + materialized[index].1, + Item::ToolCall { .. } | Item::ToolResult { .. } + ) { + index += 1; + continue; + } + let start = index; + while index < materialized.len() + && matches!( + materialized[index].1, + Item::ToolCall { .. } | Item::ToolResult { .. } + ) + { + index += 1; + } + materialized[start..index].sort_by_key(|(_, item)| match item { + Item::ToolCall { + call_id, + call_index, + .. + } => ( + 0usize, + call_index + .or_else(|| call_order.get(call_id).copied()) + .unwrap_or(usize::MAX), + ), + Item::ToolResult { call_id, .. } => ( + 1usize, + call_order.get(call_id).copied().unwrap_or(usize::MAX), + ), + _ => unreachable!("tool run contains only calls and results"), + }); + } + let mut original_to_projected_index = vec![0; materialized.len()]; for (projected_index, (original_index, _)) in materialized.iter().enumerate() { original_to_projected_index[*original_index] = projected_index; @@ -230,8 +253,7 @@ type ToolCallExecutionInfo = ( /// provider-confirmed blocks complete. struct EarlyToolExecutionBatch { batch_id: String, - next_call_index: usize, - calls: Vec, + calls: Vec<(usize, ToolCall)>, seen_call_ids: HashSet, call_info_map: HashMap, started_calls: Vec<(String, String)>, @@ -247,7 +269,6 @@ 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(), @@ -378,6 +399,37 @@ struct StreamResponseOutput { early_tools: Option, } +const RESPONSE_STREAM_PUMP_CAPACITY: usize = 64; + +struct ResponseStreamPump { + receiver: mpsc::Receiver>, + task: tokio::task::JoinHandle<()>, +} + +impl ResponseStreamPump { + fn start(mut stream: ResponseStream) -> Self { + let (sender, receiver) = mpsc::channel(RESPONSE_STREAM_PUMP_CAPACITY); + let task = tokio::spawn(async move { + while let Some(event) = stream.next().await { + if sender.send(event).await.is_err() { + break; + } + } + }); + Self { receiver, task } + } + + async fn recv(&mut self) -> Option> { + self.receiver.recv().await + } +} + +impl Drop for ResponseStreamPump { + fn drop(&mut self) { + self.task.abort(); + } +} + pub struct Engine { /// LLM client client: C, @@ -1229,67 +1281,88 @@ impl Engine { exit } - /// Check for pending tool calls (for resuming from Pause) - fn get_pending_tool_calls(&self, history: &History) -> Option> { - // Find the last ToolCall items that don't have corresponding ToolResult + /// Classify unanswered tool calls for safe resume. Calls with a persisted + /// execution identity were already handed to a provider and are never rerun. + fn get_unanswered_tool_calls(&self, history: &History) -> (Vec, Vec) { let mut pending_calls = Vec::new(); + let mut started_call_ids = Vec::new(); let mut answered_call_ids = std::collections::HashSet::new(); - // First pass: collect all answered call IDs for item in history.items() { if let Item::ToolResult { call_id, .. } = item { answered_call_ids.insert(call_id.clone()); } } - // Second pass: find unanswered tool calls for item in history.items() { if let Item::ToolCall { call_id, name, arguments, + execution_id, + status, .. } = item + && !answered_call_ids.contains(call_id) { - if !answered_call_ids.contains(call_id) { - 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(), - input, - }); + if execution_id.is_some() || *status == Some(ItemStatus::InProgress) { + started_call_ids.push(call_id.clone()); + continue; } + 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(), + input, + }); } } - if pending_calls.is_empty() { - None + (pending_calls, started_call_ids) + } + + fn append_early_tool_call( + &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + call: &ToolCall, + call_index: usize, + execution_id: Option, + ) -> Result<(), EngineError> { + let status = if execution_id.is_some() { + ItemStatus::InProgress } else { - Some(pending_calls) - } + ItemStatus::Completed + }; + let item = Item::tool_call_json(&call.id, &call.name, call.input.clone()) + .with_tool_execution_metadata(call_index, execution_id) + .with_status(status); + self.append_history_items(history, std::iter::once(item), annotate) } async fn admit_early_tool_calls( &mut self, - history: &History, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, batch: &mut EarlyToolExecutionBatch, - tool_calls: Vec, + tool_calls: Vec, ) -> Result { - for mut tool_call in tool_calls { + for collected in tool_calls { + let call_index = collected.call_index; + let mut tool_call = collected.call; 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()); + batch.calls.push((call_index, 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() { + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; batch.synthetic_results.push_back(ToolResult::error( &tool_call.id, "Tool arguments must be one complete JSON object", @@ -1299,21 +1372,11 @@ impl Engine { } 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}")), - ) - })); + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + batch.synthetic_results.push_back(ToolResult::error( + &tool_call.id, + format!("Tool not found: {}", tool_call.name), + )); batch.dispatched_any = true; continue; }; @@ -1332,16 +1395,28 @@ impl Engine { tool, context, }; - let action = self - .interceptor - .pre_tool_call(&mut info) - .await - .map_err(|error| { - EngineError::from(InterceptorFailure::new( - InterceptorPhase::PreToolCall, - error, - )) - })?; + let action = tokio::select! { + result = self.interceptor.pre_tool_call(&mut info) => { + result.map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPhase::PreToolCall, + error, + )) + })? + } + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!(call_id = %expected_tool_use_id, "Paused while awaiting pre-tool policy"); + } + return Ok(EarlyToolAdmission::Pause); + } + cancel = self.cancel_rx.recv() => { + if cancel.is_some() { + info!(call_id = %expected_tool_use_id, "Cancelled while awaiting pre-tool policy"); + } + return Err(EngineError::Cancelled); + } + }; if info.call.id != expected_tool_use_id { return Err(InterceptorFailure::new( InterceptorPhase::PreToolCall, @@ -1352,22 +1427,30 @@ impl Engine { ) .into()); } - tool_call = info.call; - let mut context = info.context; + let ToolCallInfo { + call, + meta, + tool, + mut context, + .. + } = info; + tool_call = call; 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(), - ), + (tool_call.clone(), meta, tool.clone(), context.clone()), ); match action { PreToolAction::Continue => { - let attempt_id = context.batch_id.clone(); + let attempt_id = context.execution_id(); + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + Some(attempt_id.clone()), + )?; batch .attempt_fence .register(tool_call.id.clone(), attempt_id.clone()); @@ -1376,8 +1459,7 @@ impl Engine { .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); + let (handle, terminal) = ToolExecutionHandle::start(tool, input_json, context); batch.execution_handles.insert(call_id.clone(), handle); batch.futures.push(Box::pin(async move { let result = match terminal.await { @@ -1420,17 +1502,57 @@ impl Engine { ) .into()); } + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; 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), + PreToolAction::Skip => { + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + batch.synthetic_results.push_back(ToolResult::error( + &tool_call.id, + "Tool execution skipped by interceptor", + )); + batch.dispatched_any = true; + } + PreToolAction::Abort(reason) => { + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + batch.synthetic_results.push_back(ToolResult::error( + &tool_call.id, + format!("Tool execution aborted: {reason}"), + )); + batch.dispatched_any = true; + return Err(EngineError::Aborted(reason)); + } + PreToolAction::Pause => { + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + return Ok(EarlyToolAdmission::Pause); + } } } Ok(EarlyToolAdmission::Continue) } + async fn commit_early_synthetic_results( + &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + batch: &mut EarlyToolExecutionBatch, + ) -> Result<(), EngineError> { + while let Some(result) = batch.synthetic_results.pop_front() { + self.finalize_and_commit_tool_result( + history, + annotate, + result, + None, + &batch.call_info_map, + &mut batch.attempt_fence, + &mut batch.terminal_call_ids, + ) + .await?; + } + Ok(()) + } + async fn cancel_and_terminalize_early_batch( &mut self, history: &mut History, @@ -1439,6 +1561,25 @@ impl Engine { ) -> Option { use futures::StreamExt as _; + let mut batch_error = None; + 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 + && batch_error.is_none() + { + batch_error = Some(error); + } + } + let cancellation_request_deadline = TokioInstant::now() + self.tool_execution_policy.cancellation_request_timeout; let cancellation_requests = batch @@ -1462,7 +1603,6 @@ impl Engine { } } - let mut batch_error = None; let deadline = TokioInstant::now() + self.tool_execution_policy.terminal_confirmation_timeout; while !batch.futures.is_empty() { @@ -2131,11 +2271,33 @@ impl Engine { "Starting engine run" ); - // Resume pending tool calls from a previous Pause - if let Some(tool_calls) = self.get_pending_tool_calls(history) { - info!("Resuming pending tool calls"); + // Resume only calls durably proven not started. A persisted started + // attempt without a terminal result is closed as OutcomeUnknown first. + let (pending_tool_calls, started_call_ids) = self.get_unanswered_tool_calls(history); + if !started_call_ids.is_empty() { + let call_info_map = HashMap::new(); + let mut attempt_fence = ToolExecutionAttemptFence::default(); + let mut terminal_call_ids = HashSet::new(); + for call_id in started_call_ids { + self.finalize_and_commit_tool_result( + history, + annotate, + ToolResult::outcome_unknown(&call_id), + None, + &call_info_map, + &mut attempt_fence, + &mut terminal_call_ids, + ) + .await?; + } + } + if !pending_tool_calls.is_empty() { + info!( + count = pending_tool_calls.len(), + "Resuming pending tool calls" + ); if let Some(result) = self - .execute_and_commit_tools(history, annotate, tool_calls) + .execute_and_commit_tools(history, annotate, pending_tool_calls) .await? { return Ok(result); @@ -2348,17 +2510,28 @@ impl Engine { self.request_trace_payload(&request), ); let request = self.attach_transport_trace(request, current_turn, current_llm_call); + let provider_supports_early = self.client.tool_call_completion_support() + == crate::llm_client::ToolCallCompletionSupport::PerBlock; + let interceptor_supports_early = self.interceptor.supports_early_tool_dispatch(); let early_dispatch = self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete - && self.client.tool_call_completion_support() - == crate::llm_client::ToolCallCompletionSupport::PerBlock; + && provider_supports_early + && interceptor_supports_early; if self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete - && !early_dispatch + && !provider_supports_early { self.emit_warning( "tool-call early dispatch deferred until response completion: provider does not expose a trustworthy per-call completion boundary", ); } + if self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete + && provider_supports_early + && !interceptor_supports_early + { + self.emit_warning( + "tool-call early dispatch deferred until response completion: interceptor requires the whole-response turn-end gate", + ); + } let assistant_start = history.len(); let StreamResponseOutput { completion: stream_outcome, @@ -2466,9 +2639,25 @@ 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 remaining_tool_calls = self.tool_call_collector.take_collected(); - let assistant_items = - self.build_assistant_items(&reasoning_items, &text_blocks, &remaining_tool_calls); + let remaining_collected_calls = self.tool_call_collector.take_collected_with_index(); + let mut remaining_tool_calls = remaining_collected_calls + .iter() + .map(|collected| (collected.call_index, collected.call.clone())) + .collect::>(); + remaining_tool_calls.sort_by_key(|(call_index, _)| *call_index); + let remaining_tool_calls = remaining_tool_calls + .into_iter() + .map(|(_, call)| call) + .collect::>(); + let assistant_items = self.build_assistant_items( + &reasoning_items, + &text_blocks, + if early_dispatch { + &[] + } else { + &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 @@ -2489,8 +2678,9 @@ impl Engine { let admission = self .admit_early_tool_calls( history, + annotate, early_tools.as_mut().expect("batch initialized"), - remaining_tool_calls.clone(), + remaining_collected_calls, ) .await; let admission = match admission { @@ -2513,10 +2703,27 @@ impl Engine { } return Ok(EngineResult::Paused); } + if let Some(batch) = early_tools.as_mut() + && let Err(error) = self + .commit_early_synthetic_results(history, annotate, batch) + .await + { + 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); + } } let mut tool_calls = early_tools .as_ref() - .map(|batch| batch.calls.clone()) + .map(|batch| { + let mut calls = batch.calls.clone(); + calls.sort_by_key(|(call_index, _)| *call_index); + calls.into_iter().map(|(_, call)| call).collect() + }) .unwrap_or_default(); if !early_dispatch { tool_calls = remaining_tool_calls; @@ -2797,13 +3004,15 @@ impl Engine { "Sending request to LLM" ); - let mut stream = self.open_stream_with_retry(request, turn, llm_call).await?; + let stream = self.open_stream_with_retry(request, turn, llm_call).await?; + self.tool_call_collector.begin_response(); + let mut stream = ResponseStreamPump::start(stream); let mut early_tools: Option = None; let mut event_count: usize = 0; loop { tokio::select! { - event_result = stream.next() => { + event_result = stream.recv() => { match event_result { Some(result) => { match &result { @@ -2839,7 +3048,7 @@ 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(); + let collected = self.tool_call_collector.take_collected_with_index(); if !collected.is_empty() { if early_tools.is_none() { let batch_id = format!( @@ -2854,22 +3063,23 @@ impl Engine { let batch = early_tools.as_ref().expect("batch initialized"); collected .into_iter() - .filter(|call| { - !batch.seen_call_ids.contains(&call.id) + .filter(|collected| { + !batch.seen_call_ids.contains(&collected.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. + // Commit every completed reasoning/text block observed so + // far. Admission durably commits each ToolCall with its + // started/not-started state immediately before ownership + // can transfer to a provider. 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, @@ -2892,13 +3102,38 @@ impl Engine { early_tools.as_mut().expect("batch initialized"); self.admit_early_tool_calls( history, + annotate, batch, fresh_calls, ) .await }; match admission { - Ok(EarlyToolAdmission::Continue) => {} + Ok(EarlyToolAdmission::Continue) => { + let batch = early_tools + .as_mut() + .expect("batch initialized"); + if let Err(error) = self + .commit_early_synthetic_results( + history, + annotate, + batch, + ) + .await + { + 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); + } + } Ok(EarlyToolAdmission::Pause) => { self.timeline.abort_current_block(); self.timeline.flush_usage(); @@ -2959,6 +3194,45 @@ impl Engine { None => break, } } + terminal = async { + early_tools + .as_mut() + .expect("terminal branch requires early batch") + .futures + .next() + .await + }, if early_tools.as_ref().is_some_and(|batch| !batch.futures.is_empty()) => { + let (attempt_id, result) = + terminal.expect("non-empty early terminal queue returns a result"); + let commit = { + let batch = early_tools.as_mut().expect("batch initialized"); + 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 + }; + if let Err(error) = commit { + 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); + } + } pause = self.pause_rx.recv() => { if pause.is_some() { info!("Paused during response stream"); @@ -3716,6 +3990,41 @@ mod tests { assert_eq!(projection.original_to_projected_index, [0, 1, 3, 2]); } + #[test] + fn provider_projection_groups_interleaved_calls_and_results_in_model_order() { + let items = vec![ + Item::tool_call_json("call_second", "second", serde_json::json!({})) + .with_tool_execution_metadata(1, Some("batch:call_second".to_string())) + .with_status(ItemStatus::InProgress), + Item::tool_result("call_second", "second result"), + Item::tool_call_json("call_first", "first", serde_json::json!({})) + .with_tool_execution_metadata(0, Some("batch:call_first".to_string())) + .with_status(ItemStatus::InProgress), + Item::tool_result("call_first", "first result"), + ]; + + let projection = materialize_provider_history(&items); + let order = projection + .items + .iter() + .map(|item| match item { + Item::ToolCall { call_id, .. } => format!("call:{call_id}"), + Item::ToolResult { call_id, .. } => format!("result:{call_id}"), + _ => unreachable!(), + }) + .collect::>(); + assert_eq!( + order, + [ + "call:call_first", + "call:call_second", + "result:call_first", + "result:call_second", + ] + ); + assert_eq!(projection.original_to_projected_index, [1, 3, 0, 2]); + } + #[test] fn tool_attachment_round_trips_through_durable_history_json() { let body: Arc<[u8]> = Arc::from(&b"image-body"[..]); diff --git a/crates/agen/src/interceptor.rs b/crates/agen/src/interceptor.rs index 28b2afa6b..67d34166d 100644 --- a/crates/agen/src/interceptor.rs +++ b/crates/agen/src/interceptor.rs @@ -336,6 +336,14 @@ pub struct ToolResultInfo<'a, A = ()> { /// flows, permission checks, and other trusted host adaptation. #[async_trait] pub trait Interceptor: Send + Sync { + /// Whether side effects may start before [`Self::on_assistant_turn_end`]. + /// + /// The conservative default preserves the historical whole-response gate. + /// Interceptors that do not use turn-end policy may opt in explicitly. + fn supports_early_tool_dispatch(&self) -> bool { + false + } + /// Called after receiving user input, before adding it to Engine history. async fn on_prompt_submit( &self, @@ -420,4 +428,8 @@ pub trait Interceptor: Send + Sync { pub(crate) struct DefaultInterceptor; #[async_trait] -impl Interceptor for DefaultInterceptor {} +impl Interceptor for DefaultInterceptor { + fn supports_early_tool_dispatch(&self) -> bool { + true + } +} diff --git a/crates/agen/src/llm_client/scheme/openai_responses/events.rs b/crates/agen/src/llm_client/scheme/openai_responses/events.rs index 47c8233b4..cb6ead57e 100644 --- a/crates/agen/src/llm_client/scheme/openai_responses/events.rs +++ b/crates/agen/src/llm_client/scheme/openai_responses/events.rs @@ -28,6 +28,43 @@ pub struct OpenAIResponsesState { /// 蓄積、`output_item.done`(Reasoning) で既存 reasoning_text block または /// metadata-only Thinking block に reasoning persistence material を載せる。 pending_reasoning: HashMap, + /// Tool input deltas and authoritative final values, keyed by output index. + pending_tool_inputs: HashMap, +} + +/// A tool input stream accumulated until an authoritative done payload arrives. +#[derive(Debug)] +struct PendingToolInput { + kind: ToolInputKind, + streamed: String, + saw_delta: bool, + final_value: Option, +} + +impl PendingToolInput { + fn new(kind: ToolInputKind) -> Self { + Self { + kind, + streamed: String::new(), + saw_delta: false, + final_value: None, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ToolInputKind { + FunctionCall, + CustomToolCall, +} + +impl ToolInputKind { + fn label(self) -> &'static str { + match self { + Self::FunctionCall => "function-call arguments", + Self::CustomToolCall => "custom-tool input", + } + } } /// 1 つの reasoning output_item の蓄積バッファ。 @@ -173,6 +210,8 @@ enum OutputItem { id: Option, call_id: String, name: String, + #[serde(default)] + arguments: Option, }, CustomToolCall { #[allow(dead_code)] @@ -180,6 +219,8 @@ enum OutputItem { id: Option, call_id: String, name: String, + #[serde(default)] + input: Option, }, #[serde(other)] Other, @@ -260,12 +301,24 @@ struct FunctionCallArgumentsDelta { delta: String, } +#[derive(Debug, Deserialize)] +struct FunctionCallArgumentsDone { + output_index: usize, + arguments: String, +} + #[derive(Debug, Deserialize)] struct CustomToolCallInputDelta { output_index: usize, delta: String, } +#[derive(Debug, Deserialize)] +struct CustomToolCallInputDone { + output_index: usize, + input: String, +} + #[derive(Debug, Deserialize)] struct ResponseCompleted { response: CompletedResponse, @@ -398,16 +451,20 @@ pub(crate) fn parse_sse( "response.output_item.added" => { let ev: OutputItemAdded = from_json(data)?; match ev.item { - OutputItem::FunctionCall { call_id, name, .. } - | OutputItem::CustomToolCall { call_id, name, .. } => { - let info = - state.allocate(SlotKey::OutputItem(ev.output_index), BlockType::ToolUse); - Ok(vec![Event::BlockStart(BlockStart { - index: info.flat_index, - block_type: BlockType::ToolUse, - metadata: BlockMetadata::ToolUse { id: call_id, name }, - })]) - } + OutputItem::FunctionCall { call_id, name, .. } => start_tool_output_item( + state, + ev.output_index, + ToolInputKind::FunctionCall, + call_id, + name, + ), + OutputItem::CustomToolCall { call_id, name, .. } => start_tool_output_item( + state, + ev.output_index, + ToolInputKind::CustomToolCall, + call_id, + name, + ), OutputItem::Reasoning { id, .. } => { // wrapper を確保。中身の content_part / summary_part は // 別 SlotKey で扱われ続ける（Streaming 表示は維持）。 @@ -431,19 +488,19 @@ pub(crate) fn parse_sse( id, encrypted_content, .. - } = ev.item + } = &ev.item { let mut pending = state .pending_reasoning .remove(&ev.output_index) .unwrap_or_default(); if pending.id.is_none() { - pending.id = id; + pending.id = id.clone(); } let mut stop_blocks = std::mem::take(&mut pending.deferred_thinking_stops); stop_blocks.extend(state.take_active_reasoning_slots(ev.output_index)); - let reasoning = pending.into_reasoning_data(encrypted_content); + let reasoning = pending.into_reasoning_data(encrypted_content.clone()); if stop_blocks.is_empty() { let info = @@ -478,15 +535,45 @@ pub(crate) fn parse_sse( }) .collect()); } - if let Some(info) = state.slots.remove(&SlotKey::OutputItem(ev.output_index)) { - Ok(vec![Event::BlockStop(BlockStop { - index: info.flat_index, - block_type: info.block_type, - stop_reason: None, - reasoning: None, - })]) - } else { - Ok(Vec::new()) + match ev.item { + OutputItem::FunctionCall { + call_id, + name, + arguments, + .. + } => finish_tool_output_item( + state, + ev.output_index, + ToolInputKind::FunctionCall, + call_id, + name, + arguments, + ), + OutputItem::CustomToolCall { + call_id, + name, + input, + .. + } => finish_tool_output_item( + state, + ev.output_index, + ToolInputKind::CustomToolCall, + call_id, + name, + input, + ), + _ => { + if let Some(info) = state.slots.remove(&SlotKey::OutputItem(ev.output_index)) { + Ok(vec![Event::BlockStop(BlockStop { + index: info.flat_index, + block_type: info.block_type, + stop_reason: None, + reasoning: None, + })]) + } else { + Ok(Vec::new()) + } + } } } @@ -619,30 +706,50 @@ pub(crate) fn parse_sse( "response.function_call_arguments.delta" => { let ev: FunctionCallArgumentsDelta = from_json(data)?; - Ok(ensure_and_delta( + record_tool_input_delta( state, - SlotKey::OutputItem(ev.output_index), - BlockType::ToolUse, + ev.output_index, + ToolInputKind::FunctionCall, + ev.delta, + ) + } + + "response.function_call_arguments.done" => { + let ev: FunctionCallArgumentsDone = from_json(data)?; + reconcile_tool_input( + state, + ev.output_index, + ToolInputKind::FunctionCall, + ev.arguments, BlockMetadata::ToolUse { id: String::new(), name: String::new(), }, - DeltaContent::InputJson(ev.delta), - )) + ) } "response.custom_tool_call_input.delta" => { let ev: CustomToolCallInputDelta = from_json(data)?; - Ok(ensure_and_delta( + record_tool_input_delta( state, - SlotKey::OutputItem(ev.output_index), - BlockType::ToolUse, + ev.output_index, + ToolInputKind::CustomToolCall, + ev.delta, + ) + } + + "response.custom_tool_call_input.done" => { + let ev: CustomToolCallInputDone = from_json(data)?; + reconcile_tool_input( + state, + ev.output_index, + ToolInputKind::CustomToolCall, + ev.input, BlockMetadata::ToolUse { id: String::new(), name: String::new(), }, - DeltaContent::InputJson(ev.delta), - )) + ) } "error" => { @@ -666,6 +773,181 @@ pub(crate) fn parse_sse( } } +fn start_tool_output_item( + state: &mut OpenAIResponsesState, + output_index: usize, + kind: ToolInputKind, + call_id: String, + name: String, +) -> Result, ClientError> { + ensure_tool_input_kind(state, output_index, kind)?; + let (info, just_created) = + state.get_or_allocate(SlotKey::OutputItem(output_index), BlockType::ToolUse); + if just_created { + Ok(vec![Event::BlockStart(BlockStart { + index: info.flat_index, + block_type: BlockType::ToolUse, + metadata: BlockMetadata::ToolUse { id: call_id, name }, + })]) + } else { + Ok(Vec::new()) + } +} + +fn record_tool_input_delta( + state: &mut OpenAIResponsesState, + output_index: usize, + kind: ToolInputKind, + delta: String, +) -> Result, ClientError> { + let pending = ensure_tool_input_kind(state, output_index, kind)?; + if pending.final_value.is_some() { + return Err(tool_input_error( + "tool_input_delta_after_done", + format!( + "OpenAI Responses {} delta arrived after its final value at output index {output_index}", + kind.label() + ), + )); + } + pending.saw_delta = true; + pending.streamed.push_str(&delta); + + Ok(ensure_and_delta( + state, + SlotKey::OutputItem(output_index), + BlockType::ToolUse, + BlockMetadata::ToolUse { + id: String::new(), + name: String::new(), + }, + DeltaContent::InputJson(delta), + )) +} + +fn reconcile_tool_input( + state: &mut OpenAIResponsesState, + output_index: usize, + kind: ToolInputKind, + final_value: String, + metadata: BlockMetadata, +) -> Result, ClientError> { + let pending = ensure_tool_input_kind(state, output_index, kind)?; + if pending + .final_value + .as_ref() + .is_some_and(|existing| existing != &final_value) + { + return Err(tool_input_error( + "tool_input_final_mismatch", + format!( + "OpenAI Responses authoritative {} values disagree at output index {output_index}", + kind.label() + ), + )); + } + if pending.saw_delta && pending.streamed != final_value { + return Err(tool_input_error( + "tool_input_delta_mismatch", + format!( + "OpenAI Responses streamed {} does not match its authoritative final value at output index {output_index}", + kind.label() + ), + )); + } + + let should_emit = !pending.saw_delta; + if should_emit { + pending.saw_delta = true; + pending.streamed.clone_from(&final_value); + } + pending.final_value = Some(final_value.clone()); + + if should_emit { + Ok(ensure_and_delta( + state, + SlotKey::OutputItem(output_index), + BlockType::ToolUse, + metadata, + DeltaContent::InputJson(final_value), + )) + } else { + Ok(Vec::new()) + } +} + +fn finish_tool_output_item( + state: &mut OpenAIResponsesState, + output_index: usize, + kind: ToolInputKind, + call_id: String, + name: String, + final_value: Option, +) -> Result, ClientError> { + let metadata = BlockMetadata::ToolUse { id: call_id, name }; + let mut out = if let Some(final_value) = final_value { + reconcile_tool_input(state, output_index, kind, final_value, metadata.clone())? + } else { + let pending = ensure_tool_input_kind(state, output_index, kind)?; + if pending.final_value.is_none() { + return Err(tool_input_error( + "tool_input_final_missing", + format!( + "OpenAI Responses {} ended without an authoritative final value at output index {output_index}", + kind.label() + ), + )); + } + Vec::new() + }; + + let (info, just_created) = + state.get_or_allocate(SlotKey::OutputItem(output_index), BlockType::ToolUse); + if just_created { + out.push(Event::BlockStart(BlockStart { + index: info.flat_index, + block_type: BlockType::ToolUse, + metadata, + })); + } + state.slots.remove(&SlotKey::OutputItem(output_index)); + state.pending_tool_inputs.remove(&output_index); + out.push(Event::BlockStop(BlockStop { + index: info.flat_index, + block_type: info.block_type, + stop_reason: None, + reasoning: None, + })); + Ok(out) +} + +fn ensure_tool_input_kind( + state: &mut OpenAIResponsesState, + output_index: usize, + kind: ToolInputKind, +) -> Result<&mut PendingToolInput, ClientError> { + let pending = state + .pending_tool_inputs + .entry(output_index) + .or_insert_with(|| PendingToolInput::new(kind)); + if pending.kind != kind { + return Err(tool_input_error( + "tool_input_kind_mismatch", + format!("OpenAI Responses tool input kind changed at output index {output_index}"), + )); + } + Ok(pending) +} + +fn tool_input_error(code: &str, message: String) -> ClientError { + ClientError::Api { + status: None, + code: Some(code.to_string()), + message, + retry_after: None, + } +} + fn response_failure_diagnostic(event_type: &str, ev: ResponseFailed) -> (Option, String) { let mut diagnostic = Map::new(); diagnostic.insert("event".to_string(), Value::String(event_type.to_string())); @@ -855,6 +1137,26 @@ mod tests { parse_sse(event_type, data, state).unwrap() } + fn input_json_deltas(events: &[Event]) -> Vec<&str> { + events + .iter() + .filter_map(|event| match event { + Event::BlockDelta(BlockDelta { + delta: DeltaContent::InputJson(delta), + .. + }) => Some(delta.as_str()), + _ => None, + }) + .collect() + } + + fn api_error_code(error: ClientError) -> Option { + match error { + ClientError::Api { code, .. } => code, + other => panic!("expected ClientError::Api, got {other:?}"), + } + } + #[test] fn created_emits_status_started() { let (events, _) = run("response.created", r#"{"response":{}}"#); @@ -981,6 +1283,19 @@ mod tests { if let Event::BlockDelta(d) = &ev[0] { assert!(matches!(&d.delta, DeltaContent::InputJson(j) if j == "{\"x\":")); } + let ev = with( + &mut state, + "response.function_call_arguments.delta", + r#"{"output_index":1,"item_id":"fc1","delta":"1}"}"#, + ); + assert_eq!(ev.len(), 1); + assert!(matches!( + &ev[0], + Event::BlockDelta(BlockDelta { + delta: DeltaContent::InputJson(j), + .. + }) if j == "1}" + )); // output_item.done → BlockStop let ev = with( &mut state, @@ -1012,6 +1327,155 @@ mod tests { } } + #[test] + fn function_call_done_payload_emits_missing_final_input_before_stop() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_1","name":"lookup"}}"#, + ); + + let events = with( + &mut state, + "response.output_item.done", + r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"city\":\"Tokyo\"}"}}"#, + ); + + assert_eq!(input_json_deltas(&events), vec![r#"{"city":"Tokyo"}"#]); + assert!(matches!(events.last(), Some(Event::BlockStop(_)))); + } + + #[test] + fn function_call_arguments_done_accepts_exact_stream_without_duplication() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_1","name":"lookup"}}"#, + ); + let delta = with( + &mut state, + "response.function_call_arguments.delta", + r#"{"output_index":0,"delta":"{\"city\":\"Tokyo\"}"}"#, + ); + assert_eq!(input_json_deltas(&delta), vec![r#"{"city":"Tokyo"}"#]); + + let done = with( + &mut state, + "response.function_call_arguments.done", + r#"{"output_index":0,"arguments":"{\"city\":\"Tokyo\"}"}"#, + ); + assert!(done.is_empty()); + let output_done = with( + &mut state, + "response.output_item.done", + r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"city\":\"Tokyo\"}"}}"#, + ); + assert!(input_json_deltas(&output_done).is_empty()); + assert!(matches!(output_done.as_slice(), [Event::BlockStop(_)])); + } + + #[test] + fn function_call_arguments_done_rejects_stream_mismatch() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_1","name":"lookup"}}"#, + ); + with( + &mut state, + "response.function_call_arguments.delta", + r#"{"output_index":0,"delta":"{\"city\":\"Kyoto\"}"}"#, + ); + + let error = parse_sse( + "response.function_call_arguments.done", + r#"{"output_index":0,"arguments":"{\"city\":\"Tokyo\"}"}"#, + &mut state, + ) + .unwrap_err(); + assert_eq!( + api_error_code(error).as_deref(), + Some("tool_input_delta_mismatch") + ); + } + + #[test] + fn custom_tool_input_done_event_emits_missing_final_input_once() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell"}}"#, + ); + + let done = with( + &mut state, + "response.custom_tool_call_input.done", + r#"{"output_index":0,"input":"echo ready"}"#, + ); + assert_eq!(input_json_deltas(&done), vec!["echo ready"]); + let output_done = with( + &mut state, + "response.output_item.done", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell","input":"echo ready"}}"#, + ); + assert!(input_json_deltas(&output_done).is_empty()); + assert!(matches!(output_done.as_slice(), [Event::BlockStop(_)])); + } + + #[test] + fn custom_tool_done_payload_accepts_exact_stream_without_duplication() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell"}}"#, + ); + let delta = with( + &mut state, + "response.custom_tool_call_input.delta", + r#"{"output_index":0,"delta":"echo ready"}"#, + ); + assert_eq!(input_json_deltas(&delta), vec!["echo ready"]); + + let output_done = with( + &mut state, + "response.output_item.done", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell","input":"echo ready"}}"#, + ); + assert!(input_json_deltas(&output_done).is_empty()); + assert!(matches!(output_done.as_slice(), [Event::BlockStop(_)])); + } + + #[test] + fn custom_tool_done_payload_rejects_stream_mismatch() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell"}}"#, + ); + with( + &mut state, + "response.custom_tool_call_input.delta", + r#"{"output_index":0,"delta":"echo stale"}"#, + ); + + let error = parse_sse( + "response.output_item.done", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell","input":"echo ready"}}"#, + &mut state, + ) + .unwrap_err(); + assert_eq!( + api_error_code(error).as_deref(), + Some("tool_input_delta_mismatch") + ); + } + #[test] fn reasoning_text_delta_emits_thinking() { let mut state = OpenAIResponsesState::default(); diff --git a/crates/agen/src/llm_client/types.rs b/crates/agen/src/llm_client/types.rs index 101e97731..50ed7506a 100644 --- a/crates/agen/src/llm_client/types.rs +++ b/crates/agen/src/llm_client/types.rs @@ -104,7 +104,14 @@ pub enum Item { name: String, /// Tool arguments as JSON string arguments: String, - /// Item status + /// Provider/model order among tool calls in this assistant response. + #[serde(default, skip_serializing_if = "Option::is_none")] + call_index: Option, + /// Durable identity of a tool execution attempt that has started. + #[serde(default, skip_serializing_if = "Option::is_none")] + execution_id: Option, + /// Item status. `InProgress` on a persisted ToolCall means provider + /// execution ownership was transferred and must never be inferred pending. #[serde(skip_serializing_if = "Option::is_none")] status: Option, }, @@ -234,6 +241,8 @@ impl Item { call_id: call_id.into(), name: name.into(), arguments: arguments.into(), + call_index: None, + execution_id: None, status: None, } } @@ -247,6 +256,25 @@ impl Item { Self::tool_call(call_id, name, arguments.to_string()) } + /// Attach response-local ordering and, when execution has started, its + /// durable attempt identity to a tool call. + pub fn with_tool_execution_metadata( + mut self, + call_index: usize, + execution_id: Option, + ) -> Self { + if let Self::ToolCall { + call_index: item_call_index, + execution_id: item_execution_id, + .. + } = &mut self + { + *item_call_index = Some(call_index); + *item_execution_id = execution_id; + } + self + } + /// Create a tool result item with summary only (no content). pub fn tool_result(call_id: impl Into, summary: impl Into) -> Self { Self::tool_result_item(call_id, summary, None, false) diff --git a/crates/agen/src/timeline/mod.rs b/crates/agen/src/timeline/mod.rs index a1bd24cb8..b016c2592 100644 --- a/crates/agen/src/timeline/mod.rs +++ b/crates/agen/src/timeline/mod.rs @@ -21,6 +21,7 @@ pub use event::*; pub use text_block_collector::TextBlockCollector; pub use thinking_block_collector::ThinkingBlockCollector; pub use timeline::Timeline; +pub(crate) use tool_call_collector::CollectedToolCall; pub use tool_call_collector::ToolCallCollector; // 型定義からのre-export diff --git a/crates/agen/src/timeline/tool_call_collector.rs b/crates/agen/src/timeline/tool_call_collector.rs index 3c2ed0ca2..3663a8545 100644 --- a/crates/agen/src/timeline/tool_call_collector.rs +++ b/crates/agen/src/timeline/tool_call_collector.rs @@ -17,10 +17,19 @@ pub struct CollectorState { /// 現在のツール呼び出し情報 (ブロック進行中) current_id: Option, current_name: Option, + /// Zero-based order among tool calls in this assistant response. + call_index: Option, /// 蓄積中のJSON入力 input_json_buffer: String, } +/// One provider-completed call together with its stable model order. +#[derive(Debug, Clone)] +pub(crate) struct CollectedToolCall { + pub call_index: usize, + pub call: ToolCall, +} + /// ToolCallCollector - ToolUseブロックハンドラ /// /// Timelineに登録してToolUseブロックイベントを受信し、 @@ -28,7 +37,8 @@ pub struct CollectorState { #[derive(Clone)] pub struct ToolCallCollector { /// 収集されたToolCall - collected: Arc>>, + collected: Arc>>, + next_call_index: Arc>, } impl ToolCallCollector { @@ -36,18 +46,38 @@ impl ToolCallCollector { pub fn new() -> Self { Self { collected: Arc::new(Mutex::new(Vec::new())), + next_call_index: Arc::new(Mutex::new(0)), } } /// 収集されたToolCallを取得してクリア pub fn take_collected(&self) -> Vec { + self.take_collected_with_index() + .into_iter() + .map(|collected| collected.call) + .collect() + } + + /// Drain completed calls while preserving their model-returned order. + pub(crate) fn take_collected_with_index(&self) -> Vec { let mut guard = self.collected.lock().unwrap(); std::mem::take(&mut *guard) } /// 収集されたToolCallの参照を取得 pub fn collected(&self) -> Vec { - self.collected.lock().unwrap().clone() + self.collected + .lock() + .unwrap() + .iter() + .map(|collected| collected.call.clone()) + .collect() + } + + /// Start one response-local tool-call ordering scope. + pub(crate) fn begin_response(&self) { + self.collected.lock().unwrap().clear(); + *self.next_call_index.lock().unwrap() = 0; } /// 収集されたToolCallがあるかどうか @@ -75,6 +105,11 @@ impl Handler for ToolCallCollector { ToolUseBlockEvent::Start(start) => { scope.current_id = Some(start.id.clone()); scope.current_name = Some(start.name.clone()); + if scope.call_index.is_none() { + let mut next = self.next_call_index.lock().unwrap(); + scope.call_index = Some(*next); + *next += 1; + } scope.input_json_buffer.clear(); } ToolUseBlockEvent::InputJsonDelta(delta) => { @@ -98,8 +133,15 @@ impl Handler for ToolCallCollector { }; let tool_call = ToolCall { id, name, input }; + let call_index = scope + .call_index + .take() + .expect("tool-use stop follows a started collector scope"); - self.collected.lock().unwrap().push(tool_call); + self.collected.lock().unwrap().push(CollectedToolCall { + call_index, + call: tool_call, + }); } scope.input_json_buffer.clear(); } diff --git a/crates/agen/tests/common/mod.rs b/crates/agen/tests/common/mod.rs index 5f4d4e941..277f3ec13 100644 --- a/crates/agen/tests/common/mod.rs +++ b/crates/agen/tests/common/mod.rs @@ -7,7 +7,7 @@ use std::pin::Pin; use std::sync::{Arc, Mutex}; use agen::llm_client::event::{BlockType, DeltaContent, Event}; -use agen::llm_client::{ClientError, LlmClient, Request}; +use agen::llm_client::{ClientError, LlmClient, Request, ToolCallCompletionSupport}; use agen::timeline::{Handler, TextBlockEvent, TextBlockKind, Timeline}; use async_trait::async_trait; use futures::Stream; @@ -20,6 +20,7 @@ pub struct MockLlmClient { responses: Arc>>, call_count: Arc, requests: Arc>>, + completion_support: ToolCallCompletionSupport, } impl MockLlmClient { @@ -32,6 +33,7 @@ impl MockLlmClient { responses: Arc::new(responses), call_count: Arc::new(AtomicUsize::new(0)), requests: Arc::new(Mutex::new(Vec::new())), + completion_support: ToolCallCompletionSupport::ResponseComplete, } } @@ -47,6 +49,11 @@ impl MockLlmClient { pub fn requests(&self) -> Vec { self.requests.lock().unwrap().clone() } + + pub fn with_completion_support(mut self, support: ToolCallCompletionSupport) -> Self { + self.completion_support = support; + self + } } #[async_trait] @@ -55,6 +62,10 @@ impl LlmClient for MockLlmClient { Box::new(self.clone()) } + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + self.completion_support + } + async fn stream( &self, request: Request, diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index f38f9f7cc..f6d445fb0 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -7,8 +7,9 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use agen::interceptor::{ - Interceptor, InterceptorError, InterceptorErrorCategory, InterceptorPhase, InterceptorResult, - PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo, + AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorErrorCategory, + InterceptorPhase, InterceptorResult, PostToolAction, PreToolAction, ToolCallInfo, + ToolResultInfo, TurnEndAction, }; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::{ @@ -307,6 +308,46 @@ impl LlmClient for ControlledStreamClient { } } +#[derive(Clone)] +struct BoundedStreamClient { + receiver: Arc>>>>, +} + +impl BoundedStreamClient { + fn new(capacity: usize) -> (Self, tokio::sync::mpsc::Sender>) { + let (tx, rx) = tokio::sync::mpsc::channel(capacity); + ( + Self { + receiver: Arc::new(Mutex::new(Some(rx))), + }, + tx, + ) + } +} + +#[async_trait] +impl LlmClient for BoundedStreamClient { + async fn stream(&self, _request: Request) -> Result { + let mut rx = self + .receiver + .lock() + .unwrap() + .take() + .expect("bounded stream is consumed once"); + 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 { + ToolCallCompletionSupport::PerBlock + } +} + #[derive(Clone)] struct BarrierTool { name: &'static str, @@ -400,6 +441,141 @@ async fn default_dispatch_waits_for_response_completion() { assert_eq!(probe.starts.load(Ordering::SeqCst), 1); } +#[tokio::test] +async fn turn_end_policy_incompatibility_visibly_falls_back_before_side_effects() { + struct TurnEndPause; + + #[async_trait] + impl Interceptor for TurnEndPause { + async fn on_assistant_turn_end( + &self, + _context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { + Ok(TurnEndAction::Pause) + } + } + + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = BarrierTool::new("turn_end_fallback"); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(TurnEndPause); + engine.register_tool(tool.definition()); + let (warning_tx, mut warning_rx) = tokio::sync::mpsc::unbounded_channel(); + engine.on_warning(move |warning| { + let _ = warning_tx.send(warning.to_string()); + }); + let (ping_tx, mut ping_rx) = tokio::sync::mpsc::unbounded_channel(); + engine.on_stream_event(move |_, _, event| { + if matches!(event, Event::Ping(_)) { + let _ = ping_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_turn_end_fallback", + "turn_end_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(); + ping_rx.recv().await.unwrap(); + assert_eq!(probe.starts.load(Ordering::SeqCst), 0); + assert!( + warning_rx + .recv() + .await + .is_some_and(|warning| warning.contains("whole-response turn-end gate")) + ); + + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + let (output, _history) = run.await.unwrap(); + assert!(matches!(output.result, EngineRunExit::Paused)); + assert_eq!(probe.starts.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_pause() { + struct BlockingEarlyPolicy { + entered: Arc, + release: Arc, + } + + #[async_trait] + impl Interceptor for BlockingEarlyPolicy { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn pre_tool_call( + &self, + _info: &mut ToolCallInfo<'_>, + ) -> InterceptorResult { + self.entered.notify_one(); + self.release.notified().await; + Ok(PreToolAction::Continue) + } + } + + let (client, tx) = BoundedStreamClient::new(1); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let tool = SlowTool::new("blocked_policy_tool", 0); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(BlockingEarlyPolicy { + entered: entered.clone(), + release, + }); + 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_blocked_policy", + "blocked_policy_tool", + ))) + .await + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))) + .await + .unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), entered.notified()) + .await + .expect("pre-tool policy entered"); + + tx.send(Ok(Event::ping())).await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), tx.send(Ok(Event::ping()))) + .await + .expect("background pump continues draining the bounded provider stream") + .unwrap(); + pause.send(()).await.unwrap(); + drop(tx); + + let (output, _history) = run.await.unwrap(); + assert!(matches!(output.result, EngineRunExit::Paused)); + assert_eq!(probe.call_count(), 0); +} + #[tokio::test] async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_result() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); @@ -457,6 +633,146 @@ async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_ assert_eq!((calls, results), (2, 2)); } +#[tokio::test] +async fn early_terminal_result_commits_while_response_stream_is_open() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = BarrierTool::new("commit_while_streaming"); + 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 (committed_tx, mut committed_rx) = tokio::sync::mpsc::unbounded_channel(); + engine.on_history_append(move |item| { + if matches!(item, Item::ToolResult { call_id, .. } if call_id == "call_commit_open") { + let _ = committed_tx.send(()); + } + 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_commit_open", + "commit_while_streaming", + ))) + .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 while response remains open"); + probe.release.notify_one(); + tokio::time::timeout(Duration::from_secs(1), committed_rx.recv()) + .await + .expect("terminal result is durably committed before response release") + .expect("commit notification channel remains open"); + + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + let _ = run.await.unwrap(); +} + +#[tokio::test] +async fn reverse_stop_order_preserves_model_call_indexes() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let contexts = Arc::new(Mutex::new(Vec::new())); + let mut engine = Engine::new(client); + engine.set_max_turns(Some(1)); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(ContextRecordingTool::new("ordered_a", contexts.clone()).definition()); + engine.register_tool(ContextRecordingTool::new("ordered_b", contexts.clone()).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_order_a", "ordered_a"))) + .unwrap(); + tx.send(Ok(Event::tool_use_start(1, "call_order_b", "ordered_b"))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(1, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(1))).unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + let _ = run.await.unwrap(); + + let by_id = contexts + .lock() + .unwrap() + .iter() + .map(|context| (context.call_id.clone(), context.call_index)) + .collect::>(); + assert_eq!(by_id.get("call_order_a"), Some(&0)); + assert_eq!(by_id.get("call_order_b"), Some(&1)); +} + +#[tokio::test] +async fn reverse_stop_order_projects_next_provider_request_in_model_order() { + let first_response = vec![ + Event::tool_use_start(0, "call_project_a", "project_a"), + Event::tool_use_start(1, "call_project_b", "project_b"), + Event::tool_input_delta(1, r#"{}"#), + Event::tool_use_stop(1), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]; + let second_response = vec![ + Event::text_delta(0, "done"), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]; + let client = MockLlmClient::with_responses(vec![first_response, second_response]) + .with_completion_support(ToolCallCompletionSupport::PerBlock); + let probe = client.clone(); + let contexts = Arc::new(Mutex::new(Vec::new())); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(ContextRecordingTool::new("project_a", contexts.clone()).definition()); + engine.register_tool(ContextRecordingTool::new("project_b", contexts).definition()); + let mut history = History::new(); + + let _ = engine.run(&mut history, "run").await; + + let requests = probe.requests(); + assert_eq!(requests.len(), 2); + let tool_order = requests[1] + .items + .iter() + .filter_map(|item| match item { + Item::ToolCall { call_id, .. } => Some(format!("call:{call_id}")), + Item::ToolResult { call_id, .. } => Some(format!("result:{call_id}")), + _ => None, + }) + .collect::>(); + assert_eq!( + tool_order, + [ + "call:call_project_a", + "call:call_project_b", + "result:call_project_a", + "result:call_project_b", + ] + ); +} + #[tokio::test] async fn later_history_failure_terminalizes_already_started_early_sibling() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); @@ -525,6 +841,108 @@ async fn later_history_failure_terminalizes_already_started_early_sibling() { ))); } +#[tokio::test] +async fn persisted_started_call_is_closed_unknown_and_never_resumed() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = SlowTool::new("must_not_resume_started", 0); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_max_turns(Some(1)); + engine.register_tool(tool.definition()); + let mut history = History::from_items(vec![ + Item::user_message("run"), + Item::tool_call_json( + "call_started_before_restore", + "must_not_resume_started", + serde_json::json!({}), + ) + .with_tool_execution_metadata( + 0, + Some("restored-batch:call_started_before_restore".to_string()), + ) + .with_status(agen::llm_client::ItemStatus::InProgress), + ]); + + let run = tokio::spawn(async move { + let output = engine.resume(&mut history).await; + (output, history) + }); + 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, + disposition: ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call_started_before_restore" + ))); +} + +#[tokio::test] +async fn abort_during_early_admission_commits_synthetic_terminal() { + struct AbortAdmission; + + #[async_trait] + impl Interceptor for AbortAdmission { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn pre_tool_call( + &self, + _info: &mut ToolCallInfo<'_>, + ) -> InterceptorResult { + Ok(PreToolAction::Abort("denied after completion".to_string())) + } + } + + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = SlowTool::new("must_not_run_aborted", 0); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(AbortAdmission); + 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_aborted_admission", + "must_not_run_aborted", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + drop(tx); + let (output, history) = run.await.unwrap(); + + assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(EngineError::Aborted(ref reason))) + if reason == "denied after completion" + )); + assert_eq!(probe.call_count(), 0); + assert!(history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { + call_id, + summary, + .. + } if call_id == "call_aborted_admission" && summary.contains("denied after completion") + ))); +} + #[tokio::test] async fn malformed_completed_arguments_commit_error_without_executing_tool() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); diff --git a/crates/session-store/src/logged_item.rs b/crates/session-store/src/logged_item.rs index f6547265b..6855b0973 100644 --- a/crates/session-store/src/logged_item.rs +++ b/crates/session-store/src/logged_item.rs @@ -7,13 +7,14 @@ //! save / replay boundaries. //! //! Fields kept here are limited to what is needed to reconstruct a worker -//! `Item` for replay. `id` and `status` annotations are intentionally dropped -//! (they are output-side metadata; replayed items synthesize fresh `None`). +//! `Item` for replay. `id` annotations are intentionally dropped because they +//! are output-side metadata. ToolCall execution status/order are retained because +//! they fence durable resume from replaying an already-started side effect. //! `Reasoning::encrypted_content` is preserved because OpenAI Responses ZDR //! requires it on stateless re-send. use agen::{ - llm_client::types::{ContentPart, Item, Role}, + llm_client::types::{ContentPart, Item, ItemStatus, Role}, tool::{Attachment, ImageAttachment, ToolResultDisposition}, }; use base64::{Engine as _, engine::general_purpose::STANDARD}; @@ -53,6 +54,12 @@ pub enum LoggedItem { call_id: String, name: String, arguments: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + call_index: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + execution_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + status: Option, }, ToolResult { call_id: String, @@ -119,11 +126,17 @@ impl From<&Item> for LoggedItem { call_id, name, arguments, + call_index, + execution_id, + status, .. } => Self::ToolCall { call_id: call_id.clone(), name: name.clone(), arguments: arguments.clone(), + call_index: *call_index, + execution_id: execution_id.clone(), + status: *status, }, Item::ToolResult { call_id, @@ -176,12 +189,17 @@ impl From for Item { call_id, name, arguments, + call_index, + execution_id, + status, } => Item::ToolCall { id: None, call_id, name, arguments, - status: None, + call_index, + execution_id, + status, }, LoggedItem::ToolResult { call_id, @@ -327,7 +345,9 @@ mod tests { #[test] fn round_trip_tool_call() { - let original = Item::tool_call("call_42", "get_weather", r#"{"city":"Tokyo"}"#); + let original = Item::tool_call("call_42", "get_weather", r#"{"city":"Tokyo"}"#) + .with_tool_execution_metadata(3, Some("batch:call_42".to_string())) + .with_status(ItemStatus::InProgress); let logged: LoggedItem = (&original).into(); let json = serde_json::to_string(&logged).unwrap(); let parsed: LoggedItem = serde_json::from_str(&json).unwrap(); @@ -336,11 +356,17 @@ mod tests { call_id, name, arguments, + call_index, + execution_id, + status, .. } => { assert_eq!(call_id, "call_42"); assert_eq!(name, "get_weather"); assert_eq!(arguments, r#"{"city":"Tokyo"}"#); + assert_eq!(call_index, Some(3)); + assert_eq!(execution_id.as_deref(), Some("batch:call_42")); + assert_eq!(status, Some(ItemStatus::InProgress)); } other => panic!("unexpected variant: {other:?}"), } diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index 752187064..72ec0a140 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -356,6 +356,7 @@ fn project_item(item: &LoggedItem) -> Option { call_id, name, arguments, + .. } => Some(SessionSnapshotEntryData::ToolCall { call_id: call_id.clone(), name: name.clone(), diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 802d0c7d7..2b403cd39 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -363,6 +363,10 @@ impl WorkerInterceptor { #[async_trait] impl Interceptor for WorkerInterceptor { + fn supports_early_tool_dispatch(&self) -> bool { + self.registry.on_turn_end.is_empty() + } + async fn on_prompt_submit( &self, context: PromptSubmitContext<'_, SessionHistoryMetadata>, diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 8cac5d7e5..f0b05eeec 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -1371,6 +1371,9 @@ mod tests { call_id: format!("call-{index}"), name: name.to_string(), arguments: "{}".to_string(), + call_index: None, + execution_id: None, + status: None, }, ), }); diff --git a/docs/manifest.toml b/docs/manifest.toml index 274a792e8..c1bcccc9f 100644 --- a/docs/manifest.toml +++ b/docs/manifest.toml @@ -123,8 +123,9 @@ ref = "anthropic/claude-sonnet-4-6" # 任意。デフォルト: "after_response"。 # "on_tool_call_complete" は provider が個別 ToolCall の完成を保証する場合だけ、 -# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider では -# 応答完了まで明示的に遅延する。これは直列実行モードではなく、先行 call の完了を +# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider、または +# 応答全体を見て実行前 Pause を判断する turn-end hook がある場合は、応答完了まで +# 警告付きで明示的に遅延する。これは直列実行モードではなく、先行 call の完了を # 待たず後続 call も開始するため、完了順・DB 反映順・採番順は保証しない。 # 早期開始後に stream が切断された場合は自動 continuation せず、安全に terminalize して停止する。 # tool_call_dispatch = "on_tool_call_complete"