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

fix: harden early tool call dispatch

Browse Source
This commit is contained in:
Hare
2026-09-29 12:00:36 +09:00
parent cc2f8b3422
commit dcf14ef565
13 changed files with 1481 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/src/engine.rs
+423 -114
View File
@@ -24,7 +24,7 @@ use crate::{
ToolCallInfo, ToolResultInfo, TurnEndAction,
},
llm_client::{
ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream,
ClientError, ConfigWarning, ItemStatus, LlmClient, Request, RequestConfig, ResponseStream,
ToolDefinition, error::is_retryable, event::Event, retry::RetryPolicy,
transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT,
},
@@ -110,33 +110,56 @@ fn materialize_provider_history(items: &[Item]) -> ProviderHistoryProjection {
let mut materialized: Vec<_> = items.iter().cloned().enumerate().collect();
let mut call_order = HashMap::<String, usize>::new();
let mut next_call_order = 0usize;
let mut index = 0usize;
while index < materialized.len() {
match &materialized[index].1 {
Item::ToolCall { call_id, .. } => {
call_order.insert(call_id.clone(), next_call_order);
next_call_order += 1;
index += 1;
}
Item::ToolResult { .. } => {
let start = index;
while index < materialized.len()
&& matches!(materialized[index].1, Item::ToolResult { .. })
{
index += 1;
}
materialized[start..index].sort_by_key(|(_, item)| match item {
Item::ToolResult { call_id, .. } => {
call_order.get(call_id).copied().unwrap_or(usize::MAX)
}
_ => unreachable!("tool-result run contains only ToolResult items"),
});
}
_ => index += 1,
for (_, item) in &materialized {
if let Item::ToolCall {
call_id,
call_index,
..
} = item
{
call_order.insert(call_id.clone(), call_index.unwrap_or(next_call_order));
next_call_order += 1;
}
}
let mut index = 0usize;
while index < materialized.len() {
if !matches!(
materialized[index].1,
Item::ToolCall { .. } | Item::ToolResult { .. }
) {
index += 1;
continue;
}
let start = index;
while index < materialized.len()
&& matches!(
materialized[index].1,
Item::ToolCall { .. } | Item::ToolResult { .. }
)
{
index += 1;
}
materialized[start..index].sort_by_key(|(_, item)| match item {
Item::ToolCall {
call_id,
call_index,
..
} => (
0usize,
call_index
.or_else(|| call_order.get(call_id).copied())
.unwrap_or(usize::MAX),
),
Item::ToolResult { call_id, .. } => (
1usize,
call_order.get(call_id).copied().unwrap_or(usize::MAX),
),
_ => unreachable!("tool run contains only calls and results"),
});
}
let mut original_to_projected_index = vec![0; materialized.len()];
for (projected_index, (original_index, _)) in materialized.iter().enumerate() {
original_to_projected_index[*original_index] = projected_index;
@@ -230,8 +253,7 @@ type ToolCallExecutionInfo = (
/// provider-confirmed blocks complete.
struct EarlyToolExecutionBatch {
batch_id: String,
next_call_index: usize,
calls: Vec<ToolCall>,
calls: Vec<(usize, ToolCall)>,
seen_call_ids: HashSet<String>,
call_info_map: HashMap<String, ToolCallExecutionInfo>,
started_calls: Vec<(String, String)>,
@@ -247,7 +269,6 @@ impl EarlyToolExecutionBatch {
fn new(batch_id: String) -> Self {
Self {
batch_id,
next_call_index: 0,
calls: Vec::new(),
seen_call_ids: HashSet::new(),
call_info_map: HashMap::new(),
@@ -378,6 +399,37 @@ struct StreamResponseOutput {
early_tools: Option<EarlyToolExecutionBatch>,
}
const RESPONSE_STREAM_PUMP_CAPACITY: usize = 64;
struct ResponseStreamPump {
receiver: mpsc::Receiver<Result<Event, ClientError>>,
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<Result<Event, ClientError>> {
self.receiver.recv().await
}
}
impl Drop for ResponseStreamPump {
fn drop(&mut self) {
self.task.abort();
}
}
pub struct Engine<C: LlmClient, S: EngineState = Mutable, A: Send + Sync = ()> {
/// LLM client
client: C,
@@ -1229,67 +1281,88 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
exit
}
/// Check for pending tool calls (for resuming from Pause)
fn get_pending_tool_calls(&self, history: &History<A>) -> Option<Vec<ToolCall>> {
// 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<A>) -> (Vec<ToolCall>, Vec<String>) {
let mut pending_calls = Vec::new();
let mut started_call_ids = Vec::new();
let mut answered_call_ids = std::collections::HashSet::new();
// First pass: collect all answered call IDs
for item in history.items() {
if let Item::ToolResult { call_id, .. } = item {
answered_call_ids.insert(call_id.clone());
}
}
// Second pass: find unanswered tool calls
for item in history.items() {
if let Item::ToolCall {
call_id,
name,
arguments,
execution_id,
status,
..
} = item
&& !answered_call_ids.contains(call_id)
{
if !answered_call_ids.contains(call_id) {
let input = serde_json::from_str::<Value>(arguments)
.ok()
.filter(Value::is_object)
.unwrap_or(Value::Null);
pending_calls.push(ToolCall {
id: call_id.clone(),
name: name.clone(),
input,
});
if execution_id.is_some() || *status == Some(ItemStatus::InProgress) {
started_call_ids.push(call_id.clone());
continue;
}
let input = serde_json::from_str::<Value>(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<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
call: &ToolCall,
call_index: usize,
execution_id: Option<String>,
) -> Result<(), EngineError> {
let status = if execution_id.is_some() {
ItemStatus::InProgress
} else {
Some(pending_calls)
}
ItemStatus::Completed
};
let item = Item::tool_call_json(&call.id, &call.name, call.input.clone())
.with_tool_execution_metadata(call_index, execution_id)
.with_status(status);
self.append_history_items(history, std::iter::once(item), annotate)
}
async fn admit_early_tool_calls(
&mut self,
history: &History<A>,
history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
batch: &mut EarlyToolExecutionBatch,
tool_calls: Vec<ToolCall>,
tool_calls: Vec<crate::timeline::CollectedToolCall>,
) -> Result<EarlyToolAdmission, EngineError> {
for mut tool_call in tool_calls {
for collected in tool_calls {
let call_index = collected.call_index;
let mut tool_call = collected.call;
if !batch.seen_call_ids.insert(tool_call.id.clone()) {
warn!(call_id = %tool_call.id, "Ignoring duplicate completed tool-call event");
continue;
}
let call_index = batch.next_call_index;
batch.next_call_index += 1;
batch.calls.push(tool_call.clone());
batch.calls.push((call_index, tool_call.clone()));
let expected_tool_use_id = tool_call.id.clone();
let context = ToolExecutionContext::new(&tool_call.id, &batch.batch_id, call_index);
if !tool_call.input.is_object() {
self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?;
batch.synthetic_results.push_back(ToolResult::error(
&tool_call.id,
"Tool arguments must be one complete JSON object",
@@ -1299,21 +1372,11 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}
let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) else {
let attempt_id = context.batch_id.clone();
batch
.attempt_fence
.register(tool_call.id.clone(), attempt_id.clone());
batch
.started_calls
.push((tool_call.id.clone(), attempt_id.clone()));
let call_id = tool_call.id.clone();
let name = tool_call.name.clone();
batch.futures.push(Box::pin(async move {
(
attempt_id,
ToolResult::error(&call_id, format!("Tool not found: {name}")),
)
}));
self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?;
batch.synthetic_results.push_back(ToolResult::error(
&tool_call.id,
format!("Tool not found: {}", tool_call.name),
));
batch.dispatched_any = true;
continue;
};
@@ -1332,16 +1395,28 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
tool,
context,
};
let action = self
.interceptor
.pre_tool_call(&mut info)
.await
.map_err(|error| {
EngineError::from(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
error,
))
})?;
let action = tokio::select! {
result = self.interceptor.pre_tool_call(&mut info) => {
result.map_err(|error| {
EngineError::from(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
error,
))
})?
}
pause = self.pause_rx.recv() => {
if pause.is_some() {
info!(call_id = %expected_tool_use_id, "Paused while awaiting pre-tool policy");
}
return Ok(EarlyToolAdmission::Pause);
}
cancel = self.cancel_rx.recv() => {
if cancel.is_some() {
info!(call_id = %expected_tool_use_id, "Cancelled while awaiting pre-tool policy");
}
return Err(EngineError::Cancelled);
}
};
if info.call.id != expected_tool_use_id {
return Err(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
@@ -1352,22 +1427,30 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
)
.into());
}
tool_call = info.call;
let mut context = info.context;
let ToolCallInfo {
call,
meta,
tool,
mut context,
..
} = info;
tool_call = call;
context.call_id = tool_call.id.clone();
batch.call_info_map.insert(
tool_call.id.clone(),
(
tool_call.clone(),
info.meta.clone(),
info.tool.clone(),
context.clone(),
),
(tool_call.clone(), meta, tool.clone(), context.clone()),
);
match action {
PreToolAction::Continue => {
let attempt_id = context.batch_id.clone();
let attempt_id = context.execution_id();
self.append_early_tool_call(
history,
annotate,
&tool_call,
call_index,
Some(attempt_id.clone()),
)?;
batch
.attempt_fence
.register(tool_call.id.clone(), attempt_id.clone());
@@ -1376,8 +1459,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
.push((tool_call.id.clone(), attempt_id.clone()));
let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default();
let call_id = tool_call.id.clone();
let (handle, terminal) =
ToolExecutionHandle::start(info.tool, input_json, context);
let (handle, terminal) = ToolExecutionHandle::start(tool, input_json, context);
batch.execution_handles.insert(call_id.clone(), handle);
batch.futures.push(Box::pin(async move {
let result = match terminal.await {
@@ -1420,17 +1502,57 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
)
.into());
}
self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?;
batch.synthetic_results.push_back(result);
batch.dispatched_any = true;
}
PreToolAction::Skip => {}
PreToolAction::Abort(reason) => return Err(EngineError::Aborted(reason)),
PreToolAction::Pause => return Ok(EarlyToolAdmission::Pause),
PreToolAction::Skip => {
self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?;
batch.synthetic_results.push_back(ToolResult::error(
&tool_call.id,
"Tool execution skipped by interceptor",
));
batch.dispatched_any = true;
}
PreToolAction::Abort(reason) => {
self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?;
batch.synthetic_results.push_back(ToolResult::error(
&tool_call.id,
format!("Tool execution aborted: {reason}"),
));
batch.dispatched_any = true;
return Err(EngineError::Aborted(reason));
}
PreToolAction::Pause => {
self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?;
return Ok(EarlyToolAdmission::Pause);
}
}
}
Ok(EarlyToolAdmission::Continue)
}
async fn commit_early_synthetic_results(
&mut self,
history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
batch: &mut EarlyToolExecutionBatch,
) -> Result<(), EngineError> {
while let Some(result) = batch.synthetic_results.pop_front() {
self.finalize_and_commit_tool_result(
history,
annotate,
result,
None,
&batch.call_info_map,
&mut batch.attempt_fence,
&mut batch.terminal_call_ids,
)
.await?;
}
Ok(())
}
async fn cancel_and_terminalize_early_batch(
&mut self,
history: &mut History<A>,
@@ -1439,6 +1561,25 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
) -> Option<EngineError> {
use futures::StreamExt as _;
let mut batch_error = None;
while let Some(result) = batch.synthetic_results.pop_front() {
if let Err(error) = self
.finalize_and_commit_tool_result(
history,
annotate,
result,
None,
&batch.call_info_map,
&mut batch.attempt_fence,
&mut batch.terminal_call_ids,
)
.await
&& batch_error.is_none()
{
batch_error = Some(error);
}
}
let cancellation_request_deadline =
TokioInstant::now() + self.tool_execution_policy.cancellation_request_timeout;
let cancellation_requests = batch
@@ -1462,7 +1603,6 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}
}
let mut batch_error = None;
let deadline =
TokioInstant::now() + self.tool_execution_policy.terminal_confirmation_timeout;
while !batch.futures.is_empty() {
@@ -2131,11 +2271,33 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
"Starting engine run"
);
// Resume pending tool calls from a previous Pause
if let Some(tool_calls) = self.get_pending_tool_calls(history) {
info!("Resuming pending tool calls");
// Resume only calls durably proven not started. A persisted started
// attempt without a terminal result is closed as OutcomeUnknown first.
let (pending_tool_calls, started_call_ids) = self.get_unanswered_tool_calls(history);
if !started_call_ids.is_empty() {
let call_info_map = HashMap::new();
let mut attempt_fence = ToolExecutionAttemptFence::default();
let mut terminal_call_ids = HashSet::new();
for call_id in started_call_ids {
self.finalize_and_commit_tool_result(
history,
annotate,
ToolResult::outcome_unknown(&call_id),
None,
&call_info_map,
&mut attempt_fence,
&mut terminal_call_ids,
)
.await?;
}
}
if !pending_tool_calls.is_empty() {
info!(
count = pending_tool_calls.len(),
"Resuming pending tool calls"
);
if let Some(result) = self
.execute_and_commit_tools(history, annotate, tool_calls)
.execute_and_commit_tools(history, annotate, pending_tool_calls)
.await?
{
return Ok(result);
@@ -2348,17 +2510,28 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
self.request_trace_payload(&request),
);
let request = self.attach_transport_trace(request, current_turn, current_llm_call);
let provider_supports_early = self.client.tool_call_completion_support()
== crate::llm_client::ToolCallCompletionSupport::PerBlock;
let interceptor_supports_early = self.interceptor.supports_early_tool_dispatch();
let early_dispatch = self.tool_call_dispatch_mode
== ToolCallDispatchMode::OnToolCallComplete
&& self.client.tool_call_completion_support()
== crate::llm_client::ToolCallCompletionSupport::PerBlock;
&& provider_supports_early
&& interceptor_supports_early;
if self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete
&& !early_dispatch
&& !provider_supports_early
{
self.emit_warning(
"tool-call early dispatch deferred until response completion: provider does not expose a trustworthy per-call completion boundary",
);
}
if self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete
&& provider_supports_early
&& !interceptor_supports_early
{
self.emit_warning(
"tool-call early dispatch deferred until response completion: interceptor requires the whole-response turn-end gate",
);
}
let assistant_start = history.len();
let StreamResponseOutput {
completion: stream_outcome,
@@ -2466,9 +2639,25 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
// `append_history_items` so observers see each item as it lands.
let reasoning_items = self.thinking_block_collector.take_collected();
let text_blocks = self.text_block_collector.take_collected();
let remaining_tool_calls = self.tool_call_collector.take_collected();
let assistant_items =
self.build_assistant_items(&reasoning_items, &text_blocks, &remaining_tool_calls);
let remaining_collected_calls = self.tool_call_collector.take_collected_with_index();
let mut remaining_tool_calls = remaining_collected_calls
.iter()
.map(|collected| (collected.call_index, collected.call.clone()))
.collect::<Vec<_>>();
remaining_tool_calls.sort_by_key(|(call_index, _)| *call_index);
let remaining_tool_calls = remaining_tool_calls
.into_iter()
.map(|(_, call)| call)
.collect::<Vec<_>>();
let assistant_items = self.build_assistant_items(
&reasoning_items,
&text_blocks,
if early_dispatch {
&[]
} else {
&remaining_tool_calls
},
);
if let Err(error) = self.append_history_items(history, assistant_items, annotate) {
if let Some(batch) = early_tools.take()
&& let Err(cleanup_error) = self
@@ -2489,8 +2678,9 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
let admission = self
.admit_early_tool_calls(
history,
annotate,
early_tools.as_mut().expect("batch initialized"),
remaining_tool_calls.clone(),
remaining_collected_calls,
)
.await;
let admission = match admission {
@@ -2513,10 +2703,27 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}
return Ok(EngineResult::Paused);
}
if let Some(batch) = early_tools.as_mut()
&& let Err(error) = self
.commit_early_synthetic_results(history, annotate, batch)
.await
{
if let Some(mut batch) = early_tools.take() {
let cleanup_error = self
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
.await;
return Err(cleanup_error.unwrap_or(error));
}
return Err(error);
}
}
let mut tool_calls = early_tools
.as_ref()
.map(|batch| batch.calls.clone())
.map(|batch| {
let mut calls = batch.calls.clone();
calls.sort_by_key(|(call_index, _)| *call_index);
calls.into_iter().map(|(_, call)| call).collect()
})
.unwrap_or_default();
if !early_dispatch {
tool_calls = remaining_tool_calls;
@@ -2797,13 +3004,15 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
"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<EarlyToolExecutionBatch> = None;
let mut event_count: usize = 0;
loop {
tokio::select! {
event_result = stream.next() => {
event_result = stream.recv() => {
match event_result {
Some(result) => {
match &result {
@@ -2839,7 +3048,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
self.emit_stream_event(turn, llm_call, &event);
self.timeline.dispatch(&event);
if early_dispatch {
let collected = self.tool_call_collector.take_collected();
let collected = self.tool_call_collector.take_collected_with_index();
if !collected.is_empty() {
if early_tools.is_none() {
let batch_id = format!(
@@ -2854,22 +3063,23 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
let batch = early_tools.as_ref().expect("batch initialized");
collected
.into_iter()
.filter(|call| {
!batch.seen_call_ids.contains(&call.id)
.filter(|collected| {
!batch.seen_call_ids.contains(&collected.call.id)
})
.collect::<Vec<_>>()
};
if !fresh_calls.is_empty() {
// Commit every completed assistant block observed so far,
// including the ToolCall itself, before provider ownership
// can start. This is the durable fail-closed gate.
// Commit every completed reasoning/text block observed so
// far. Admission durably commits each ToolCall with its
// started/not-started state immediately before ownership
// can transfer to a provider.
let reasoning =
self.thinking_block_collector.take_collected();
let text = self.text_block_collector.take_collected();
let assistant_items = self.build_assistant_items(
&reasoning,
&text,
&fresh_calls,
&[],
);
if let Err(error) = self.append_history_items(
history,
@@ -2892,13 +3102,38 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
early_tools.as_mut().expect("batch initialized");
self.admit_early_tool_calls(
history,
annotate,
batch,
fresh_calls,
)
.await
};
match admission {
Ok(EarlyToolAdmission::Continue) => {}
Ok(EarlyToolAdmission::Continue) => {
let batch = early_tools
.as_mut()
.expect("batch initialized");
if let Err(error) = self
.commit_early_synthetic_results(
history,
annotate,
batch,
)
.await
{
if let Some(mut batch) = early_tools.take() {
let cleanup_error = self
.cancel_and_terminalize_early_batch(
history,
annotate,
&mut batch,
)
.await;
return Err(cleanup_error.unwrap_or(error));
}
return Err(error);
}
}
Ok(EarlyToolAdmission::Pause) => {
self.timeline.abort_current_block();
self.timeline.flush_usage();
@@ -2959,6 +3194,45 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
None => break,
}
}
terminal = async {
early_tools
.as_mut()
.expect("terminal branch requires early batch")
.futures
.next()
.await
}, if early_tools.as_ref().is_some_and(|batch| !batch.futures.is_empty()) => {
let (attempt_id, result) =
terminal.expect("non-empty early terminal queue returns a result");
let commit = {
let batch = early_tools.as_mut().expect("batch initialized");
self.finalize_and_commit_tool_result(
history,
annotate,
result,
Some(&attempt_id),
&batch.call_info_map,
&mut batch.attempt_fence,
&mut batch.terminal_call_ids,
)
.await
};
if let Err(error) = commit {
self.timeline.abort_current_block();
self.timeline.flush_usage();
if let Some(mut batch) = early_tools.take() {
let cleanup_error = self
.cancel_and_terminalize_early_batch(
history,
annotate,
&mut batch,
)
.await;
return Err(cleanup_error.unwrap_or(error));
}
return Err(error);
}
}
pause = self.pause_rx.recv() => {
if pause.is_some() {
info!("Paused during response stream");
@@ -3716,6 +3990,41 @@ mod tests {
assert_eq!(projection.original_to_projected_index, [0, 1, 3, 2]);
}
#[test]
fn provider_projection_groups_interleaved_calls_and_results_in_model_order() {
let items = vec![
Item::tool_call_json("call_second", "second", serde_json::json!({}))
.with_tool_execution_metadata(1, Some("batch:call_second".to_string()))
.with_status(ItemStatus::InProgress),
Item::tool_result("call_second", "second result"),
Item::tool_call_json("call_first", "first", serde_json::json!({}))
.with_tool_execution_metadata(0, Some("batch:call_first".to_string()))
.with_status(ItemStatus::InProgress),
Item::tool_result("call_first", "first result"),
];
let projection = materialize_provider_history(&items);
let order = projection
.items
.iter()
.map(|item| match item {
Item::ToolCall { call_id, .. } => format!("call:{call_id}"),
Item::ToolResult { call_id, .. } => format!("result:{call_id}"),
_ => unreachable!(),
})
.collect::<Vec<_>>();
assert_eq!(
order,
[
"call:call_first",
"call:call_second",
"result:call_first",
"result:call_second",
]
);
assert_eq!(projection.original_to_projected_index, [1, 3, 0, 2]);
}
#[test]
fn tool_attachment_round_trips_through_durable_history_json() {
let body: Arc<[u8]> = Arc::from(&b"image-body"[..]);
crates/agen/src/interceptor.rs
+13 -1
View File
@@ -336,6 +336,14 @@ pub struct ToolResultInfo<'a, A = ()> {
/// flows, permission checks, and other trusted host adaptation.
#[async_trait]
pub trait Interceptor<A: Send + Sync = ()>: 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<A: Send + Sync = ()>: Send + Sync {
pub(crate) struct DefaultInterceptor;
#[async_trait]
impl<A: Send + Sync> Interceptor<A> for DefaultInterceptor {}
impl<A: Send + Sync> Interceptor<A> for DefaultInterceptor {
fn supports_early_tool_dispatch(&self) -> bool {
true
}
}
crates/agen/src/llm_client/scheme/openai_responses/events.rs
+496 -32
View File
@@ -28,6 +28,43 @@ pub struct OpenAIResponsesState {
/// 蓄積、`output_item.done`(Reasoning) で既存 reasoning_text block または
/// metadata-only Thinking block に reasoning persistence material を載せる。
pending_reasoning: HashMap<usize, PendingReasoning>,
/// Tool input deltas and authoritative final values, keyed by output index.
pending_tool_inputs: HashMap<usize, PendingToolInput>,
}
/// 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<String>,
}
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<String>,
call_id: String,
name: String,
#[serde(default)]
arguments: Option<String>,
},
CustomToolCall {
#[allow(dead_code)]
@@ -180,6 +219,8 @@ enum OutputItem {
id: Option<String>,
call_id: String,
name: String,
#[serde(default)]
input: Option<String>,
},
#[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<Vec<Event>, 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<Vec<Event>, 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<Vec<Event>, 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<String>,
) -> Result<Vec<Event>, 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>, 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<String> {
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();
crates/agen/src/llm_client/types.rs
+29 -1
View File
@@ -104,7 +104,14 @@ pub enum Item {
name: String,
/// Tool arguments as JSON string
arguments: String,
/// Item status
/// Provider/model order among tool calls in this assistant response.
#[serde(default, skip_serializing_if = "Option::is_none")]
call_index: Option<usize>,
/// Durable identity of a tool execution attempt that has started.
#[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.
#[serde(skip_serializing_if = "Option::is_none")]
status: Option<ItemStatus>,
},
@@ -234,6 +241,8 @@ impl Item {
call_id: call_id.into(),
name: name.into(),
arguments: arguments.into(),
call_index: None,
execution_id: None,
status: None,
}
}
@@ -247,6 +256,25 @@ impl Item {
Self::tool_call(call_id, name, arguments.to_string())
}
/// Attach response-local ordering and, when execution has started, its
/// durable attempt identity to a tool call.
pub fn with_tool_execution_metadata(
mut self,
call_index: usize,
execution_id: Option<String>,
) -> 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<String>, summary: impl Into<String>) -> Self {
Self::tool_result_item(call_id, summary, None, false)
crates/agen/src/timeline/mod.rs
+1
View File
@@ -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
crates/agen/src/timeline/tool_call_collector.rs
+45 -3
View File
@@ -17,10 +17,19 @@ pub struct CollectorState {
/// 現在のツール呼び出し情報 (ブロック進行中)
current_id: Option<String>,
current_name: Option<String>,
/// Zero-based order among tool calls in this assistant response.
call_index: Option<usize>,
/// 蓄積中のJSON入力
input_json_buffer: String,
}
/// One provider-completed call together with its stable model order.
#[derive(Debug, Clone)]
pub(crate) struct CollectedToolCall {
pub call_index: usize,
pub call: ToolCall,
}
/// ToolCallCollector - ToolUseブロックハンドラ
///
/// Timelineに登録してToolUseブロックイベントを受信し、
@@ -28,7 +37,8 @@ pub struct CollectorState {
#[derive(Clone)]
pub struct ToolCallCollector {
/// 収集されたToolCall
collected: Arc<Mutex<Vec<ToolCall>>>,
collected: Arc<Mutex<Vec<CollectedToolCall>>>,
next_call_index: Arc<Mutex<usize>>,
}
impl ToolCallCollector {
@@ -36,18 +46,38 @@ impl ToolCallCollector {
pub fn new() -> Self {
Self {
collected: Arc::new(Mutex::new(Vec::new())),
next_call_index: Arc::new(Mutex::new(0)),
}
}
/// 収集されたToolCallを取得してクリア
pub fn take_collected(&self) -> Vec<ToolCall> {
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<CollectedToolCall> {
let mut guard = self.collected.lock().unwrap();
std::mem::take(&mut *guard)
}
/// 収集されたToolCallの参照を取得
pub fn collected(&self) -> Vec<ToolCall> {
self.collected.lock().unwrap().clone()
self.collected
.lock()
.unwrap()
.iter()
.map(|collected| collected.call.clone())
.collect()
}
/// Start one response-local tool-call ordering scope.
pub(crate) fn begin_response(&self) {
self.collected.lock().unwrap().clear();
*self.next_call_index.lock().unwrap() = 0;
}
/// 収集されたToolCallがあるかどうか
@@ -75,6 +105,11 @@ impl Handler<ToolUseBlockKind> for ToolCallCollector {
ToolUseBlockEvent::Start(start) => {
scope.current_id = Some(start.id.clone());
scope.current_name = Some(start.name.clone());
if scope.call_index.is_none() {
let mut next = self.next_call_index.lock().unwrap();
scope.call_index = Some(*next);
*next += 1;
}
scope.input_json_buffer.clear();
}
ToolUseBlockEvent::InputJsonDelta(delta) => {
@@ -98,8 +133,15 @@ impl Handler<ToolUseBlockKind> for ToolCallCollector {
};
let tool_call = ToolCall { id, name, input };
let call_index = scope
.call_index
.take()
.expect("tool-use stop follows a started collector scope");
self.collected.lock().unwrap().push(tool_call);
self.collected.lock().unwrap().push(CollectedToolCall {
call_index,
call: tool_call,
});
}
scope.input_json_buffer.clear();
}
crates/agen/tests/common/mod.rs
+12 -1
View File
@@ -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<Vec<Vec<Event>>>,
call_count: Arc<AtomicUsize>,
requests: Arc<Mutex<Vec<Request>>>,
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<Request> {
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,
crates/agen/tests/parallel_execution_test.rs
+420 -2
View File
@@ -7,8 +7,9 @@ use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use agen::interceptor::{
Interceptor, InterceptorError, InterceptorErrorCategory, InterceptorPhase, InterceptorResult,
PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo,
AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorErrorCategory,
InterceptorPhase, InterceptorResult, PostToolAction, PreToolAction, ToolCallInfo,
ToolResultInfo, TurnEndAction,
};
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::llm_client::{
@@ -307,6 +308,46 @@ impl LlmClient for ControlledStreamClient {
}
}
#[derive(Clone)]
struct BoundedStreamClient {
receiver: Arc<Mutex<Option<tokio::sync::mpsc::Receiver<Result<Event, ClientError>>>>>,
}
impl BoundedStreamClient {
fn new(capacity: usize) -> (Self, tokio::sync::mpsc::Sender<Result<Event, ClientError>>) {
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<ResponseStream, ClientError> {
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<dyn LlmClient> {
Box::new(self.clone())
}
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
ToolCallCompletionSupport::PerBlock
}
}
#[derive(Clone)]
struct BarrierTool {
name: &'static str,
@@ -400,6 +441,141 @@ async fn default_dispatch_waits_for_response_completion() {
assert_eq!(probe.starts.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn turn_end_policy_incompatibility_visibly_falls_back_before_side_effects() {
struct TurnEndPause;
#[async_trait]
impl Interceptor for TurnEndPause {
async fn on_assistant_turn_end(
&self,
_context: AssistantTurnEndContext<'_>,
) -> InterceptorResult<TurnEndAction> {
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<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("blocked_policy_tool", 0);
let probe = tool.clone();
let mut engine = Engine::new(client);
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
engine.set_interceptor(BlockingEarlyPolicy {
entered: entered.clone(),
release,
});
engine.register_tool(tool.definition());
let pause = engine.pause_sender();
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
tx.send(Ok(Event::tool_use_start(
0,
"call_blocked_policy",
"blocked_policy_tool",
)))
.await
.unwrap();
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#)))
.await
.unwrap();
tx.send(Ok(Event::tool_use_stop(0))).await.unwrap();
tokio::time::timeout(Duration::from_secs(1), entered.notified())
.await
.expect("pre-tool policy entered");
tx.send(Ok(Event::ping())).await.unwrap();
tokio::time::timeout(Duration::from_secs(1), tx.send(Ok(Event::ping())))
.await
.expect("background pump continues draining the bounded provider stream")
.unwrap();
pause.send(()).await.unwrap();
drop(tx);
let (output, _history) = run.await.unwrap();
assert!(matches!(output.result, EngineRunExit::Paused));
assert_eq!(probe.call_count(), 0);
}
#[tokio::test]
async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_result() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
@@ -457,6 +633,146 @@ async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_
assert_eq!((calls, results), (2, 2));
}
#[tokio::test]
async fn early_terminal_result_commits_while_response_stream_is_open() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let tool = BarrierTool::new("commit_while_streaming");
let probe = tool.clone();
let mut engine = Engine::new(client);
engine.set_max_turns(Some(1));
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
engine.register_tool(tool.definition());
let (committed_tx, mut committed_rx) = tokio::sync::mpsc::unbounded_channel();
engine.on_history_append(move |item| {
if matches!(item, Item::ToolResult { call_id, .. } if call_id == "call_commit_open") {
let _ = committed_tx.send(());
}
Ok(())
});
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
tx.send(Ok(Event::tool_use_start(
0,
"call_commit_open",
"commit_while_streaming",
)))
.unwrap();
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
tokio::time::timeout(Duration::from_secs(1), probe.started.notified())
.await
.expect("tool starts while response remains open");
probe.release.notify_one();
tokio::time::timeout(Duration::from_secs(1), committed_rx.recv())
.await
.expect("terminal result is durably committed before response release")
.expect("commit notification channel remains open");
tx.send(Ok(Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})))
.unwrap();
drop(tx);
let _ = run.await.unwrap();
}
#[tokio::test]
async fn reverse_stop_order_preserves_model_call_indexes() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let contexts = Arc::new(Mutex::new(Vec::new()));
let mut engine = Engine::new(client);
engine.set_max_turns(Some(1));
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
engine.register_tool(ContextRecordingTool::new("ordered_a", contexts.clone()).definition());
engine.register_tool(ContextRecordingTool::new("ordered_b", contexts.clone()).definition());
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
tx.send(Ok(Event::tool_use_start(0, "call_order_a", "ordered_a")))
.unwrap();
tx.send(Ok(Event::tool_use_start(1, "call_order_b", "ordered_b")))
.unwrap();
tx.send(Ok(Event::tool_input_delta(1, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(1))).unwrap();
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
tx.send(Ok(Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})))
.unwrap();
drop(tx);
let _ = run.await.unwrap();
let by_id = contexts
.lock()
.unwrap()
.iter()
.map(|context| (context.call_id.clone(), context.call_index))
.collect::<std::collections::HashMap<_, _>>();
assert_eq!(by_id.get("call_order_a"), Some(&0));
assert_eq!(by_id.get("call_order_b"), Some(&1));
}
#[tokio::test]
async fn reverse_stop_order_projects_next_provider_request_in_model_order() {
let first_response = vec![
Event::tool_use_start(0, "call_project_a", "project_a"),
Event::tool_use_start(1, "call_project_b", "project_b"),
Event::tool_input_delta(1, r#"{}"#),
Event::tool_use_stop(1),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
];
let second_response = vec![
Event::text_delta(0, "done"),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
];
let client = MockLlmClient::with_responses(vec![first_response, second_response])
.with_completion_support(ToolCallCompletionSupport::PerBlock);
let probe = client.clone();
let contexts = Arc::new(Mutex::new(Vec::new()));
let mut engine = Engine::new(client);
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
engine.register_tool(ContextRecordingTool::new("project_a", contexts.clone()).definition());
engine.register_tool(ContextRecordingTool::new("project_b", contexts).definition());
let mut history = History::new();
let _ = engine.run(&mut history, "run").await;
let requests = probe.requests();
assert_eq!(requests.len(), 2);
let tool_order = requests[1]
.items
.iter()
.filter_map(|item| match item {
Item::ToolCall { call_id, .. } => Some(format!("call:{call_id}")),
Item::ToolResult { call_id, .. } => Some(format!("result:{call_id}")),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(
tool_order,
[
"call:call_project_a",
"call:call_project_b",
"result:call_project_a",
"result:call_project_b",
]
);
}
#[tokio::test]
async fn later_history_failure_terminalizes_already_started_early_sibling() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
@@ -525,6 +841,108 @@ async fn later_history_failure_terminalizes_already_started_early_sibling() {
)));
}
#[tokio::test]
async fn persisted_started_call_is_closed_unknown_and_never_resumed() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let tool = SlowTool::new("must_not_resume_started", 0);
let probe = tool.clone();
let mut engine = Engine::new(client);
engine.set_max_turns(Some(1));
engine.register_tool(tool.definition());
let mut history = History::from_items(vec![
Item::user_message("run"),
Item::tool_call_json(
"call_started_before_restore",
"must_not_resume_started",
serde_json::json!({}),
)
.with_tool_execution_metadata(
0,
Some("restored-batch:call_started_before_restore".to_string()),
)
.with_status(agen::llm_client::ItemStatus::InProgress),
]);
let run = tokio::spawn(async move {
let output = engine.resume(&mut history).await;
(output, history)
});
tx.send(Ok(Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})))
.unwrap();
drop(tx);
let (_output, history) = run.await.unwrap();
assert_eq!(probe.call_count(), 0);
assert!(history.iter().any(|entry| matches!(
&entry.item,
Item::ToolResult {
call_id,
disposition: ToolResultDisposition::OutcomeUnknown,
..
} if call_id == "call_started_before_restore"
)));
}
#[tokio::test]
async fn abort_during_early_admission_commits_synthetic_terminal() {
struct AbortAdmission;
#[async_trait]
impl Interceptor for AbortAdmission {
fn supports_early_tool_dispatch(&self) -> bool {
true
}
async fn pre_tool_call(
&self,
_info: &mut ToolCallInfo<'_>,
) -> InterceptorResult<PreToolAction> {
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);
crates/session-store/src/logged_item.rs
+31 -5
View File
@@ -7,13 +7,14 @@
//! save / replay boundaries.
//!
//! Fields kept here are limited to what is needed to reconstruct a worker
//! `Item` for replay. `id` and `status` annotations are intentionally dropped
//! (they are output-side metadata; replayed items synthesize fresh `None`).
//! `Item` for replay. `id` annotations are intentionally dropped because they
//! are output-side metadata. ToolCall execution status/order are retained because
//! they fence durable resume from replaying an already-started side effect.
//! `Reasoning::encrypted_content` is preserved because OpenAI Responses ZDR
//! requires it on stateless re-send.
use agen::{
llm_client::types::{ContentPart, Item, Role},
llm_client::types::{ContentPart, Item, ItemStatus, Role},
tool::{Attachment, ImageAttachment, ToolResultDisposition},
};
use base64::{Engine as _, engine::general_purpose::STANDARD};
@@ -53,6 +54,12 @@ pub enum LoggedItem {
call_id: String,
name: String,
arguments: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
call_index: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
execution_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
status: Option<ItemStatus>,
},
ToolResult {
call_id: String,
@@ -119,11 +126,17 @@ impl From<&Item> for LoggedItem {
call_id,
name,
arguments,
call_index,
execution_id,
status,
..
} => Self::ToolCall {
call_id: call_id.clone(),
name: name.clone(),
arguments: arguments.clone(),
call_index: *call_index,
execution_id: execution_id.clone(),
status: *status,
},
Item::ToolResult {
call_id,
@@ -176,12 +189,17 @@ impl From<LoggedItem> for Item {
call_id,
name,
arguments,
call_index,
execution_id,
status,
} => Item::ToolCall {
id: None,
call_id,
name,
arguments,
status: None,
call_index,
execution_id,
status,
},
LoggedItem::ToolResult {
call_id,
@@ -327,7 +345,9 @@ mod tests {
#[test]
fn round_trip_tool_call() {
let original = Item::tool_call("call_42", "get_weather", r#"{"city":"Tokyo"}"#);
let original = Item::tool_call("call_42", "get_weather", r#"{"city":"Tokyo"}"#)
.with_tool_execution_metadata(3, Some("batch:call_42".to_string()))
.with_status(ItemStatus::InProgress);
let logged: LoggedItem = (&original).into();
let json = serde_json::to_string(&logged).unwrap();
let parsed: LoggedItem = serde_json::from_str(&json).unwrap();
@@ -336,11 +356,17 @@ mod tests {
call_id,
name,
arguments,
call_index,
execution_id,
status,
..
} => {
assert_eq!(call_id, "call_42");
assert_eq!(name, "get_weather");
assert_eq!(arguments, r#"{"city":"Tokyo"}"#);
assert_eq!(call_index, Some(3));
assert_eq!(execution_id.as_deref(), Some("batch:call_42"));
assert_eq!(status, Some(ItemStatus::InProgress));
}
other => panic!("unexpected variant: {other:?}"),
}
crates/session-store/src/public_snapshot.rs
+1
View File
@@ -356,6 +356,7 @@ fn project_item(item: &LoggedItem) -> Option<SessionSnapshotEntryData> {
call_id,
name,
arguments,
..
} => Some(SessionSnapshotEntryData::ToolCall {
call_id: call_id.clone(),
name: name.clone(),
crates/worker/src/ipc/interceptor.rs
+4
View File
@@ -363,6 +363,10 @@ impl WorkerInterceptor {
#[async_trait]
impl Interceptor<SessionHistoryMetadata> 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>,
crates/worker/src/spawn/registry.rs
+3
View File
@@ -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,
},
),
});
docs/manifest.toml
+3 -2
View File
@@ -123,8 +123,9 @@ ref = "anthropic/claude-sonnet-4-6"
# 任意。デフォルト: "after_response"。
# "on_tool_call_complete" は provider が個別 ToolCall の完成を保証する場合だけ、
# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider では
# 応答完了まで明示的に遅延する。これは直列実行モードではなく、先行 call の完了を
# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider、または
# 応答全体を見て実行前 Pause を判断する turn-end hook がある場合は、応答完了まで
# 警告付きで明示的に遅延する。これは直列実行モードではなく、先行 call の完了を
# 待たず後続 call も開始するため、完了順・DB 反映順・採番順は保証しない。
# 早期開始後に stream が切断された場合は自動 continuation せず、安全に terminalize して停止する。
# tool_call_dispatch = "on_tool_call_complete"