diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 1fa4fafd0..45a28293d 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; @@ -23,9 +24,9 @@ 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, types::parse_tool_arguments, + transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT, }, state::{EngineState, Locked, Mutable}, timeline::event::{ErrorEvent, StatusEvent, UsageEvent}, @@ -78,59 +79,198 @@ 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. /// /// Runtime history intentionally retains completion order so every result can -/// be committed without waiting for slower siblings. The provider projection -/// is deterministic within each contiguous result batch and does not rewrite -/// the committed transcript. +/// be committed without waiting for slower siblings. An internal durable +/// response boundary lets this projection move every early result after the +/// complete assistant response, even when trailing text/reasoning arrived after +/// a fast result. Boundaries are never exposed to providers. struct ProviderHistoryProjection { items: Vec, original_to_projected_index: Vec, } -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; +fn normalize_tool_run(run: &mut [(usize, Item)], call_order: &HashMap) { + run.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"), + }); +} +fn normalize_legacy_history( + materialized: &mut [(usize, Item)], + call_order: &HashMap, +) { + 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, + 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; + } + normalize_tool_run(&mut materialized[start..index], call_order); + } +} + +fn normalize_response_history(response: &mut [(usize, Item)], call_order: &HashMap) { + // Provider-confirmed call completions can be observed out of index order. + // Reorder only the values occupying ToolCall slots so surrounding assistant + // text/reasoning retains its streamed order. + let mut calls = response + .iter() + .filter_map(|(_, item)| match item { + Item::ToolCall { .. } => Some(item.clone()), + _ => None, + }) + .collect::>(); + calls.sort_by_key(|item| match item { + Item::ToolCall { + call_id, + call_index, + .. + } => call_index + .or_else(|| call_order.get(call_id).copied()) + .unwrap_or(usize::MAX), + _ => unreachable!("filtered ToolCall list"), + }); + let mut calls = calls.into_iter(); + for (_, item) in response.iter_mut() { + if matches!(item, Item::ToolCall { .. }) { + *item = calls.next().expect("one replacement per ToolCall slot"); } } - let mut original_to_projected_index = vec![0; materialized.len()]; + // A host-injected user/system item is a hard provider turn boundary. Do not + // move a result across it; normalize each assistant-side region separately. + let mut start = 0usize; + for index in 0..=response.len() { + let at_boundary = index == response.len() + || matches!( + response[index].1, + Item::Message { + role: crate::llm_client::types::Role::User + | crate::llm_client::types::Role::System, + .. + } + ); + if at_boundary { + response[start..index].sort_by_key(|(_, item)| match item { + Item::ToolResult { call_id, .. } => ( + 1usize, + call_order.get(call_id).copied().unwrap_or(usize::MAX), + ), + _ => (0usize, 0usize), + }); + start = index.saturating_add(1); + } + } +} + +fn materialize_provider_history(items: &[Item]) -> ProviderHistoryProjection { + let mut call_order = HashMap::::new(); + let mut next_call_order = 0usize; + for item in items { + 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 first_boundary = items + .iter() + .position(|item| matches!(item, Item::AssistantResponseBoundary { .. })); + let legacy_end = first_boundary.unwrap_or(items.len()); + let mut legacy = items[..legacy_end] + .iter() + .cloned() + .enumerate() + .collect::>(); + normalize_legacy_history(&mut legacy, &call_order); + + let mut materialized = legacy; + let mut boundary_indices = Vec::new(); + let mut index = legacy_end; + while index < items.len() { + debug_assert!(matches!( + items[index], + Item::AssistantResponseBoundary { .. } + )); + boundary_indices.push(index); + let start = index + 1; + let end = items[start..] + .iter() + .position(|item| matches!(item, Item::AssistantResponseBoundary { .. })) + .map(|offset| start + offset) + .unwrap_or(items.len()); + let mut response = items[start..end] + .iter() + .cloned() + .enumerate() + .map(|(offset, item)| (start + offset, item)) + .collect::>(); + normalize_response_history(&mut response, &call_order); + materialized.extend(response); + index = end; + } + + let mut original_to_projected_index = vec![0; items.len()]; for (projected_index, (original_index, _)) in materialized.iter().enumerate() { original_to_projected_index[*original_index] = projected_index; } + for boundary_index in boundary_indices { + let preceding_projected = materialized + .iter() + .take_while(|(original_index, _)| *original_index < boundary_index) + .count() + .saturating_sub(1); + original_to_projected_index[boundary_index] = preceding_projected; + } ProviderHistoryProjection { items: materialized.into_iter().map(|(_, item)| item).collect(), @@ -196,6 +336,7 @@ pub struct EngineRunOutput { } /// Internal: tool execution result +#[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ToolExecutionResult { Completed, Paused, @@ -207,6 +348,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, + turn_id: usize, + calls: Vec<(usize, ToolCall)>, + 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, turn_id: usize) -> Self { + Self { + batch_id, + turn_id, + 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 +503,42 @@ enum StreamCompletion { Interrupted { reason: String }, } +struct StreamResponseOutput { + completion: StreamCompletion, + 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, @@ -362,6 +589,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) @@ -1161,44 +1390,530 @@ 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 = parse_tool_arguments(arguments); - 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: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + batch: &mut EarlyToolExecutionBatch, + tool_calls: Vec, + ) -> Result { + 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; + } + 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, + Some(context.execution_id()), + )?; + 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 { + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + Some(context.execution_id()), + )?; + batch.synthetic_results.push_back(ToolResult::error( + &tool_call.id, + format!("Tool not found: {}", tool_call.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 = 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"); + } + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + None, + )?; + batch.dispatched_any = true; + 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"); + } + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + None, + )?; + batch.dispatched_any = true; + return Err(EngineError::Cancelled); + } + }; + 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()); + } + 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(), meta, tool.clone(), context.clone()), + ); + + match action { + PreToolAction::Continue => { + 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()); + 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(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()); + } + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + Some(context.execution_id()), + )?; + batch.synthetic_results.push_back(result); + batch.dispatched_any = true; + } + PreToolAction::Skip => { + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + Some(context.execution_id()), + )?; + 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, + Some(context.execution_id()), + )?; + 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)?; + batch.dispatched_any = true; + 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.turn_id, + true, + &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, + annotate: &mut impl FnMut(&Item) -> Result, + batch: &mut EarlyToolExecutionBatch, + ) -> 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.turn_id, + false, + &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 + .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 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.turn_id, + false, + &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.turn_id, + false, + &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.turn_id, + true, + &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.turn_id, + true, + &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 @@ -1221,6 +1936,7 @@ impl Engine { // Map from tool call ID to (ToolCall, Meta, Tool, Context) // Retained because it's needed for PostToolCall hooks + let turn_id = self.turn_count.saturating_sub(1); let mut call_info_map = HashMap::new(); let mut synthetic_results = Vec::new(); let batch_id = format!("tool-batch-{}", self.tool_execution_batch_count); @@ -1231,6 +1947,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, @@ -1398,6 +2121,8 @@ impl Engine { annotate, result, None, + turn_id, + true, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -1428,6 +2153,8 @@ impl Engine { annotate, result, Some(&attempt_id), + turn_id, + true, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -1507,6 +2234,8 @@ impl Engine { annotate, result, Some(&attempt_id), + turn_id, + false, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -1533,6 +2262,8 @@ impl Engine { annotate, ToolResult::outcome_unknown(call_id), Some(attempt_id), + turn_id, + false, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -1581,6 +2312,8 @@ impl Engine { annotate: &mut impl FnMut(&Item) -> Result, mut tool_result: ToolResult, execution_attempt_id: Option<&str>, + turn_id: usize, + invoke_post_hook: bool, call_info_map: &HashMap< String, ( @@ -1669,10 +2402,10 @@ impl Engine { ); self.emit_tool_result(&tool_result); - if let Some((tool_call, meta, tool, context)) = call_info { + if invoke_post_hook && let Some((tool_call, meta, tool, context)) = call_info { let invocation = self.interceptor_invocation( InterceptorPhase::PostToolCall, - Some(self.turn_count.saturating_sub(1)), + Some(turn_id), Some(InterceptorCallId::Tool(tool_call.id.clone())), context.call_index, ); @@ -1685,16 +2418,28 @@ impl Engine { tool: tool.clone(), context: context.clone(), }; - let post_tool_action = - self.interceptor - .post_tool_call(&info) - .await - .map_err(|error| { + let post_tool_action = tokio::select! { + result = self.interceptor.post_tool_call(&info) => { + result.map_err(|error| { EngineError::from(InterceptorFailure::new( InterceptorPhase::PostToolCall, error, )) - })?; + })? + } + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!(call_id = %tool_call.id, "Paused while awaiting post-tool policy"); + } + return Err(EngineError::PauseRequested); + } + cancel = self.cancel_rx.recv() => { + if cancel.is_some() { + info!(call_id = %tool_call.id, "Cancelled while awaiting post-tool policy"); + } + return Err(EngineError::Cancelled); + } + }; if let PostToolAction::Abort(reason) = post_tool_action { return Err(EngineError::Aborted(reason)); } @@ -1717,11 +2462,35 @@ 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, + self.turn_count, + false, + &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); @@ -1934,8 +2703,49 @@ 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 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 + && provider_supports_early + && interceptor_supports_early; + if self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete + && !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", + ); + } + if early_dispatch { + let response_id = format!("assistant-response-{current_turn}-{current_llm_call}"); + self.append_history_items( + history, + std::iter::once(Item::assistant_response_boundary(response_id)), + annotate, + )?; + } + 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 +2754,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 +2840,95 @@ 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 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)?; + 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 + .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, current_turn)); + } + let admission = self + .admit_early_tool_calls( + history, + annotate, + early_tools.as_mut().expect("batch initialized"), + remaining_collected_calls, + ) + .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); + } + 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| { + 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; + } let assistant_invocation = self.interceptor_invocation( InterceptorPhase::AssistantTurnEnd, @@ -2011,29 +2946,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 +3194,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(), @@ -2232,12 +3205,16 @@ 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 response_completed = false; let mut event_count: usize = 0; loop { tokio::select! { - event_result = stream.next() => { + event_result = stream.recv() => { match event_result { Some(result) => { match &result { @@ -2254,8 +3231,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, }); } }; @@ -2267,11 +3247,152 @@ impl Engine { json!({}), ); } + if matches!( + &event, + Event::Status(StatusEvent { + status: crate::llm_client::event::ResponseStatus::Completed, + }) + ) { + response_completed = true; + } self.emit_stream_event(turn, llm_call, &event); self.timeline.dispatch(&event); + if early_dispatch { + let collected = self.tool_call_collector.take_collected_with_index(); + 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, turn)); + } + let fresh_calls = { + let batch = early_tools.as_ref().expect("batch initialized"); + collected + .into_iter() + .filter(|collected| { + !batch.seen_call_ids.contains(&collected.call.id) + }) + .collect::>() + }; + if !fresh_calls.is_empty() { + // 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, + &[], + ); + 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, + annotate, + batch, + fresh_calls, + ) + .await + }; + match admission { + 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(); + 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(), @@ -2280,7 +3401,64 @@ impl Engine { })); } } - None => break, + None if !early_dispatch || response_completed => break, + None => { + // Providers that support only response-complete tool calls use + // stream EOF as their established completion boundary. Early + // dispatch is never enabled for them. Once per-block dispatch is + // active, however, transport EOF alone cannot prove that the + // assistant response completed after a side effect may have begun. + // Require the provider's explicit terminal status in that mode. + self.timeline.flush_usage(); + return Ok(StreamResponseOutput { + completion: StreamCompletion::Interrupted { + reason: "LLM stream ended before provider completion status" + .to_string(), + }, + early_tools, + }); + } + } + } + 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.turn_id, + true, + &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() => { @@ -2291,6 +3469,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 +3481,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 +3496,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 +3554,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 +3755,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 +3843,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 +4027,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, @@ -2964,6 +4169,7 @@ fn items_trace_payload( fn item_kind(item: &Item) -> &'static str { match item { Item::Message { .. } => "message", + Item::AssistantResponseBoundary { .. } => "assistant_response_boundary", Item::ToolCall { .. } => "tool_call", Item::ToolResult { .. } => "tool_result", Item::Reasoning { .. } => "reasoning", @@ -3013,6 +4219,86 @@ 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 provider_projection_keeps_complete_response_before_early_results() { + let items = vec![ + Item::assistant_response_boundary("response-1"), + 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::assistant_message("trailing text"), + Item::reasoning("trailing reasoning"), + 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::Message { .. } => "text".to_string(), + Item::Reasoning { .. } => "reasoning".to_string(), + Item::ToolCall { call_id, .. } => format!("call:{call_id}"), + Item::ToolResult { call_id, .. } => format!("result:{call_id}"), + Item::AssistantResponseBoundary { .. } => panic!("boundary leaked"), + }) + .collect::>(); + assert_eq!( + order, + [ + "call:call_first", + "text", + "reasoning", + "call:call_second", + "result:call_first", + "result:call_second", + ] + ); + assert_eq!( + projection.original_to_projected_index, + [0, 0, 5, 1, 2, 3, 4] + ); + } + #[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/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/request.rs b/crates/agen/src/llm_client/scheme/anthropic/request.rs index 62fac8d17..b546f4c94 100644 --- a/crates/agen/src/llm_client/scheme/anthropic/request.rs +++ b/crates/agen/src/llm_client/scheme/anthropic/request.rs @@ -263,6 +263,7 @@ impl AnthropicScheme { for (i, item) in items.iter().enumerate() { match item { + Item::AssistantResponseBoundary { .. } => {} Item::Message { role, content, .. } => { let parts: Vec = content .iter() 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/gemini/request.rs b/crates/agen/src/llm_client/scheme/gemini/request.rs index e2fac9179..c1eb6d82c 100644 --- a/crates/agen/src/llm_client/scheme/gemini/request.rs +++ b/crates/agen/src/llm_client/scheme/gemini/request.rs @@ -242,6 +242,7 @@ impl GeminiScheme { for item in items { match item { + Item::AssistantResponseBoundary { .. } => {} Item::Message { role, content, .. } => { // Flush pending parts self.flush_pending_parts( 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_chat/request.rs b/crates/agen/src/llm_client/scheme/openai_chat/request.rs index d971ba8b0..e8a87496e 100644 --- a/crates/agen/src/llm_client/scheme/openai_chat/request.rs +++ b/crates/agen/src/llm_client/scheme/openai_chat/request.rs @@ -221,6 +221,7 @@ impl OpenAIScheme { ); } match item { + Item::AssistantResponseBoundary { .. } => {} Item::Message { role, content, .. } => { // Flush pending tool calls self.flush_pending_assistant( 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/scheme/openai_responses/request.rs b/crates/agen/src/llm_client/scheme/openai_responses/request.rs index 63df50248..e26b60a51 100644 --- a/crates/agen/src/llm_client/scheme/openai_responses/request.rs +++ b/crates/agen/src/llm_client/scheme/openai_responses/request.rs @@ -254,6 +254,7 @@ fn convert_items_to_input(items: &[Item], supports_images: bool) -> Vec {} Item::Message { role, content, .. } => { let (role_str, text_variant): (&'static str, fn(String) -> InputContent) = match role { 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/llm_client/types.rs b/crates/agen/src/llm_client/types.rs index 101e97731..1a9d3ca24 100644 --- a/crates/agen/src/llm_client/types.rs +++ b/crates/agen/src/llm_client/types.rs @@ -93,6 +93,12 @@ pub enum Item { status: Option, }, + /// Internal durable boundary used to reconstruct one assistant response when + /// early ToolResults are committed before the response stream has ended. + /// Provider request projections remove this item before serialization. + #[doc(hidden)] + AssistantResponseBoundary { response_id: String }, + /// Tool call from the assistant ToolCall { /// Optional item ID @@ -104,7 +110,15 @@ 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 an execution or synthetic terminalization that + /// owns this call. Once present, restore must never infer the call pending. + #[serde(default, skip_serializing_if = "Option::is_none")] + execution_id: Option, + /// Item status. `InProgress` on a persisted ToolCall means execution or + /// terminalization ownership was transferred and must never be inferred pending. #[serde(skip_serializing_if = "Option::is_none")] status: Option, }, @@ -219,6 +233,14 @@ impl Item { } } + /// Create an internal durable assistant-response boundary. + #[doc(hidden)] + pub fn assistant_response_boundary(response_id: impl Into) -> Self { + Self::AssistantResponseBoundary { + response_id: response_id.into(), + } + } + // ======================================================================== // Tool call constructors // ======================================================================== @@ -234,6 +256,8 @@ impl Item { call_id: call_id.into(), name: name.into(), arguments: arguments.into(), + call_index: None, + execution_id: None, status: None, } } @@ -247,6 +271,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) @@ -377,6 +420,7 @@ impl Item { pub fn with_id(mut self, id: impl Into) -> Self { match &mut self { Self::Message { id: item_id, .. } => *item_id = Some(id.into()), + Self::AssistantResponseBoundary { .. } => {} Self::ToolCall { id: item_id, .. } => *item_id = Some(id.into()), Self::ToolResult { id: item_id, .. } => *item_id = Some(id.into()), Self::Reasoning { id: item_id, .. } => *item_id = Some(id.into()), @@ -388,6 +432,7 @@ impl Item { pub fn with_status(mut self, new_status: ItemStatus) -> Self { match &mut self { Self::Message { status, .. } => *status = Some(new_status), + Self::AssistantResponseBoundary { .. } => {} Self::ToolCall { status, .. } => *status = Some(new_status), Self::ToolResult { .. } => {} // Result items don't have status Self::Reasoning { status, .. } => *status = Some(new_status), @@ -403,6 +448,7 @@ impl Item { pub fn id(&self) -> Option<&str> { match self { Self::Message { id, .. } => id.as_deref(), + Self::AssistantResponseBoundary { .. } => None, Self::ToolCall { id, .. } => id.as_deref(), Self::ToolResult { id, .. } => id.as_deref(), Self::Reasoning { id, .. } => id.as_deref(), @@ -413,6 +459,7 @@ impl Item { pub fn item_type(&self) -> &'static str { match self { Self::Message { .. } => "message", + Self::AssistantResponseBoundary { .. } => "assistant_response_boundary", Self::ToolCall { .. } => "tool_call", Self::ToolResult { .. } => "tool_result", Self::Reasoning { .. } => "reasoning", 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/timeline.rs b/crates/agen/src/timeline/timeline.rs index f972ec31a..6dfe80b2c 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() } } @@ -379,7 +374,9 @@ pub struct Timeline { thinking_block_handlers: Vec>, tool_use_block_handlers: Vec>, - // 現在アクティブなブロック + // 現在アクティブなブロック。provider は index ごとに複数 scope を + // interleave できるため、単一の current block ではなく全 scope を追跡する。 + active_blocks: HashMap<(BlockType, usize), ()>, current_block: Option, // 1リクエスト内で受信した Usage event の集約バッファ。 @@ -405,6 +402,7 @@ impl Timeline { text_block_handlers: Vec::new(), thinking_block_handlers: Vec::new(), tool_use_block_handlers: Vec::new(), + active_blocks: HashMap::new(), current_block: None, pending_usage: None, } @@ -559,6 +557,8 @@ impl Timeline { } fn handle_block_start(&mut self, start: &BlockStart) { + self.active_blocks + .insert((start.block_type, start.index), ()); self.current_block = Some(start.block_type); let handlers = self.get_block_handlers_mut(start.block_type); @@ -572,10 +572,11 @@ impl Timeline { let block_type = delta.delta.block_type(); // OpenAIなどのプロバイダはBlockStartを送らない場合があるため、 - // Deltaが来たときにスコープがなければ暗黙的に開始する - if self.current_block.is_none() { - self.current_block = Some(block_type); - } + // Deltaが来たときにその index のスコープがなければ暗黙的に開始する + self.active_blocks + .entry((block_type, delta.index)) + .or_insert(()); + self.current_block = Some(block_type); let handlers = self.get_block_handlers_mut(block_type); for handler in handlers { @@ -593,7 +594,12 @@ impl Timeline { handler.dispatch_stop(stop); handler.end_scope(); } - self.current_block = None; + self.active_blocks.remove(&(stop.block_type, stop.index)); + self.current_block = self + .active_blocks + .keys() + .next() + .map(|(block_type, _)| *block_type); } fn handle_block_abort(&mut self, abort: &BlockAbort) { @@ -602,7 +608,12 @@ impl Timeline { handler.dispatch_abort(abort); handler.end_scope(); } - self.current_block = None; + self.active_blocks.remove(&(abort.block_type, abort.index)); + self.current_block = self + .active_blocks + .keys() + .next() + .map(|(block_type, _)| *block_type); } fn get_block_handlers_mut( @@ -622,19 +633,22 @@ impl Timeline { self.current_block } - /// 現在アクティブなブロックを中断する + /// 現在アクティブな全ブロックを中断する /// - /// キャンセルやエラー時に呼び出し、進行中のブロックに対して + /// キャンセルやエラー時に呼び出し、index ごとに進行中の全ブロックへ /// BlockAbortイベントを発火してスコープをクリーンアップする。 pub fn abort_current_block(&mut self) { - if let Some(block_type) = self.current_block { + let active = self.active_blocks.keys().copied().collect::>(); + for (block_type, index) in active { let abort = crate::timeline::event::BlockAbort { - index: 0, // インデックスは不明なので0 + index, block_type, reason: "Cancelled".to_string(), }; self.handle_block_abort(&abort); } + self.active_blocks.clear(); + self.current_block = None; } } @@ -679,6 +693,50 @@ mod tests { assert!(calls.lock().unwrap().is_empty()); } + #[test] + fn abort_clears_every_indexed_scope_before_indices_are_reused() { + #[derive(Default)] + struct Scope { + started: bool, + } + struct IndexedToolHandler { + reused_cleanly: Arc>>, + } + impl Handler for IndexedToolHandler { + type Scope = Scope; + + fn on_event(&mut self, scope: &mut Scope, event: &ToolUseBlockEvent) { + if matches!(event, ToolUseBlockEvent::Start(_)) { + self.reused_cleanly.lock().unwrap().push(!scope.started); + scope.started = true; + } + } + } + + let reused_cleanly = Arc::new(Mutex::new(Vec::new())); + let mut timeline = Timeline::new(); + timeline.on_tool_use_block(IndexedToolHandler { + reused_cleanly: reused_cleanly.clone(), + }); + + timeline.dispatch(&Event::tool_use_start(0, "first", "tool")); + timeline.dispatch(&Event::tool_use_start(1, "second", "tool")); + // Stopping one interleaved block must not hide its still-active sibling. + timeline.dispatch(&Event::tool_use_stop(0)); + assert_eq!(timeline.current_block(), Some(BlockType::ToolUse)); + timeline.abort_current_block(); + assert!(timeline.current_block().is_none()); + + // A later provider response may reuse the same content-block index. + timeline.dispatch(&Event::tool_use_start(1, "replacement", "tool")); + timeline.dispatch(&Event::tool_use_stop(1)); + + assert_eq!( + reused_cleanly.lock().unwrap().as_slice(), + [true, true, true] + ); + } + #[test] fn test_meta_event_dispatch() { // シンプルなテスト用構造体 diff --git a/crates/agen/src/timeline/tool_call_collector.rs b/crates/agen/src/timeline/tool_call_collector.rs index fa04b0ede..3663a8545 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}; @@ -18,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ブロックイベントを受信し、 @@ -29,7 +37,8 @@ pub struct CollectorState { #[derive(Clone)] pub struct ToolCallCollector { /// 収集されたToolCall - collected: Arc>>, + collected: Arc>>, + next_call_index: Arc>, } impl ToolCallCollector { @@ -37,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があるかどうか @@ -76,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) => { @@ -85,11 +119,29 @@ 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 }; + 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(); } @@ -144,6 +196,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/annotated_history_test.rs b/crates/agen/tests/annotated_history_test.rs index cf1be3c08..b7c391965 100644 --- a/crates/agen/tests/annotated_history_test.rs +++ b/crates/agen/tests/annotated_history_test.rs @@ -36,6 +36,7 @@ async fn run_preserves_item_annotations_without_projecting_them() { Role::Assistant => "assistant", Role::System => "system", }, + Item::AssistantResponseBoundary { .. } => "assistant_response_boundary", Item::ToolCall { .. } => "tool_call", Item::ToolResult { .. } => "tool_result", Item::Reasoning { .. } => "reasoning", 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 5df9f386e..279c65eeb 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -7,16 +7,21 @@ 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::{ + 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 +256,1437 @@ impl Tool for ContextRecordingTool { } } +#[derive(Clone)] +struct ControlledStreamClient { + receiver: Arc>>>>, + completion_support: ToolCallCompletionSupport, + stream_calls: Arc, +} + +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, + stream_calls: Arc::new(AtomicUsize::new(0)), + }, + tx, + ) + } + + fn stream_count(&self) -> usize { + self.stream_calls.load(Ordering::SeqCst) + } +} + +#[async_trait] +impl LlmClient for ControlledStreamClient { + async fn stream(&self, _request: Request) -> Result { + self.stream_calls.fetch_add(1, Ordering::SeqCst); + 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 MultiControlledStreamClient { + receivers: Arc< + Mutex< + std::collections::VecDeque< + tokio::sync::mpsc::UnboundedReceiver>, + >, + >, + >, +} + +impl MultiControlledStreamClient { + fn new( + response_count: usize, + ) -> ( + Self, + Vec>>, + ) { + let mut senders = Vec::with_capacity(response_count); + let mut receivers = std::collections::VecDeque::with_capacity(response_count); + for _ in 0..response_count { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + senders.push(tx); + receivers.push_back(rx); + } + ( + Self { + receivers: Arc::new(Mutex::new(receivers)), + }, + senders, + ) + } +} + +#[async_trait] +impl LlmClient for MultiControlledStreamClient { + async fn stream(&self, _request: Request) -> Result { + let mut rx = + self.receivers + .lock() + .unwrap() + .pop_front() + .ok_or_else(|| ClientError::Api { + status: Some(500), + code: Some("multi_controlled_stream_exhausted".to_string()), + message: "controlled responses exhausted".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 { + ToolCallCompletionSupport::PerBlock + } +} + +#[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, + 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 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, mut history) = run.await.unwrap(); + assert!(matches!(output.result, EngineRunExit::Paused)); + assert_eq!(probe.call_count(), 0); + assert!(history.items().any(|item| matches!( + item, + Item::ToolCall { + call_id, + execution_id: None, + status: Some(agen::llm_client::ItemStatus::Completed), + .. + } if call_id == "call_blocked_policy" + ))); + + let mut restored = Engine::new(MockLlmClient::new(vec![Event::Status(StatusEvent { + status: ResponseStatus::Completed, + })])); + restored.register_tool(probe.definition()); + let _ = restored.resume(&mut history).await; + assert_eq!(probe.call_count(), 1); + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_blocked_policy" + ))); +} + +#[tokio::test] +async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_cancel() { + 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("cancelled_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 cancel = engine.cancel_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_cancelled_policy", + "cancelled_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(); + cancel.send(()).await.unwrap(); + drop(tx); + + let (output, mut history) = tokio::time::timeout(Duration::from_secs(1), run) + .await + .expect("cancel does not deadlock the bounded stream pump") + .unwrap(); + assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) + )); + assert_eq!(probe.call_count(), 0); + assert!(history.items().any(|item| matches!( + item, + Item::ToolCall { + call_id, + execution_id: None, + status: Some(agen::llm_client::ItemStatus::Completed), + .. + } if call_id == "call_cancelled_policy" + ))); + + let mut restored = Engine::new(MockLlmClient::new(vec![Event::Status(StatusEvent { + status: ResponseStatus::Completed, + })])); + restored.register_tool(probe.definition()); + let _ = restored.resume(&mut history).await; + assert_eq!(probe.call_count(), 1); + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_cancelled_policy" + ))); +} + +#[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 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(); +} + +#[derive(Clone, Copy)] +enum BlockingPostControl { + Pause, + Cancel, +} + +async fn assert_blocking_post_hook_services_control(control: BlockingPostControl) { + struct BlockingPostPolicy { + entered: Arc, + release: Arc, + } + + #[async_trait] + impl Interceptor for BlockingPostPolicy { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn post_tool_call( + &self, + _info: &ToolResultInfo<'_>, + ) -> InterceptorResult { + self.entered.notify_one(); + self.release.notified().await; + Ok(PostToolAction::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("blocking_post_tool", 0); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(BlockingPostPolicy { + entered: entered.clone(), + release, + }); + engine.register_tool(tool.definition()); + let pause = engine.pause_sender(); + let cancel = engine.cancel_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_blocking_post", + "blocking_post_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("post-tool policy entered after the terminal result was committed"); + + 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 while post-tool policy blocks") + .unwrap(); + match control { + BlockingPostControl::Pause => pause.send(()).await.unwrap(), + BlockingPostControl::Cancel => cancel.send(()).await.unwrap(), + } + drop(tx); + + let (output, history) = tokio::time::timeout(Duration::from_secs(1), run) + .await + .expect("control request interrupts the blocked post-tool policy") + .unwrap(); + match control { + BlockingPostControl::Pause => assert!(matches!(output.result, EngineRunExit::Paused)), + BlockingPostControl::Cancel => assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) + )), + } + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_blocking_post" + ))); +} + +#[tokio::test] +async fn blocking_post_tool_policy_services_pause_after_result_commit() { + assert_blocking_post_hook_services_control(BlockingPostControl::Pause).await; +} + +#[tokio::test] +async fn blocking_post_tool_policy_services_cancel_after_result_commit() { + assert_blocking_post_hook_services_control(BlockingPostControl::Cancel).await; +} + +#[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 early_pre_and_post_hooks_share_the_exact_response_turn_identity() { + struct TurnRecordingPolicy { + events: tokio::sync::mpsc::UnboundedSender<(InterceptorPhase, String, u64)>, + } + + #[async_trait] + impl Interceptor for TurnRecordingPolicy { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_>, + ) -> InterceptorResult { + self.events + .send(( + InterceptorPhase::PreToolCall, + info.call.id.clone(), + info.invocation.turn_id.expect("pre hook turn").0, + )) + .unwrap(); + Ok(PreToolAction::Continue) + } + + async fn post_tool_call( + &self, + info: &ToolResultInfo<'_>, + ) -> InterceptorResult { + self.events + .send(( + InterceptorPhase::PostToolCall, + info.call.id.clone(), + info.invocation.turn_id.expect("post hook turn").0, + )) + .unwrap(); + Ok(PostToolAction::Continue) + } + } + + let (client, mut responses) = MultiControlledStreamClient::new(3); + let first = responses.remove(0); + let second = responses.remove(0); + let third = responses.remove(0); + let (event_tx, mut event_rx) = tokio::sync::mpsc::unbounded_channel(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(TurnRecordingPolicy { events: event_tx }); + engine.register_tool(SlowTool::new("turn_identity", 0).definition()); + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + + first + .send(Ok(Event::tool_use_start(0, "call_turn_0", "turn_identity"))) + .unwrap(); + first.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + first.send(Ok(Event::tool_use_stop(0))).unwrap(); + let pre_0 = event_rx.recv().await.unwrap(); + let post_0 = event_rx.recv().await.unwrap(); + assert_eq!( + pre_0, + (InterceptorPhase::PreToolCall, "call_turn_0".into(), 0) + ); + assert_eq!( + post_0, + (InterceptorPhase::PostToolCall, "call_turn_0".into(), 0) + ); + first + .send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(first); + + second + .send(Ok(Event::tool_use_start(0, "call_turn_1", "turn_identity"))) + .unwrap(); + second + .send(Ok(Event::tool_input_delta(0, r#"{}"#))) + .unwrap(); + second.send(Ok(Event::tool_use_stop(0))).unwrap(); + let pre_1 = event_rx.recv().await.unwrap(); + let post_1 = event_rx.recv().await.unwrap(); + assert_eq!( + pre_1, + (InterceptorPhase::PreToolCall, "call_turn_1".into(), 1) + ); + assert_eq!( + post_1, + (InterceptorPhase::PostToolCall, "call_turn_1".into(), 1) + ); + second + .send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(second); + + third.send(Ok(Event::text_delta(0, "done"))).unwrap(); + third + .send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(third); + let (output, _) = run.await.unwrap(); + assert!(matches!(output.result, EngineRunExit::Finished)); +} + +#[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); + 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 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 rejected_synthetic_result_commit_restores_as_non_runnable() { + struct SyntheticDenial; + + #[async_trait] + impl Interceptor for SyntheticDenial { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_>, + ) -> InterceptorResult { + Ok(PreToolAction::SyntheticResult(ToolResult::error( + &info.call.id, + "permission denied", + ))) + } + } + + let client = MockLlmClient::with_responses(vec![ + vec![ + Event::tool_use_start(0, "call_denied_restore", "denied_restore"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ], + vec![ + Event::text_delta(0, "restored safely"), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ], + ]) + .with_completion_support(ToolCallCompletionSupport::PerBlock); + let tool = SlowTool::new("denied_restore", 0); + let probe = tool.clone(); + let reject_once = Arc::new(AtomicUsize::new(1)); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(SyntheticDenial); + engine.register_tool(tool.definition()); + let reject_result = reject_once.clone(); + engine.on_history_append(move |item| { + if matches!(item, Item::ToolResult { call_id, .. } if call_id == "call_denied_restore") + && reject_result.swap(0, Ordering::SeqCst) == 1 + { + return Err("injected synthetic result commit failure".to_string()); + } + Ok(()) + }); + let mut history = History::new(); + + let output = engine.run(&mut history, "run denied tool").await; + assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + EngineError::HistoryAppend(ref message) + )) if message == "injected synthetic result commit failure" + )); + assert_eq!(probe.call_count(), 0); + assert!(history.items().any(|item| matches!( + item, + Item::ToolCall { + call_id, + execution_id: Some(_), + status: Some(agen::llm_client::ItemStatus::InProgress), + .. + } if call_id == "call_denied_restore" + ))); + assert!(!history.items().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_denied_restore" + ))); + + let mut engine = output.engine; + let resumed = engine.resume(&mut history).await; + assert!(matches!(resumed, EngineRunExit::Finished)); + assert_eq!( + probe.call_count(), + 0, + "denied side effect must never run on restore" + ); + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call_denied_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); + 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 clean_eof_without_provider_completion_after_early_start_stops_without_continuation() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let client_probe = client.clone(); + let tool = BarrierTool::new("barrier_clean_eof"); + 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_clean_eof", + "barrier_clean_eof", + ))) + .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 clean transport EOF"); + probe.release.notify_one(); + drop(tx); + + let (output, 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!( + client_probe.stream_count(), + 1, + "must not open a continuation request" + ); + assert_eq!(probe.starts.load(Ordering::SeqCst), 1); + assert_eq!( + history + .items() + .filter(|item| matches!(item, Item::ToolResult { call_id, .. } if call_id == "call_clean_eof")) + .count(), + 1 + ); +} + +#[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") + })); + + 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/agen/tests/response_completion_compatibility_test.rs b/crates/agen/tests/response_completion_compatibility_test.rs new file mode 100644 index 000000000..c7b7d8fac --- /dev/null +++ b/crates/agen/tests/response_completion_compatibility_test.rs @@ -0,0 +1,142 @@ +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use agen::llm_client::{ + ClientError, LlmClient, Request, ResponseStream, ToolCallCompletionSupport, + scheme::{Scheme, gemini::GeminiScheme, openai_chat::OpenAIScheme}, +}; +use agen::{Engine, EngineRunExit, History, ToolCallDispatchMode}; +use async_trait::async_trait; + +#[derive(Clone)] +struct SseFrame { + event_type: &'static str, + data: &'static str, +} + +#[derive(Clone)] +struct SchemeReplayClient { + scheme: S, + responses: Arc>>, + stream_calls: Arc, +} + +impl SchemeReplayClient { + fn new(scheme: S, responses: Vec>) -> Self { + Self { + scheme, + responses: Arc::new(responses), + stream_calls: Arc::new(AtomicUsize::new(0)), + } + } + + fn stream_count(&self) -> usize { + self.stream_calls.load(Ordering::SeqCst) + } +} + +#[async_trait] +impl LlmClient for SchemeReplayClient { + async fn stream(&self, _request: Request) -> Result { + let response_index = self.stream_calls.fetch_add(1, Ordering::SeqCst); + let frames = self + .responses + .get(response_index) + .ok_or_else(|| ClientError::Config("scheme replay responses exhausted".to_string()))?; + let mut state = S::State::default(); + let mut output = Vec::new(); + for frame in frames { + match self + .scheme + .parse_sse(frame.event_type, frame.data, &mut state) + { + Ok(events) => output.extend(events.into_iter().map(Ok)), + Err(error) => { + output.push(Err(error)); + break; + } + } + } + Ok(Box::pin(futures::stream::iter(output))) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + self.scheme.tool_call_completion_support() + } +} + +async fn assert_response_complete_scheme_accepts_clean_eof( + scheme: S, + frames: Vec, +) { + assert_eq!( + scheme.tool_call_completion_support(), + ToolCallCompletionSupport::ResponseComplete + ); + + for mode in [ + ToolCallDispatchMode::AfterResponse, + ToolCallDispatchMode::OnToolCallComplete, + ] { + let client = SchemeReplayClient::new(scheme.clone(), vec![frames.clone()]); + let probe = client.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(mode); + let mut history = History::new(); + + let output = engine.run(&mut history, "hello").await; + + assert!( + matches!(output.result, EngineRunExit::Finished), + "{mode:?} should accept the response-complete scheme's clean EOF, got {:?}", + output.result + ); + assert_eq!( + probe.stream_count(), + 1, + "{mode:?} must not open a continuation request after the adapter's normal EOF" + ); + } +} + +#[tokio::test] +async fn openai_chat_clean_eof_completes_in_default_and_opt_in_fallback_modes() { + assert_response_complete_scheme_accepts_clean_eof( + OpenAIScheme::new(), + vec![ + SseFrame { + event_type: "message", + data: r#"{"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":"stop"}]}"#, + }, + SseFrame { + event_type: "message", + data: "[DONE]", + }, + ], + ) + .await; +} + +#[tokio::test] +async fn gemini_clean_eof_completes_in_default_and_opt_in_fallback_modes() { + assert_response_complete_scheme_accepts_clean_eof( + GeminiScheme::new(), + vec![ + SseFrame { + event_type: "message", + data: r#"{"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"},"finishReason":"STOP","index":0}]}"#, + }, + SseFrame { + event_type: "message", + data: "[DONE]", + }, + ], + ) + .await; +} 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/memory/src/extract/input.rs b/crates/memory/src/extract/input.rs index 9cc2b0934..e2169df46 100644 --- a/crates/memory/src/extract/input.rs +++ b/crates/memory/src/extract/input.rs @@ -43,7 +43,7 @@ fn render_items(items: &[Item]) -> String { Item::ToolResult { summary, .. } => { lines.push(format!("[ToolResult] {summary}")); } - Item::Reasoning { .. } => {} + Item::AssistantResponseBoundary { .. } | Item::Reasoning { .. } => {} } } lines.join("\n\n") diff --git a/crates/session-store/src/logged_item.rs b/crates/session-store/src/logged_item.rs index f6547265b..546110baa 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}; @@ -49,10 +50,19 @@ pub enum LoggedItem { role: LoggedRole, content: Vec, }, + AssistantResponseBoundary { + response_id: String, + }, ToolCall { 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, @@ -115,15 +125,24 @@ impl From<&Item> for LoggedItem { role: (*role).into(), content: content.iter().map(LoggedContentPart::from).collect(), }, + Item::AssistantResponseBoundary { response_id } => Self::AssistantResponseBoundary { + response_id: response_id.clone(), + }, Item::ToolCall { 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, @@ -172,16 +191,24 @@ impl From for Item { content: content.into_iter().map(Into::into).collect(), status: None, }, + LoggedItem::AssistantResponseBoundary { response_id } => { + Item::assistant_response_boundary(response_id) + } LoggedItem::ToolCall { 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, @@ -325,9 +352,23 @@ mod tests { } } + #[test] + fn assistant_response_boundary_round_trips_for_provider_projection() { + let original = Item::assistant_response_boundary("response-7"); + let logged: LoggedItem = (&original).into(); + let json = serde_json::to_string(&logged).unwrap(); + assert!(json.contains(r#""kind":"assistant_response_boundary""#)); + assert_eq!( + Item::from(serde_json::from_str::(&json).unwrap()), + original + ); + } + #[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 +377,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..9fae33a15 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(), @@ -385,8 +386,8 @@ fn project_item(item: &LoggedItem) -> Option { }) .collect(), }), - // Hidden model reasoning is never observable. - LoggedItem::Reasoning { .. } => None, + // Internal response grouping and hidden model reasoning are never observable. + LoggedItem::AssistantResponseBoundary { .. } | LoggedItem::Reasoning { .. } => None, } } diff --git a/crates/worker/src/feature/builtin/flow_transition.rs b/crates/worker/src/feature/builtin/flow_transition.rs index 163ac8928..b96d1919d 100644 --- a/crates/worker/src/feature/builtin/flow_transition.rs +++ b/crates/worker/src/feature/builtin/flow_transition.rs @@ -636,7 +636,7 @@ mod tests { use agen::llm_client::error::ClientError; use agen::llm_client::event::{ BlockDelta, BlockMetadata, BlockStart, BlockStop, BlockType, DeltaContent, - Event as LlmEvent, StopReason, + Event as LlmEvent, ResponseStatus, StatusEvent, StopReason, }; use agen::llm_client::types::Request; use agen::tool::ToolExecutionContext; @@ -906,6 +906,9 @@ mod tests { reasoning: None, stop_reason: Some(StopReason::ToolUse), })), + Ok(LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + })), ]))) } else { Ok(Box::pin(stream::iter(vec![ @@ -924,6 +927,9 @@ mod tests { reasoning: None, stop_reason: Some(StopReason::EndTurn), })), + Ok(LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + })), ]))) } } 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/session_capture.rs b/crates/worker/src/session_capture.rs index 92f145185..01125c332 100644 --- a/crates/worker/src/session_capture.rs +++ b/crates/worker/src/session_capture.rs @@ -415,7 +415,7 @@ impl SessionCapture { search_text: text, }); } - Item::Reasoning { .. } => {} + Item::AssistantResponseBoundary { .. } | Item::Reasoning { .. } => {} } } @@ -713,6 +713,9 @@ fn render_item( ), } } + Item::AssistantResponseBoundary { .. } => { + format!("[{} Assistant response boundary omitted]", entry.id) + } Item::Reasoning { .. } => format!("[{} Reasoning omitted]", entry.id), }; truncate_chars(&text, max_bytes) 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/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..666eddec1 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -1264,7 +1264,10 @@ fn active_run_checkpoint_entry( fn is_ai_materialized_item(item: &Item) -> bool { match item { Item::Message { role, .. } => *role == Role::Assistant, - Item::ToolCall { .. } | Item::ToolResult { .. } | Item::Reasoning { .. } => true, + Item::AssistantResponseBoundary { .. } + | Item::ToolCall { .. } + | Item::ToolResult { .. } + | Item::Reasoning { .. } => true, } } @@ -1302,7 +1305,11 @@ where let subject = worker_subject(annotation_writer.state.location().session_id); let origin = if item.is_tool_result() { WorkerHistoryProvenance::ToolOutput { worker: subject } - } else if item.is_assistant_message() || item.is_tool_call() || item.is_reasoning() { + } else if item.is_assistant_message() + || matches!(item, Item::AssistantResponseBoundary { .. }) + || item.is_tool_call() + || item.is_reasoning() + { WorkerHistoryProvenance::ModelOutput { worker: subject } } else { // Unknown user/system append paths fail closed. Trusted system @@ -7138,6 +7145,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 @@ -7551,6 +7564,7 @@ fn write_overview_header(items: &[Item], out: &mut String) { for item in items { match item { Item::Message { .. } => messages += 1, + Item::AssistantResponseBoundary { .. } => {} Item::ToolCall { .. } => tool_calls += 1, Item::ToolResult { .. } => tool_results += 1, Item::Reasoning { .. } => reasoning += 1, @@ -8721,6 +8735,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 +8747,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/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 2bc9c314c..abd2b5ec4 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -17076,13 +17076,19 @@ fn workdir_removal_source_actor( fn contract_workdir_removal_source_actor( context: &server_api::ServerRequestContext, ) -> ApiResult { - workdir_removal_source_actor( - context - .worker_source - .as_ref() - .map(|source| (source.runtime_id.as_str(), source.worker_id.as_str())), - context.actor.as_ref(), - ) + let worker_source = context + .worker_source + .as_ref() + .map(|source| (source.runtime_id.as_str(), source.worker_id.as_str())) + .or_else(|| { + context.runtime_source.as_ref().and_then(|source| { + source + .worker_id + .as_deref() + .map(|worker_id| (source.runtime_id.as_str(), worker_id)) + }) + }); + workdir_removal_source_actor(worker_source, context.actor.as_ref()) } async fn cleanup_runtime_working_directory_contract( @@ -29527,6 +29533,36 @@ mod tests { .unwrap(); } + fn runtime_source_request( + identity: &worker_runtime::auth::RuntimeIdentityMaterial, + worker_id: Option<&str>, + method: &str, + path: &str, + body: Vec, + ) -> Request { + let proof = worker_runtime::auth::RuntimeRequestSourceSigner::from_identity(identity) + .issue( + "server-test", + TEST_WORKSPACE_ID, + worker_id, + worker_runtime::auth::WORKSPACE_REQUEST_PERMISSION, + method, + path, + &body, + i64::try_from(worker_runtime::auth::unix_now_seconds()).unwrap_or(i64::MAX), + 30, + ) + .unwrap(); + let mut request = Request::builder().method(method).uri(path).header( + worker_runtime::auth::RUNTIME_REQUEST_SOURCE_PROOF_HEADER, + proof, + ); + if !body.is_empty() { + request = request.header(CONTENT_TYPE, "application/json"); + } + request.body(Body::from(body)).unwrap() + } + fn runtime_resource_fetch_request( api: &WorkspaceApi, identity: &worker_runtime::auth::RuntimeIdentityMaterial, @@ -38967,6 +39003,390 @@ mod tests { ); } + #[test] + fn workdir_removal_source_actor_preserves_verified_worker_and_account_sources() { + let actor = RequestActor { + user_id: "account-user".to_string(), + account_id: "account-test".to_string(), + handle: "account".to_string(), + display_name: "Account".to_string(), + auth_method: ActorAuthMethod::ApiToken, + }; + let context = |worker_source, runtime_source, actor| server_api::ServerRequestContext { + actor, + worker_source, + runtime_source, + origin: None, + transport_headers: Vec::new(), + }; + + assert_eq!( + contract_workdir_removal_source_actor(&context( + Some(server_api::ServerWorkerSource { + runtime_id: "legacy-runtime".to_string(), + worker_id: "legacy-worker".to_string(), + }), + Some(server_api::ServerRuntimeSource { + runtime_id: "runtime-source".to_string(), + worker_id: Some("runtime-worker".to_string()), + }), + Some(actor.clone()), + )) + .unwrap(), + "worker:legacy-runtime:legacy-worker" + ); + assert_eq!( + contract_workdir_removal_source_actor(&context( + None, + Some(server_api::ServerRuntimeSource { + runtime_id: "runtime-source".to_string(), + worker_id: Some("runtime-worker".to_string()), + }), + Some(actor.clone()), + )) + .unwrap(), + "worker:runtime-source:runtime-worker" + ); + assert_eq!( + contract_workdir_removal_source_actor(&context( + None, + Some(server_api::ServerRuntimeSource { + runtime_id: "runtime-only".to_string(), + worker_id: None, + }), + Some(actor.clone()), + )) + .unwrap(), + "account:account-test" + ); + let error = contract_workdir_removal_source_actor(&context( + None, + Some(server_api::ServerRuntimeSource { + runtime_id: "runtime-only".to_string(), + worker_id: None, + }), + None, + )) + .unwrap_err(); + assert_eq!(error.into_response().status(), StatusCode::FORBIDDEN); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn signed_worker_workdir_removal_reaches_both_generated_routes_and_records_actor() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let mut api = test_api(workspace.path()).await; + let identity = + worker_runtime::auth::RuntimeIdentityMaterial::generate("runtime-source").unwrap(); + configure_runtime_request_auth(&mut api, &identity, "runtime-source"); + seed_worker_source_member(&api, "runtime-source", "worker-source"); + let provider_runtime_id = "runtime-provider"; + let fixtures = [ + ( + "workspace-route-workdir", + format!("/api/w/{TEST_WORKSPACE_ID}/working-directories/workspace-route-workdir"), + "remove through workspace route", + ), + ( + "runtime-route-workdir", + format!( + "/api/w/{TEST_WORKSPACE_ID}/runtimes/{provider_runtime_id}/working-directories/runtime-route-workdir" + ), + "remove through Runtime route", + ), + ]; + let runtime = worker_runtime::Runtime::with_execution_backend( + worker_runtime::RuntimeOptions::default(), + Arc::new(DeterministicExecutionBackend::default()), + ) + .unwrap(); + let repository_source = server_api::RepositorySource { + kind: server_api::RepositorySourceKind::LocalPath, + uri: workspace.path().display().to_string(), + }; + let source_fingerprint = repository_source_fingerprint(&repository_source); + for (workdir_id, _, _) in &fixtures { + runtime + .create_working_directory(WorkingDirectoryRequest { + repository: WorkingDirectoryRepository { + id: "repo-test".to_string(), + provider: "git".to_string(), + source: repository_source.clone(), + source_revision: 1, + source_fingerprint: source_fingerprint.clone(), + selector: Some(RuntimeRepositorySelector("HEAD".to_string())), + }, + display_name: None, + materializer: Default::default(), + backend_workdir_id: Some((*workdir_id).to_string()), + materialization: None, + }) + .unwrap(); + seed_cleanup_workdir_for_runtime( + &api, + workdir_id, + provider_runtime_id, + "present", + "clean", + ); + } + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let provider_base_url = format!("http://{}", listener.local_addr().unwrap()); + let provider_server = tokio::spawn(serve_runtime_http_with_injected_test_auth( + runtime, listener, + )); + let now = now_registry_timestamp(); + let provider_identity = RuntimeIdentityMaterial::generate(provider_runtime_id).unwrap(); + let provider_binding = WorkspaceRuntimeBinding { + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: provider_runtime_id.to_string(), + display_name: "Runtime Provider".to_string(), + base_url: provider_base_url.clone(), + public_key: provider_identity.public_key, + public_key_fingerprint: String::new(), + binding_revision: 1, + state: StoredRuntimeBindingState::Verified, + authentication_mode: StoredRuntimeAuthenticationMode::LegacyServerIssuer, + workspace_key_id: None, + workspace_key_generation: None, + created_at: now.clone(), + updated_at: now, + revoked_at: None, + }; + api.store + .upsert_workspace_runtime_binding_record(provider_binding.clone(), false) + .await + .unwrap(); + api.runtime.register_or_replace( + RemoteWorkerRuntime::new( + RemoteRuntimeConfig { + runtime_id: provider_runtime_id.to_string(), + workspace_id: Some(TEST_WORKSPACE_ID.to_string()), + display_name: "Runtime Provider".to_string(), + base_url: provider_base_url, + bearer_token: Some(TEST_RUNTIME_HTTP_TOKEN.to_string()), + workspace_authorization: None, + strict_public_egress: false, + cached_worker_creation_available: true, + cached_os: "linux".to_string(), + cached_arch: "x86_64".to_string(), + cached_status: "active".to_string(), + timeout: std::time::Duration::from_secs(2), + }, + TEST_WORKSPACE_ID.to_string(), + "http://127.0.0.1:1".to_string(), + ) + .unwrap(), + ); + let observed = api + .runtime + .working_directory(provider_runtime_id, fixtures[0].0) + .unwrap(); + assert_eq!( + observed.state, + InternalWorkerOperationState::Accepted, + "{observed:?}" + ); + let app = build_router(api.clone()); + + for (workdir_id, path, reason) in fixtures { + let body = serde_json::to_vec(&json!({ "reason": reason })).unwrap(); + let response = app + .clone() + .oneshot(runtime_source_request( + &identity, + Some("worker-source"), + "DELETE", + &path, + body, + )) + .await + .unwrap(); + let status = response.status(); + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + assert_eq!( + status, + StatusCode::OK, + "{path}: {}", + String::from_utf8_lossy(&bytes) + ); + let response: WorkingDirectoryRemovalResponse = serde_json::from_slice(&bytes).unwrap(); + assert_eq!( + response.disposition, + WorkingDirectoryRemovalDisposition::Removed, + "{response:?}" + ); + assert!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, workdir_id) + .unwrap() + .is_none() + ); + let operation = api + .config_store + .find_workdir_removal_operation_by_intent( + TEST_WORKSPACE_ID, + workdir_id, + "worker:runtime-source:worker-source", + reason, + ) + .unwrap() + .unwrap(); + assert_eq!( + operation.source_actor, + "worker:runtime-source:worker-source" + ); + assert_eq!(operation.state, WorkdirRemovalOperationState::Completed); + assert_eq!( + operation.disposition, + Some(WorkdirRemovalDisposition::Removed) + ); + } + provider_server.abort(); + } + + #[tokio::test] + async fn workdir_removal_rejects_unverified_or_runtime_only_sources_before_removal() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let mut api = test_api(workspace.path()).await; + let identity = + worker_runtime::auth::RuntimeIdentityMaterial::generate("runtime-source").unwrap(); + configure_runtime_request_auth(&mut api, &identity, "runtime-source"); + let fixtures = [ + ("runtime-only-workdir", "runtime only"), + ("invalid-proof-workdir", "invalid proof"), + ("unauthenticated-workdir", "unauthenticated"), + ]; + for (workdir_id, _) in fixtures { + seed_cleanup_workdir_for_runtime( + &api, + workdir_id, + EMBEDDED_WORKER_RUNTIME_ID, + "present", + "clean", + ); + } + let app = build_router(api.clone()); + + let runtime_only_path = + format!("/api/w/{TEST_WORKSPACE_ID}/working-directories/runtime-only-workdir"); + let runtime_only = app + .clone() + .oneshot(runtime_source_request( + &identity, + None, + "DELETE", + &runtime_only_path, + serde_json::to_vec(&json!({ "reason": "runtime only" })).unwrap(), + )) + .await + .unwrap(); + assert_eq!(runtime_only.status(), StatusCode::FORBIDDEN); + + let invalid_path = + format!("/api/w/{TEST_WORKSPACE_ID}/working-directories/invalid-proof-workdir"); + let invalid = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri(&invalid_path) + .header(CONTENT_TYPE, "application/json") + .header( + worker_runtime::auth::RUNTIME_REQUEST_SOURCE_PROOF_HEADER, + "not-a-valid-proof", + ) + .body(Body::from(json!({ "reason": "invalid proof" }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(invalid.status(), StatusCode::UNAUTHORIZED); + + let unauthenticated_path = + format!("/api/w/{TEST_WORKSPACE_ID}/working-directories/unauthenticated-workdir"); + let unauthenticated = app + .oneshot( + Request::builder() + .method("DELETE") + .uri(&unauthenticated_path) + .header(CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ "reason": "unauthenticated" }).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(unauthenticated.status(), StatusCode::UNAUTHORIZED); + + for (workdir_id, _) in fixtures { + assert!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, workdir_id) + .unwrap() + .is_some() + ); + } + } + + #[tokio::test] + async fn account_workdir_removal_still_records_account_actor_and_preserves_guards() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + let token = seed_test_api_token(api.store.as_ref(), "workdir-removal"); + let workdir_id = "account-occupied-workdir"; + seed_cleanup_workdir_for_runtime( + &api, + workdir_id, + EMBEDDED_WORKER_RUNTIME_ID, + "present", + "clean", + ); + let occupant = seed_cleanup_worker(&api, 27, "pinned"); + seed_cleanup_link(&api, &occupant, workdir_id); + let reason = "retain occupied Workdir"; + let path = format!("/api/w/{TEST_WORKSPACE_ID}/working-directories/{workdir_id}"); + + let response = request_json_authenticated( + build_router(api.clone()), + "DELETE", + &path, + Some(json!({ "reason": reason })), + &token, + StatusCode::OK, + ) + .await; + assert_eq!(response["disposition"], "retained"); + assert!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, workdir_id) + .unwrap() + .is_some() + ); + let operation = api + .config_store + .find_workdir_removal_operation_by_intent( + TEST_WORKSPACE_ID, + workdir_id, + "account:account-workdir-removal", + reason, + ) + .unwrap() + .unwrap(); + assert_eq!(operation.source_actor, "account:account-workdir-removal"); + assert_eq!( + operation.disposition, + Some(WorkdirRemovalDisposition::Retained) + ); + assert_eq!( + operation.failure_category.as_deref(), + Some("blocked_by_live_authority") + ); + } + #[test] fn workdir_session_owner_is_only_sent_for_same_runtime_worker() { let embedded_worker = RuntimeWorkerRef::new("embedded-worker-runtime", "5"); diff --git a/docs/manifest.toml b/docs/manifest.toml index 92b671e6e..c1bcccc9f 100644 --- a/docs/manifest.toml +++ b/docs/manifest.toml @@ -121,6 +121,15 @@ ref = "anthropic/claude-sonnet-4-6" # reasoning = "medium" # reasoning = -1 +# 任意。デフォルト: "after_response"。 +# "on_tool_call_complete" は provider が個別 ToolCall の完成を保証する場合だけ、 +# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider、または +# 応答全体を見て実行前 Pause を判断する turn-end hook がある場合は、応答完了まで +# 警告付きで明示的に遅延する。これは直列実行モードではなく、先行 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]