diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index fa4ab25a6..8dd5825a3 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -98,31 +98,39 @@ pub struct EngineConfig { /// 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; - - for (_, item) in &materialized { - if let Item::ToolCall { +fn normalize_tool_run(run: &mut [(usize, Item)], call_order: &HashMap) { + run.sort_by_key(|(_, item)| match item { + Item::ToolCall { call_id, call_index, .. - } = item - { - call_order.insert(call_id.clone(), call_index.unwrap_or(next_call_order)); - next_call_order += 1; - } - } + } => ( + 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() { if !matches!( @@ -141,29 +149,128 @@ fn materialize_provider_history(items: &[Item]) -> ProviderHistoryProjection { { index += 1; } - materialized[start..index].sort_by_key(|(_, item)| match item { - Item::ToolCall { - call_id, - call_index, - .. - } => ( - 0usize, - call_index - .or_else(|| call_order.get(call_id).copied()) - .unwrap_or(usize::MAX), - ), - Item::ToolResult { call_id, .. } => ( - 1usize, - call_order.get(call_id).copied().unwrap_or(usize::MAX), - ), - _ => unreachable!("tool run contains only calls and results"), - }); + 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(), @@ -253,6 +360,7 @@ type ToolCallExecutionInfo = ( /// provider-confirmed blocks complete. struct EarlyToolExecutionBatch { batch_id: String, + turn_id: usize, calls: Vec<(usize, ToolCall)>, seen_call_ids: HashSet, call_info_map: HashMap, @@ -266,9 +374,10 @@ struct EarlyToolExecutionBatch { } impl EarlyToolExecutionBatch { - fn new(batch_id: String) -> Self { + 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(), @@ -1362,7 +1471,13 @@ impl Engine { let context = ToolExecutionContext::new(&tool_call.id, &batch.batch_id, call_index); if !tool_call.input.is_object() { - self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + 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", @@ -1372,7 +1487,13 @@ impl Engine { } let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) else { - self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + 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), @@ -1502,12 +1623,24 @@ impl Engine { ) .into()); } - self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + 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, None)?; + 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", @@ -1515,7 +1648,13 @@ impl Engine { batch.dispatched_any = true; } PreToolAction::Abort(reason) => { - self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + 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}"), @@ -1544,6 +1683,8 @@ impl Engine { annotate, result, None, + batch.turn_id, + true, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1569,6 +1710,8 @@ impl Engine { annotate, result, None, + batch.turn_id, + false, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1615,6 +1758,8 @@ impl Engine { annotate, result, Some(&attempt_id), + batch.turn_id, + false, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1636,6 +1781,8 @@ impl Engine { 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, @@ -1665,6 +1812,8 @@ impl Engine { annotate, result, None, + batch.turn_id, + true, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1701,6 +1850,8 @@ impl Engine { annotate, result, Some(&attempt_id), + batch.turn_id, + true, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1768,6 +1919,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); @@ -1952,6 +2104,8 @@ impl Engine { annotate, result, None, + turn_id, + true, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -1982,6 +2136,8 @@ impl Engine { annotate, result, Some(&attempt_id), + turn_id, + true, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -2061,6 +2217,8 @@ impl Engine { annotate, result, Some(&attempt_id), + turn_id, + false, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -2087,6 +2245,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, @@ -2135,6 +2295,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, ( @@ -2223,10 +2385,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, ); @@ -2239,16 +2401,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)); } @@ -2284,6 +2458,8 @@ impl Engine { annotate, ToolResult::outcome_unknown(&call_id), None, + self.turn_count, + false, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -2532,6 +2708,14 @@ impl Engine { "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, @@ -2673,7 +2857,7 @@ impl Engine { 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)); + early_tools = Some(EarlyToolExecutionBatch::new(batch_id, current_turn)); } let admission = self .admit_early_tool_calls( @@ -3057,7 +3241,7 @@ impl Engine { ); self.tool_execution_batch_count += 1; early_tools = - Some(EarlyToolExecutionBatch::new(batch_id)); + Some(EarlyToolExecutionBatch::new(batch_id, turn)); } let fresh_calls = { let batch = early_tools.as_ref().expect("batch initialized"); @@ -3211,6 +3395,8 @@ impl Engine { annotate, result, Some(&attempt_id), + batch.turn_id, + true, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -3941,6 +4127,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", @@ -4025,6 +4212,51 @@ mod tests { 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/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/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/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/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/types.rs b/crates/agen/src/llm_client/types.rs index 50ed7506a..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 @@ -107,11 +113,12 @@ pub enum Item { /// Provider/model order among tool calls in this assistant response. #[serde(default, skip_serializing_if = "Option::is_none")] call_index: Option, - /// Durable identity of a tool execution attempt that has started. + /// 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 provider - /// execution ownership was transferred and must never be inferred pending. + /// 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, }, @@ -226,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 // ======================================================================== @@ -405,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()), @@ -416,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), @@ -431,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(), @@ -441,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/timeline.rs b/crates/agen/src/timeline/timeline.rs index fcdac1fc3..6dfe80b2c 100644 --- a/crates/agen/src/timeline/timeline.rs +++ b/crates/agen/src/timeline/timeline.rs @@ -374,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 の集約バッファ。 @@ -400,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, } @@ -554,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); @@ -567,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 { @@ -588,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) { @@ -597,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( @@ -617,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; } } @@ -674,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/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/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index f6d445fb0..744e4111a 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -308,6 +308,68 @@ impl LlmClient for ControlledStreamClient { } } +#[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>>>>, @@ -576,6 +638,82 @@ async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_pause() { assert_eq!(probe.call_count(), 0); } +#[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, _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); +} + #[tokio::test] async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_result() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); @@ -680,6 +818,106 @@ async fn early_terminal_result_commits_while_response_stream_is_open() { 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); @@ -720,6 +958,119 @@ async fn reverse_stop_order_preserves_model_call_indexes() { 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![ @@ -885,6 +1236,102 @@ async fn persisted_started_call_is_closed_unknown_and_never_resumed() { ))); } +#[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; 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 6855b0973..546110baa 100644 --- a/crates/session-store/src/logged_item.rs +++ b/crates/session-store/src/logged_item.rs @@ -50,6 +50,9 @@ pub enum LoggedItem { role: LoggedRole, content: Vec, }, + AssistantResponseBoundary { + response_id: String, + }, ToolCall { call_id: String, name: String, @@ -122,6 +125,9 @@ 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, @@ -185,6 +191,9 @@ 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, @@ -343,6 +352,18 @@ 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"}"#) diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index 72ec0a140..9fae33a15 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -386,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/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/worker.rs b/crates/worker/src/worker.rs index 6ab90b0d0..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 @@ -7557,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,