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

fix: close early dispatch recovery gaps

Browse Source
This commit is contained in:
Hare
2026-09-29 12:37:13 +09:00
parent dcf14ef565
commit 1980da2880
14 changed files with 868 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/src/engine.rs
+271 -39
View File
@@ -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<Item>,
original_to_projected_index: Vec<usize>,
}
fn materialize_provider_history(items: &[Item]) -> ProviderHistoryProjection {
let mut materialized: Vec<_> = items.iter().cloned().enumerate().collect();
let mut call_order = HashMap::<String, usize>::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<String, usize>) {
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<String, usize>,
) {
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 {
normalize_tool_run(&mut materialized[start..index], call_order);
}
}
fn normalize_response_history(response: &mut [(usize, Item)], call_order: &HashMap<String, usize>) {
// 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::<Vec<_>>();
calls.sort_by_key(|item| match item {
Item::ToolCall {
call_id,
call_index,
..
} => (
0usize,
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");
}
}
// 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),
),
_ => unreachable!("tool run contains only calls and results"),
_ => (0usize, 0usize),
});
start = index.saturating_add(1);
}
}
}
fn materialize_provider_history(items: &[Item]) -> ProviderHistoryProjection {
let mut call_order = HashMap::<String, usize>::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 mut original_to_projected_index = vec![0; materialized.len()];
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::<Vec<_>>();
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::<Vec<_>>();
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<String>,
call_info_map: HashMap<String, ToolCallExecutionInfo>,
@@ -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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
)
.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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
// 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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
annotate,
result,
None,
turn_id,
true,
&call_info_map,
&mut attempt_fence,
&mut terminal_call_ids,
@@ -1982,6 +2136,8 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
annotate,
result,
Some(&attempt_id),
turn_id,
true,
&call_info_map,
&mut attempt_fence,
&mut terminal_call_ids,
@@ -2061,6 +2217,8 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
annotate,
result,
Some(&attempt_id),
turn_id,
false,
&call_info_map,
&mut attempt_fence,
&mut terminal_call_ids,
@@ -2087,6 +2245,8 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
);
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
"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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
);
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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::<Vec<_>>();
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"[..]);
crates/agen/src/llm_client/scheme/anthropic/request.rs
+1
View File
@@ -263,6 +263,7 @@ impl AnthropicScheme {
for (i, item) in items.iter().enumerate() {
match item {
Item::AssistantResponseBoundary { .. } => {}
Item::Message { role, content, .. } => {
let parts: Vec<AnthropicContentPart> = content
.iter()
crates/agen/src/llm_client/scheme/gemini/request.rs
+1
View File
@@ -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(
crates/agen/src/llm_client/scheme/openai_chat/request.rs
+1
View File
@@ -221,6 +221,7 @@ impl OpenAIScheme {
);
}
match item {
Item::AssistantResponseBoundary { .. } => {}
Item::Message { role, content, .. } => {
// Flush pending tool calls
self.flush_pending_assistant(
crates/agen/src/llm_client/scheme/openai_responses/request.rs
+1
View File
@@ -254,6 +254,7 @@ fn convert_items_to_input(items: &[Item], supports_images: bool) -> Vec<InputIte
let mut out = Vec::with_capacity(items.len());
for item in items {
match item {
Item::AssistantResponseBoundary { .. } => {}
Item::Message { role, content, .. } => {
let (role_str, text_variant): (&'static str, fn(String) -> InputContent) =
match role {
crates/agen/src/llm_client/types.rs
+22 -3
View File
@@ -93,6 +93,12 @@ pub enum Item {
status: Option<ItemStatus>,
},
/// 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<usize>,
/// 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<String>,
/// 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<ItemStatus>,
},
@@ -226,6 +233,14 @@ impl Item {
}
}
/// Create an internal durable assistant-response boundary.
#[doc(hidden)]
pub fn assistant_response_boundary(response_id: impl Into<String>) -> 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<String>) -> 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",
crates/agen/src/timeline/timeline.rs
+73 -10
View File
@@ -374,7 +374,9 @@ pub struct Timeline {
thinking_block_handlers: Vec<Box<dyn ErasedBlockHandler>>,
tool_use_block_handlers: Vec<Box<dyn ErasedBlockHandler>>,
// 現在アクティブなブロック
// 現在アクティブなブロック。provider は index ごとに複数 scope を
// interleave できるため、単一の current block ではなく全 scope を追跡する。
active_blocks: HashMap<(BlockType, usize), ()>,
current_block: Option<BlockType>,
// 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() {
// 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::<Vec<_>>();
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<Mutex<Vec<bool>>>,
}
impl Handler<ToolUseBlockKind> 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() {
// シンプルなテスト用構造体
crates/agen/tests/annotated_history_test.rs
+1
View File
@@ -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",
crates/agen/tests/parallel_execution_test.rs
+447
View File
@@ -308,6 +308,68 @@ impl LlmClient for ControlledStreamClient {
}
}
#[derive(Clone)]
struct MultiControlledStreamClient {
receivers: Arc<
Mutex<
std::collections::VecDeque<
tokio::sync::mpsc::UnboundedReceiver<Result<Event, ClientError>>,
>,
>,
>,
}
impl MultiControlledStreamClient {
fn new(
response_count: usize,
) -> (
Self,
Vec<tokio::sync::mpsc::UnboundedSender<Result<Event, ClientError>>>,
) {
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<ResponseStream, ClientError> {
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<dyn LlmClient> {
Box::new(self.clone())
}
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
ToolCallCompletionSupport::PerBlock
}
}
#[derive(Clone)]
struct BoundedStreamClient {
receiver: Arc<Mutex<Option<tokio::sync::mpsc::Receiver<Result<Event, ClientError>>>>>,
@@ -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<tokio::sync::Notify>,
release: Arc<tokio::sync::Notify>,
}
#[async_trait]
impl Interceptor for BlockingEarlyPolicy {
fn supports_early_tool_dispatch(&self) -> bool {
true
}
async fn pre_tool_call(
&self,
_info: &mut ToolCallInfo<'_>,
) -> InterceptorResult<PreToolAction> {
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<tokio::sync::Notify>,
release: Arc<tokio::sync::Notify>,
}
#[async_trait]
impl Interceptor for BlockingPostPolicy {
fn supports_early_tool_dispatch(&self) -> bool {
true
}
async fn post_tool_call(
&self,
_info: &ToolResultInfo<'_>,
) -> InterceptorResult<PostToolAction> {
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<PreToolAction> {
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<PostToolAction> {
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<PreToolAction> {
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;
crates/memory/src/extract/input.rs
+1 -1
View File
@@ -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")
crates/session-store/src/logged_item.rs
+21
View File
@@ -50,6 +50,9 @@ pub enum LoggedItem {
role: LoggedRole,
content: Vec<LoggedContentPart>,
},
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<LoggedItem> 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::<LoggedItem>(&json).unwrap()),
original
);
}
#[test]
fn round_trip_tool_call() {
let original = Item::tool_call("call_42", "get_weather", r#"{"city":"Tokyo"}"#)
crates/session-store/src/public_snapshot.rs
+2 -2
View File
@@ -386,8 +386,8 @@ fn project_item(item: &LoggedItem) -> Option<SessionSnapshotEntryData> {
})
.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,
}
}
crates/worker/src/session_capture.rs
+4 -1
View File
@@ -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)
crates/worker/src/worker.rs
+10 -2
View File
@@ -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,