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

feat: integrate early tool call dispatch

Browse Source
This commit is contained in:
Hare
2026-09-29 14:02:11 +09:00
parent d81c9a8154 8ff7df7edf
commit 38c2a78b1c
33 changed files with 3962 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/src/engine.rs
+1371 -85
View File
File diff suppressed because it is too large Load Diff
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/lib.rs
+1 -1
View File
@@ -22,7 +22,7 @@ pub use agen_macros::{description, tool, tool_registry};
pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope};
pub use engine::{
Engine, EngineConfig, EngineError, EngineResult, EngineRunExit, EngineRunOutput,
LlmRetryNotice, RunInterruptionReason, ToolRegistryError,
LlmRetryNotice, RunInterruptionReason, ToolCallDispatchMode, ToolRegistryError,
};
pub use handler::ToolUseBlockStart;
pub use history::{History, HistoryEntry};
crates/agen/src/llm_client/client.rs
+19
View File
@@ -38,6 +38,16 @@ impl std::fmt::Display for ConfigWarning {
pub type ResponseStream = Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>;
/// Whether normalized `BlockStop(ToolUse)` events prove one call is complete.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum ToolCallCompletionSupport {
/// Individual tool-call completion events are provider-authoritative.
PerBlock,
/// Tool calls are complete only once the response stream itself completes.
#[default]
ResponseComplete,
}
/// LLMクライアントのtrait
///
/// 各プロバイダはこのtraitを実装し、統一されたインターフェースを提供する。
@@ -59,6 +69,11 @@ pub trait LlmClient: Send + Sync {
/// compaction) without access to the original construction parameters.
fn clone_boxed(&self) -> Box<dyn LlmClient>;
/// Describe the strongest safe tool-call completion boundary.
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
ToolCallCompletionSupport::ResponseComplete
}
/// 設定をバリデーションし、未サポートの設定があれば警告を返す
///
/// # Arguments
@@ -92,6 +107,10 @@ impl LlmClient for Box<dyn LlmClient> {
(**self).clone_boxed()
}
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
(**self).tool_call_completion_support()
}
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
(**self).validate_config(config)
}
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/anthropic/scheme_impl.rs
+5
View File
@@ -9,6 +9,7 @@ use crate::llm_client::{
ClientError,
auth::AuthRequirement,
capability::ModelCapability,
client::ToolCallCompletionSupport,
event::{BlockType, Event, ReasoningBlockData},
scheme::Scheme,
types::Request,
@@ -104,4 +105,8 @@ impl Scheme for AnthropicScheme {
fn default_capability(&self) -> ModelCapability {
super::capability::default_capability()
}
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
ToolCallCompletionSupport::PerBlock
}
}
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/mod.rs
+32 -1
View File
@@ -16,7 +16,7 @@ use serde_json::Value;
use super::auth::AuthRequirement;
use super::capability::ModelCapability;
use super::client::ConfigWarning;
use super::client::{ConfigWarning, ToolCallCompletionSupport};
use super::error::ClientError;
use super::event::Event;
use super::types::{Request, RequestConfig};
@@ -81,6 +81,11 @@ pub trait Scheme: Clone + Send + Sync + 'static {
/// 高レベルの client 構築層の責務で、scheme はここには関与しない。
fn default_capability(&self) -> ModelCapability;
/// Strongest safe completion boundary for streamed tool calls.
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
ToolCallCompletionSupport::ResponseComplete
}
/// scheme 側でサポートしていない `RequestConfig` フィールドを
/// 警告として返す（例: OpenAI Chat は `top_k` 非対応）。
/// デフォルトは空 Vec。
@@ -89,3 +94,29 @@ pub trait Scheme: Clone + Send + Sync + 'static {
Vec::new()
}
}
#[cfg(test)]
mod completion_support_tests {
use super::*;
use crate::llm_client::ToolCallCompletionSupport;
#[test]
fn only_schemes_with_authoritative_per_call_done_events_enable_early_dispatch() {
assert_eq!(
anthropic::AnthropicScheme::new().tool_call_completion_support(),
ToolCallCompletionSupport::PerBlock
);
assert_eq!(
openai_responses::OpenAIResponsesScheme::new().tool_call_completion_support(),
ToolCallCompletionSupport::PerBlock
);
assert_eq!(
openai_chat::OpenAIScheme::new().tool_call_completion_support(),
ToolCallCompletionSupport::ResponseComplete
);
assert_eq!(
gemini::GeminiScheme::new().tool_call_completion_support(),
ToolCallCompletionSupport::ResponseComplete
);
}
}
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/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/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/scheme/openai_responses/scheme_impl.rs
+5 -1
View File
@@ -6,7 +6,7 @@ use crate::llm_client::{
ClientError,
auth::AuthRequirement,
capability::ModelCapability,
client::ConfigWarning,
client::{ConfigWarning, ToolCallCompletionSupport},
event::Event,
scheme::Scheme,
types::{Request, RequestConfig},
@@ -56,6 +56,10 @@ impl Scheme for OpenAIResponsesScheme {
super::capability::default_capability()
}
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
ToolCallCompletionSupport::PerBlock
}
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
let mut warnings = Vec::new();
// Some compatible backends reject `max_output_tokens` with HTTP 400.
crates/agen/src/llm_client/transport.rs
+4
View File
@@ -510,6 +510,10 @@ impl<S: Scheme + Clone + 'static> LlmClient for HttpTransport<S> {
self.scheme.validate_config(config)
}
fn tool_call_completion_support(&self) -> super::ToolCallCompletionSupport {
self.scheme.tool_call_completion_support()
}
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError> {
let total_started = Instant::now();
let path = self.scheme.path(&self.model_id);
crates/agen/src/llm_client/types.rs
+48 -1
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
@@ -104,7 +110,15 @@ pub enum Item {
name: String,
/// Tool arguments as JSON string
arguments: String,
/// Item status
/// Provider/model order among tool calls in this assistant response.
#[serde(default, skip_serializing_if = "Option::is_none")]
call_index: Option<usize>,
/// 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 execution or
/// terminalization ownership was transferred and must never be inferred pending.
#[serde(skip_serializing_if = "Option::is_none")]
status: Option<ItemStatus>,
},
@@ -219,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
// ========================================================================
@@ -234,6 +256,8 @@ impl Item {
call_id: call_id.into(),
name: name.into(),
arguments: arguments.into(),
call_index: None,
execution_id: None,
status: None,
}
}
@@ -247,6 +271,25 @@ impl Item {
Self::tool_call(call_id, name, arguments.to_string())
}
/// Attach response-local ordering and, when execution has started, its
/// durable attempt identity to a tool call.
pub fn with_tool_execution_metadata(
mut self,
call_index: usize,
execution_id: Option<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)
@@ -377,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()),
@@ -388,6 +432,7 @@ impl Item {
pub fn with_status(mut self, new_status: ItemStatus) -> Self {
match &mut self {
Self::Message { status, .. } => *status = Some(new_status),
Self::AssistantResponseBoundary { .. } => {}
Self::ToolCall { status, .. } => *status = Some(new_status),
Self::ToolResult { .. } => {} // Result items don't have status
Self::Reasoning { status, .. } => *status = Some(new_status),
@@ -403,6 +448,7 @@ impl Item {
pub fn id(&self) -> Option<&str> {
match self {
Self::Message { id, .. } => id.as_deref(),
Self::AssistantResponseBoundary { .. } => None,
Self::ToolCall { id, .. } => id.as_deref(),
Self::ToolResult { id, .. } => id.as_deref(),
Self::Reasoning { id, .. } => id.as_deref(),
@@ -413,6 +459,7 @@ impl Item {
pub fn item_type(&self) -> &'static str {
match self {
Self::Message { .. } => "message",
Self::AssistantResponseBoundary { .. } => "assistant_response_boundary",
Self::ToolCall { .. } => "tool_call",
Self::ToolResult { .. } => "tool_result",
Self::Reasoning { .. } => "reasoning",
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/timeline.rs
+111 -53
View File
@@ -256,8 +256,8 @@ where
H: Handler<ToolUseBlockKind>,
{
handler: H,
scope: Option<H::Scope>,
current_tool: Option<(String, String)>, // (id, name)
scopes: HashMap<usize, H::Scope>,
tools: HashMap<usize, (String, String)>, // index -> (id, name)
}
impl<H> ToolUseBlockHandlerWrapper<H>
@@ -267,8 +267,8 @@ where
fn new(handler: H) -> Self {
Self {
handler,
scope: None,
current_tool: None,
scopes: HashMap::new(),
tools: HashMap::new(),
}
}
}
@@ -279,60 +279,55 @@ where
H::Scope: Send + Sync,
{
fn dispatch_start(&mut self, start: &BlockStart) {
if let Some(scope) = &mut self.scope {
if let BlockMetadata::ToolUse { id, name } = &start.metadata {
self.current_tool = Some((id.clone(), name.clone()));
self.handler.on_event(
scope,
&ToolUseBlockEvent::Start(ToolUseBlockStart {
index: start.index,
id: id.clone(),
name: name.clone(),
}),
);
}
if let BlockMetadata::ToolUse { id, name } = &start.metadata {
let scope = self.scopes.entry(start.index).or_default();
self.tools.insert(start.index, (id.clone(), name.clone()));
self.handler.on_event(
scope,
&ToolUseBlockEvent::Start(ToolUseBlockStart {
index: start.index,
id: id.clone(),
name: name.clone(),
}),
);
}
}
fn dispatch_delta(&mut self, delta: &BlockDelta) {
if let Some(scope) = &mut self.scope {
if let DeltaContent::InputJson(json) = &delta.delta {
self.handler
.on_event(scope, &ToolUseBlockEvent::InputJsonDelta(json.clone()));
}
if let DeltaContent::InputJson(json) = &delta.delta {
let scope = self.scopes.entry(delta.index).or_default();
self.handler
.on_event(scope, &ToolUseBlockEvent::InputJsonDelta(json.clone()));
}
}
fn dispatch_stop(&mut self, stop: &BlockStop) {
if let Some(scope) = &mut self.scope {
if let Some((id, name)) = self.current_tool.take() {
self.handler.on_event(
scope,
&ToolUseBlockEvent::Stop(ToolUseBlockStop {
index: stop.index,
id,
name,
}),
);
}
if let (Some(mut scope), Some((id, name))) = (
self.scopes.remove(&stop.index),
self.tools.remove(&stop.index),
) {
self.handler.on_event(
&mut scope,
&ToolUseBlockEvent::Stop(ToolUseBlockStop {
index: stop.index,
id,
name,
}),
);
}
}
fn dispatch_abort(&mut self, _abort: &BlockAbort) {
self.current_tool = None;
fn dispatch_abort(&mut self, abort: &BlockAbort) {
self.scopes.remove(&abort.index);
self.tools.remove(&abort.index);
}
fn start_scope(&mut self) {
self.scope = Some(H::Scope::default());
}
fn start_scope(&mut self) {}
fn end_scope(&mut self) {
self.scope = None;
self.current_tool = None;
}
fn end_scope(&mut self) {}
fn has_scope(&self) -> bool {
self.scope.is_some()
!self.scopes.is_empty()
}
}
@@ -379,7 +374,9 @@ pub struct Timeline {
thinking_block_handlers: Vec<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 の集約バッファ。
@@ -405,6 +402,7 @@ impl Timeline {
text_block_handlers: Vec::new(),
thinking_block_handlers: Vec::new(),
tool_use_block_handlers: Vec::new(),
active_blocks: HashMap::new(),
current_block: None,
pending_usage: None,
}
@@ -559,6 +557,8 @@ impl Timeline {
}
fn handle_block_start(&mut self, start: &BlockStart) {
self.active_blocks
.insert((start.block_type, start.index), ());
self.current_block = Some(start.block_type);
let handlers = self.get_block_handlers_mut(start.block_type);
@@ -572,10 +572,11 @@ impl Timeline {
let block_type = delta.delta.block_type();
// OpenAIなどのプロバイダはBlockStartを送らない場合があるため、
// Deltaが来たときにスコープがなければ暗黙的に開始する
if self.current_block.is_none() {
self.current_block = Some(block_type);
}
// Deltaが来たときにその index のスコープがなければ暗黙的に開始する
self.active_blocks
.entry((block_type, delta.index))
.or_insert(());
self.current_block = Some(block_type);
let handlers = self.get_block_handlers_mut(block_type);
for handler in handlers {
@@ -593,7 +594,12 @@ impl Timeline {
handler.dispatch_stop(stop);
handler.end_scope();
}
self.current_block = None;
self.active_blocks.remove(&(stop.block_type, stop.index));
self.current_block = self
.active_blocks
.keys()
.next()
.map(|(block_type, _)| *block_type);
}
fn handle_block_abort(&mut self, abort: &BlockAbort) {
@@ -602,7 +608,12 @@ impl Timeline {
handler.dispatch_abort(abort);
handler.end_scope();
}
self.current_block = None;
self.active_blocks.remove(&(abort.block_type, abort.index));
self.current_block = self
.active_blocks
.keys()
.next()
.map(|(block_type, _)| *block_type);
}
fn get_block_handlers_mut(
@@ -622,19 +633,22 @@ impl Timeline {
self.current_block
}
/// 現在アクティブなブロックを中断する
/// 現在アクティブな全ブロックを中断する
///
/// キャンセルやエラー時に呼び出し、進行中のブロックに対して
/// キャンセルやエラー時に呼び出し、index ごとに進行中の全ブロックへ
/// BlockAbortイベントを発火してスコープをクリーンアップする。
pub fn abort_current_block(&mut self) {
if let Some(block_type) = self.current_block {
let active = self.active_blocks.keys().copied().collect::<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;
}
}
@@ -679,6 +693,50 @@ mod tests {
assert!(calls.lock().unwrap().is_empty());
}
#[test]
fn abort_clears_every_indexed_scope_before_indices_are_reused() {
#[derive(Default)]
struct Scope {
started: bool,
}
struct IndexedToolHandler {
reused_cleanly: Arc<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/src/timeline/tool_call_collector.rs
+94 -5
View File
@@ -5,7 +5,6 @@
use crate::{
handler::{Handler, ToolUseBlockEvent, ToolUseBlockKind},
llm_client::types::parse_tool_arguments,
tool::ToolCall,
};
use std::sync::{Arc, Mutex};
@@ -18,10 +17,19 @@ pub struct CollectorState {
/// 現在のツール呼び出し情報 (ブロック進行中)
current_id: Option<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ブロックイベントを受信し、
@@ -29,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 {
@@ -37,18 +46,38 @@ impl ToolCallCollector {
pub fn new() -> Self {
Self {
collected: Arc::new(Mutex::new(Vec::new())),
next_call_index: Arc::new(Mutex::new(0)),
}
}
/// 収集されたToolCallを取得してクリア
pub fn take_collected(&self) -> Vec<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があるかどうか
@@ -76,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) => {
@@ -85,11 +119,29 @@ impl Handler<ToolUseBlockKind> for ToolCallCollector {
// ブロック完了時にToolCallを確定
if let (Some(id), Some(name)) = (scope.current_id.take(), scope.current_name.take())
{
let input = parse_tool_arguments(&scope.input_json_buffer);
// A Stop event proves the provider finished the block, but
// malformed/non-object arguments are still not executable.
// Preserve an explicit null sentinel so the engine can commit
// a terminal validation error without guessing `{}`.
let input = if scope.input_json_buffer.trim().is_empty() {
serde_json::Value::Object(serde_json::Map::new())
} else {
serde_json::from_str::<serde_json::Value>(&scope.input_json_buffer)
.ok()
.filter(serde_json::Value::is_object)
.unwrap_or(serde_json::Value::Null)
};
let tool_call = ToolCall { id, name, input };
let call_index = scope
.call_index
.take()
.expect("tool-use stop follows a started collector scope");
self.collected.lock().unwrap().push(tool_call);
self.collected.lock().unwrap().push(CollectedToolCall {
call_index,
call: tool_call,
});
}
scope.input_json_buffer.clear();
}
@@ -144,6 +196,43 @@ mod tests {
);
}
#[test]
fn test_collect_interleaved_tool_call_indexes_and_ignore_duplicate_stop() {
let collector = ToolCallCollector::new();
let mut timeline = Timeline::new();
timeline.on_tool_use_block(collector.clone());
timeline.dispatch(&Event::tool_use_start(0, "call_0", "tool_0"));
timeline.dispatch(&Event::tool_input_delta(0, r#"{"a":"#));
timeline.dispatch(&Event::tool_use_start(1, "call_1", "tool_1"));
timeline.dispatch(&Event::tool_input_delta(1, r#"{"b":2}"#));
timeline.dispatch(&Event::tool_input_delta(0, r#"1}"#));
timeline.dispatch(&Event::tool_use_stop(1));
timeline.dispatch(&Event::tool_use_stop(1));
timeline.dispatch(&Event::tool_use_stop(0));
let calls = collector.take_collected();
assert_eq!(calls.len(), 2);
assert_eq!(calls[0].id, "call_1");
assert_eq!(calls[0].input, serde_json::json!({"b": 2}));
assert_eq!(calls[1].id, "call_0");
assert_eq!(calls[1].input, serde_json::json!({"a": 1}));
}
#[test]
fn malformed_completed_arguments_are_not_normalized_to_executable_object() {
let collector = ToolCallCollector::new();
let mut timeline = Timeline::new();
timeline.on_tool_use_block(collector.clone());
timeline.dispatch(&Event::tool_use_start(0, "bad", "mutate"));
timeline.dispatch(&Event::tool_input_delta(0, r#"{"unterminated":"#));
timeline.dispatch(&Event::tool_use_stop(0));
let calls = collector.take_collected();
assert_eq!(calls.len(), 1);
assert!(calls[0].input.is_null());
}
#[test]
fn test_collect_multiple_tool_calls() {
let collector = ToolCallCollector::new();
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/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
+1435 -3
View File
File diff suppressed because it is too large Load Diff
crates/agen/tests/response_completion_compatibility_test.rs
+142
View File
@@ -0,0 +1,142 @@
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use agen::llm_client::{
ClientError, LlmClient, Request, ResponseStream, ToolCallCompletionSupport,
scheme::{Scheme, gemini::GeminiScheme, openai_chat::OpenAIScheme},
};
use agen::{Engine, EngineRunExit, History, ToolCallDispatchMode};
use async_trait::async_trait;
#[derive(Clone)]
struct SseFrame {
event_type: &'static str,
data: &'static str,
}
#[derive(Clone)]
struct SchemeReplayClient<S> {
scheme: S,
responses: Arc<Vec<Vec<SseFrame>>>,
stream_calls: Arc<AtomicUsize>,
}
impl<S> SchemeReplayClient<S> {
fn new(scheme: S, responses: Vec<Vec<SseFrame>>) -> Self {
Self {
scheme,
responses: Arc::new(responses),
stream_calls: Arc::new(AtomicUsize::new(0)),
}
}
fn stream_count(&self) -> usize {
self.stream_calls.load(Ordering::SeqCst)
}
}
#[async_trait]
impl<S: Scheme> LlmClient for SchemeReplayClient<S> {
async fn stream(&self, _request: Request) -> Result<ResponseStream, ClientError> {
let response_index = self.stream_calls.fetch_add(1, Ordering::SeqCst);
let frames = self
.responses
.get(response_index)
.ok_or_else(|| ClientError::Config("scheme replay responses exhausted".to_string()))?;
let mut state = S::State::default();
let mut output = Vec::new();
for frame in frames {
match self
.scheme
.parse_sse(frame.event_type, frame.data, &mut state)
{
Ok(events) => output.extend(events.into_iter().map(Ok)),
Err(error) => {
output.push(Err(error));
break;
}
}
}
Ok(Box::pin(futures::stream::iter(output)))
}
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
self.scheme.tool_call_completion_support()
}
}
async fn assert_response_complete_scheme_accepts_clean_eof<S: Scheme>(
scheme: S,
frames: Vec<SseFrame>,
) {
assert_eq!(
scheme.tool_call_completion_support(),
ToolCallCompletionSupport::ResponseComplete
);
for mode in [
ToolCallDispatchMode::AfterResponse,
ToolCallDispatchMode::OnToolCallComplete,
] {
let client = SchemeReplayClient::new(scheme.clone(), vec![frames.clone()]);
let probe = client.clone();
let mut engine = Engine::new(client);
engine.set_tool_call_dispatch_mode(mode);
let mut history = History::new();
let output = engine.run(&mut history, "hello").await;
assert!(
matches!(output.result, EngineRunExit::Finished),
"{mode:?} should accept the response-complete scheme's clean EOF, got {:?}",
output.result
);
assert_eq!(
probe.stream_count(),
1,
"{mode:?} must not open a continuation request after the adapter's normal EOF"
);
}
}
#[tokio::test]
async fn openai_chat_clean_eof_completes_in_default_and_opt_in_fallback_modes() {
assert_response_complete_scheme_accepts_clean_eof(
OpenAIScheme::new(),
vec![
SseFrame {
event_type: "message",
data: r#"{"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":"stop"}]}"#,
},
SseFrame {
event_type: "message",
data: "[DONE]",
},
],
)
.await;
}
#[tokio::test]
async fn gemini_clean_eof_completes_in_default_and_opt_in_fallback_modes() {
assert_response_complete_scheme_accepts_clean_eof(
GeminiScheme::new(),
vec![
SseFrame {
event_type: "message",
data: r#"{"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"},"finishReason":"STOP","index":0}]}"#,
},
SseFrame {
event_type: "message",
data: "[DONE]",
},
],
)
.await;
}
crates/manifest/src/config.rs
+37
View File
@@ -542,6 +542,8 @@ pub struct EngineManifestConfig {
#[serde(default)]
pub reasoning: Option<ReasoningControl>,
#[serde(default)]
pub tool_call_dispatch: Option<crate::ToolCallDispatchMode>,
#[serde(default)]
pub tool_output: ToolOutputLimitsPartial,
#[serde(default)]
pub file_upload: FileUploadLimitsPartial,
@@ -868,6 +870,7 @@ impl EngineManifestConfig {
top_k: upper.top_k.or(self.top_k),
stop_sequences: upper.stop_sequences.or(self.stop_sequences),
reasoning: upper.reasoning.or(self.reasoning),
tool_call_dispatch: upper.tool_call_dispatch.or(self.tool_call_dispatch),
tool_output: self.tool_output.merge(upper.tool_output),
file_upload: self.file_upload.merge(upper.file_upload),
}
@@ -1171,6 +1174,7 @@ impl TryFrom<WorkerManifestConfig> for WorkerManifest {
top_k: cfg.engine.top_k,
stop_sequences: cfg.engine.stop_sequences.unwrap_or_default(),
reasoning: cfg.engine.reasoning,
tool_call_dispatch: cfg.engine.tool_call_dispatch.unwrap_or_default(),
tool_output: ToolOutputLimits {
default_max_bytes: cfg
.engine
@@ -1343,6 +1347,39 @@ mod tests {
}
}
#[test]
fn tool_call_dispatch_defaults_merges_and_round_trips() {
let manifest: WorkerManifest = minimal_valid().try_into().unwrap();
assert_eq!(
manifest.engine.tool_call_dispatch,
crate::ToolCallDispatchMode::AfterResponse
);
let lower = WorkerManifestConfig {
engine: EngineManifestConfig {
tool_call_dispatch: Some(crate::ToolCallDispatchMode::OnToolCallComplete),
..Default::default()
},
..Default::default()
};
let upper = WorkerManifestConfig::default();
assert_eq!(
lower.merge(upper).engine.tool_call_dispatch,
Some(crate::ToolCallDispatchMode::OnToolCallComplete)
);
let mut configured = minimal_valid();
configured.engine.tool_call_dispatch =
Some(crate::ToolCallDispatchMode::OnToolCallComplete);
let manifest: WorkerManifest = configured.try_into().unwrap();
let snapshot = serde_json::to_value(&manifest).unwrap();
let restored: WorkerManifest = serde_json::from_value(snapshot).unwrap();
assert_eq!(
restored.engine.tool_call_dispatch,
crate::ToolCallDispatchMode::OnToolCallComplete
);
}
#[test]
fn resolve_minimal_succeeds() {
let manifest: WorkerManifest = minimal_valid().try_into().unwrap();
crates/manifest/src/lib.rs
+17
View File
@@ -587,6 +587,20 @@ pub struct WorkerMeta {
pub name: String,
}
/// When the engine may begin executing provider-confirmed tool calls.
///
/// `OnToolCallComplete` advances only the start time: calls may still overlap,
/// so completion order and externally allocated identifiers are not ordered.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ToolCallDispatchMode {
/// Preserve the historical behavior: wait for the whole response stream.
#[default]
AfterResponse,
/// Start each call once the provider emits a trustworthy per-call completion.
OnToolCallComplete,
}
/// Worker-level configuration embedded in the manifest.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EngineManifest {
@@ -614,6 +628,9 @@ pub struct EngineManifest {
pub stop_sequences: Vec<String>,
#[serde(default)]
pub reasoning: Option<ReasoningControl>,
/// Tool-call start timing. Omitted manifests preserve `after_response`.
#[serde(default)]
pub tool_call_dispatch: ToolCallDispatchMode,
/// Byte-size caps applied to tool `content` before it reaches the
/// conversation history. The section is optional in TOML — when
/// omitted, `ToolOutputLimits::default()` (64 KiB default cap, no
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
+52 -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};
@@ -49,10 +50,19 @@ pub enum LoggedItem {
role: LoggedRole,
content: Vec<LoggedContentPart>,
},
AssistantResponseBoundary {
response_id: String,
},
ToolCall {
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,
@@ -115,15 +125,24 @@ impl From<&Item> for LoggedItem {
role: (*role).into(),
content: content.iter().map(LoggedContentPart::from).collect(),
},
Item::AssistantResponseBoundary { response_id } => Self::AssistantResponseBoundary {
response_id: response_id.clone(),
},
Item::ToolCall {
call_id,
name,
arguments,
call_index,
execution_id,
status,
..
} => Self::ToolCall {
call_id: call_id.clone(),
name: name.clone(),
arguments: arguments.clone(),
call_index: *call_index,
execution_id: execution_id.clone(),
status: *status,
},
Item::ToolResult {
call_id,
@@ -172,16 +191,24 @@ impl From<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,
arguments,
call_index,
execution_id,
status,
} => Item::ToolCall {
id: None,
call_id,
name,
arguments,
status: None,
call_index,
execution_id,
status,
},
LoggedItem::ToolResult {
call_id,
@@ -325,9 +352,23 @@ mod tests {
}
}
#[test]
fn assistant_response_boundary_round_trips_for_provider_projection() {
let original = Item::assistant_response_boundary("response-7");
let logged: LoggedItem = (&original).into();
let json = serde_json::to_string(&logged).unwrap();
assert!(json.contains(r#""kind":"assistant_response_boundary""#));
assert_eq!(
Item::from(serde_json::from_str::<LoggedItem>(&json).unwrap()),
original
);
}
#[test]
fn round_trip_tool_call() {
let original = Item::tool_call("call_42", "get_weather", r#"{"city":"Tokyo"}"#);
let original = Item::tool_call("call_42", "get_weather", r#"{"city":"Tokyo"}"#)
.with_tool_execution_metadata(3, Some("batch:call_42".to_string()))
.with_status(ItemStatus::InProgress);
let logged: LoggedItem = (&original).into();
let json = serde_json::to_string(&logged).unwrap();
let parsed: LoggedItem = serde_json::from_str(&json).unwrap();
@@ -336,11 +377,17 @@ mod tests {
call_id,
name,
arguments,
call_index,
execution_id,
status,
..
} => {
assert_eq!(call_id, "call_42");
assert_eq!(name, "get_weather");
assert_eq!(arguments, r#"{"city":"Tokyo"}"#);
assert_eq!(call_index, Some(3));
assert_eq!(execution_id.as_deref(), Some("batch:call_42"));
assert_eq!(status, Some(ItemStatus::InProgress));
}
other => panic!("unexpected variant: {other:?}"),
}
crates/session-store/src/public_snapshot.rs
+3 -2
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(),
@@ -385,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/feature/builtin/flow_transition.rs
+7 -1
View File
@@ -636,7 +636,7 @@ mod tests {
use agen::llm_client::error::ClientError;
use agen::llm_client::event::{
BlockDelta, BlockMetadata, BlockStart, BlockStop, BlockType, DeltaContent,
Event as LlmEvent, StopReason,
Event as LlmEvent, ResponseStatus, StatusEvent, StopReason,
};
use agen::llm_client::types::Request;
use agen::tool::ToolExecutionContext;
@@ -906,6 +906,9 @@ mod tests {
reasoning: None,
stop_reason: Some(StopReason::ToolUse),
})),
Ok(LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
})),
])))
} else {
Ok(Box::pin(stream::iter(vec![
@@ -924,6 +927,9 @@ mod tests {
reasoning: None,
stop_reason: Some(StopReason::EndTurn),
})),
Ok(LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
})),
])))
}
}
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/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/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,
},
),
});
crates/worker/src/spawn/tool.rs
+7 -1
View File
@@ -965,6 +965,7 @@ fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfi
stop_sequences: (!manifest.engine.stop_sequences.is_empty())
.then_some(manifest.engine.stop_sequences.clone()),
reasoning: manifest.engine.reasoning.clone(),
tool_call_dispatch: Some(manifest.engine.tool_call_dispatch),
tool_output: ToolOutputLimitsPartial {
default_max_bytes: Some(manifest.engine.tool_output.default_max_bytes),
per_tool: manifest.engine.tool_output.per_tool.clone(),
@@ -2226,7 +2227,8 @@ max_tokens = 3333
Some("reviewer"),
&[("reviewer", "reviewer.toml", REVIEWER_PROFILE)],
);
let parent = parent_manifest(&parent_root, Some(&parent_root.join("deny")));
let mut parent = parent_manifest(&parent_root, Some(&parent_root.join("deny")));
parent.engine.tool_call_dispatch = manifest::ToolCallDispatchMode::OnToolCallComplete;
let scope = vec![abs_rule(&delegated, Permission::Read)];
let profile_config = child_config_from_profile(
@@ -2248,6 +2250,10 @@ max_tokens = 3333
SpawnProfileSelector::Inherit,
);
assert_eq!(
inherit_config.engine.tool_call_dispatch,
Some(manifest::ToolCallDispatchMode::OnToolCallComplete)
);
for config in [profile_config, inherit_config] {
assert_eq!(config.scope.allow, scope);
assert!(config.scope.deny.is_empty());
crates/worker/src/worker.rs
+24 -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
@@ -7138,6 +7145,12 @@ pub fn apply_worker_manifest<C: LlmClient + 'static, A: Send + Sync>(
) {
worker.set_request_config(request_config_from_engine_manifest(wm));
worker.set_max_turns(wm.max_turns.map(|n| n.get()));
worker.set_tool_call_dispatch_mode(match wm.tool_call_dispatch {
manifest::ToolCallDispatchMode::AfterResponse => agen::ToolCallDispatchMode::AfterResponse,
manifest::ToolCallDispatchMode::OnToolCallComplete => {
agen::ToolCallDispatchMode::OnToolCallComplete
}
});
// Worker owns the lifecycle strategy for already-started tool operations.
// The provider must first accept cooperative cancellation, then confirm a
// terminal result before this bounded deadline; Agen handles only the
@@ -7551,6 +7564,7 @@ fn write_overview_header(items: &[Item], out: &mut String) {
for item in items {
match item {
Item::Message { .. } => messages += 1,
Item::AssistantResponseBoundary { .. } => {}
Item::ToolCall { .. } => tool_calls += 1,
Item::ToolResult { .. } => tool_results += 1,
Item::Reasoning { .. } => reasoning += 1,
@@ -8721,6 +8735,7 @@ mod build_summary_prompt_tests {
top_k: Some(40),
stop_sequences: vec!["\n\n".into(), "</stop>".into()],
reasoning: None,
tool_call_dispatch: manifest::ToolCallDispatchMode::OnToolCallComplete,
tool_output: manifest::ToolOutputLimits::default(),
file_upload: manifest::FileUploadLimits::default(),
};
@@ -8732,6 +8747,13 @@ mod build_summary_prompt_tests {
assert_eq!(config.top_p, Some(0.9));
assert_eq!(config.top_k, Some(40));
assert_eq!(config.stop_sequences, vec!["\n\n", "</stop>"]);
let mut engine = Engine::new(NoopClient);
apply_worker_manifest(&mut engine, &manifest);
assert_eq!(
engine.tool_call_dispatch_mode(),
agen::ToolCallDispatchMode::OnToolCallComplete
);
}
#[test]
docs/manifest.toml
+9
View File
@@ -121,6 +121,15 @@ ref = "anthropic/claude-sonnet-4-6"
# reasoning = "medium"
# reasoning = -1
# 任意。デフォルト: "after_response"。
# "on_tool_call_complete" は provider が個別 ToolCall の完成を保証する場合だけ、
# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider、または
# 応答全体を見て実行前 Pause を判断する turn-end hook がある場合は、応答完了まで
# 警告付きで明示的に遅延する。これは直列実行モードではなく、先行 call の完了を
# 待たず後続 call も開始するため、完了順・DB 反映順・採番順は保証しない。
# 早期開始後に stream が切断された場合は自動 continuation せず、安全に terminalize して停止する。
# tool_call_dispatch = "on_tool_call_complete"
# 任意。tool 実行 content の byte 長キャップ。
# セクション省略時は default_max_bytes = 64 * 1024、per_tool 空。
# [worker.tool_output]