feat: add opt-in early tool call dispatch
This commit is contained in:
+727
-24
@@ -1,7 +1,8 @@
|
|||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet, VecDeque};
|
||||||
use std::{future::Future, marker::PhantomData, pin::Pin, sync::Arc, time::Instant};
|
use std::{future::Future, marker::PhantomData, pin::Pin, sync::Arc, time::Instant};
|
||||||
|
|
||||||
use futures::StreamExt;
|
use futures::StreamExt;
|
||||||
|
use futures::stream::FuturesUnordered;
|
||||||
use serde_json::{Value, json};
|
use serde_json::{Value, json};
|
||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
use tokio::time::Instant as TokioInstant;
|
use tokio::time::Instant as TokioInstant;
|
||||||
@@ -25,7 +26,7 @@ use crate::{
|
|||||||
llm_client::{
|
llm_client::{
|
||||||
ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream,
|
ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream,
|
||||||
ToolDefinition, error::is_retryable, event::Event, retry::RetryPolicy,
|
ToolDefinition, error::is_retryable, event::Event, retry::RetryPolicy,
|
||||||
transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT, types::parse_tool_arguments,
|
transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT,
|
||||||
},
|
},
|
||||||
state::{EngineState, Locked, Mutable},
|
state::{EngineState, Locked, Mutable},
|
||||||
timeline::event::{ErrorEvent, StatusEvent, UsageEvent},
|
timeline::event::{ErrorEvent, StatusEvent, UsageEvent},
|
||||||
@@ -78,11 +79,20 @@ pub enum ToolRegistryError {
|
|||||||
DuplicateName(String),
|
DuplicateName(String),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// When provider-confirmed tool calls may begin execution.
|
||||||
|
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||||
|
pub enum ToolCallDispatchMode {
|
||||||
|
/// Wait for the complete assistant response before starting any tools.
|
||||||
|
#[default]
|
||||||
|
AfterResponse,
|
||||||
|
/// Start calls at trustworthy provider per-call completion boundaries.
|
||||||
|
OnToolCallComplete,
|
||||||
|
}
|
||||||
|
|
||||||
/// Engine configuration
|
/// Engine configuration
|
||||||
#[derive(Debug, Clone, Default)]
|
#[derive(Debug, Clone, Default)]
|
||||||
pub struct EngineConfig {
|
pub struct EngineConfig {
|
||||||
// Reserved for future extensions (currently empty)
|
pub tool_call_dispatch: ToolCallDispatchMode,
|
||||||
_private: (),
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Project terminal tool outputs into the assistant's original ToolCall order.
|
/// Project terminal tool outputs into the assistant's original ToolCall order.
|
||||||
@@ -196,6 +206,7 @@ pub struct EngineRunOutput<C: LlmClient, A: Send + Sync = ()> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Internal: tool execution result
|
/// Internal: tool execution result
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
enum ToolExecutionResult {
|
enum ToolExecutionResult {
|
||||||
Completed,
|
Completed,
|
||||||
Paused,
|
Paused,
|
||||||
@@ -207,6 +218,56 @@ struct ToolExecutionAttempt {
|
|||||||
terminal: bool,
|
terminal: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type ToolTerminalTask = Pin<Box<dyn Future<Output = (String, ToolResult)> + Send>>;
|
||||||
|
type ToolCallExecutionInfo = (
|
||||||
|
ToolCall,
|
||||||
|
crate::tool::ToolMeta,
|
||||||
|
Arc<dyn crate::tool::Tool>,
|
||||||
|
ToolExecutionContext,
|
||||||
|
);
|
||||||
|
|
||||||
|
/// One response-scoped execution batch whose calls are admitted as their
|
||||||
|
/// provider-confirmed blocks complete.
|
||||||
|
struct EarlyToolExecutionBatch {
|
||||||
|
batch_id: String,
|
||||||
|
next_call_index: usize,
|
||||||
|
calls: Vec<ToolCall>,
|
||||||
|
seen_call_ids: HashSet<String>,
|
||||||
|
call_info_map: HashMap<String, ToolCallExecutionInfo>,
|
||||||
|
started_calls: Vec<(String, String)>,
|
||||||
|
attempt_fence: ToolExecutionAttemptFence,
|
||||||
|
futures: FuturesUnordered<ToolTerminalTask>,
|
||||||
|
execution_handles: HashMap<String, ToolExecutionHandle>,
|
||||||
|
synthetic_results: VecDeque<ToolResult>,
|
||||||
|
terminal_call_ids: HashSet<String>,
|
||||||
|
dispatched_any: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
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(),
|
||||||
|
started_calls: Vec::new(),
|
||||||
|
attempt_fence: ToolExecutionAttemptFence::default(),
|
||||||
|
futures: FuturesUnordered::new(),
|
||||||
|
execution_handles: HashMap::new(),
|
||||||
|
synthetic_results: VecDeque::new(),
|
||||||
|
terminal_call_ids: HashSet::new(),
|
||||||
|
dispatched_any: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
enum EarlyToolAdmission {
|
||||||
|
Continue,
|
||||||
|
Pause,
|
||||||
|
}
|
||||||
|
|
||||||
/// Per-batch compare-and-set fence for terminal ToolResult commits.
|
/// Per-batch compare-and-set fence for terminal ToolResult commits.
|
||||||
///
|
///
|
||||||
/// A completion may commit only when its attempt id still matches the active
|
/// A completion may commit only when its attempt id still matches the active
|
||||||
@@ -312,6 +373,11 @@ enum StreamCompletion {
|
|||||||
Interrupted { reason: String },
|
Interrupted { reason: String },
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct StreamResponseOutput {
|
||||||
|
completion: StreamCompletion,
|
||||||
|
early_tools: Option<EarlyToolExecutionBatch>,
|
||||||
|
}
|
||||||
|
|
||||||
pub struct Engine<C: LlmClient, S: EngineState = Mutable, A: Send + Sync = ()> {
|
pub struct Engine<C: LlmClient, S: EngineState = Mutable, A: Send + Sync = ()> {
|
||||||
/// LLM client
|
/// LLM client
|
||||||
client: C,
|
client: C,
|
||||||
@@ -362,6 +428,8 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable, A: Send + Sync = ()> {
|
|||||||
tool_execution_batch_count: usize,
|
tool_execution_batch_count: usize,
|
||||||
/// Maximum number of AgentTurns (None = unlimited)
|
/// Maximum number of AgentTurns (None = unlimited)
|
||||||
max_turns: Option<u32>,
|
max_turns: Option<u32>,
|
||||||
|
/// Selected tool-call execution start boundary.
|
||||||
|
tool_call_dispatch_mode: ToolCallDispatchMode,
|
||||||
/// Caller-selected policy for interrupting started provider operations.
|
/// Caller-selected policy for interrupting started provider operations.
|
||||||
tool_execution_policy: ToolExecutionPolicy,
|
tool_execution_policy: ToolExecutionPolicy,
|
||||||
/// AgentTurn-start callbacks (1:1 with LlmCall today)
|
/// AgentTurn-start callbacks (1:1 with LlmCall today)
|
||||||
@@ -1184,7 +1252,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
} = item
|
} = item
|
||||||
{
|
{
|
||||||
if !answered_call_ids.contains(call_id) {
|
if !answered_call_ids.contains(call_id) {
|
||||||
let input = parse_tool_arguments(arguments);
|
let input = serde_json::from_str::<Value>(arguments)
|
||||||
|
.ok()
|
||||||
|
.filter(Value::is_object)
|
||||||
|
.unwrap_or(Value::Null);
|
||||||
pending_calls.push(ToolCall {
|
pending_calls.push(ToolCall {
|
||||||
id: call_id.clone(),
|
id: call_id.clone(),
|
||||||
name: name.clone(),
|
name: name.clone(),
|
||||||
@@ -1201,6 +1272,342 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn admit_early_tool_calls(
|
||||||
|
&mut self,
|
||||||
|
history: &History<A>,
|
||||||
|
batch: &mut EarlyToolExecutionBatch,
|
||||||
|
tool_calls: Vec<ToolCall>,
|
||||||
|
) -> Result<EarlyToolAdmission, EngineError> {
|
||||||
|
for mut tool_call in tool_calls {
|
||||||
|
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());
|
||||||
|
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() {
|
||||||
|
batch.synthetic_results.push_back(ToolResult::error(
|
||||||
|
&tool_call.id,
|
||||||
|
"Tool arguments must be one complete JSON object",
|
||||||
|
));
|
||||||
|
batch.dispatched_any = true;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) else {
|
||||||
|
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}")),
|
||||||
|
)
|
||||||
|
}));
|
||||||
|
batch.dispatched_any = true;
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
|
||||||
|
let invocation = self.interceptor_invocation(
|
||||||
|
InterceptorPhase::PreToolCall,
|
||||||
|
Some(self.turn_count),
|
||||||
|
Some(InterceptorCallId::Tool(expected_tool_use_id.clone())),
|
||||||
|
call_index,
|
||||||
|
);
|
||||||
|
let mut info = ToolCallInfo {
|
||||||
|
invocation,
|
||||||
|
history: history.entries(),
|
||||||
|
call: tool_call.clone(),
|
||||||
|
meta,
|
||||||
|
tool,
|
||||||
|
context,
|
||||||
|
};
|
||||||
|
let action = self
|
||||||
|
.interceptor
|
||||||
|
.pre_tool_call(&mut info)
|
||||||
|
.await
|
||||||
|
.map_err(|error| {
|
||||||
|
EngineError::from(InterceptorFailure::new(
|
||||||
|
InterceptorPhase::PreToolCall,
|
||||||
|
error,
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
if info.call.id != expected_tool_use_id {
|
||||||
|
return Err(InterceptorFailure::new(
|
||||||
|
InterceptorPhase::PreToolCall,
|
||||||
|
InterceptorError::new(
|
||||||
|
InterceptorErrorCategory::ContractViolation,
|
||||||
|
"pre-tool interceptor changed immutable tool call identity",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.into());
|
||||||
|
}
|
||||||
|
tool_call = info.call;
|
||||||
|
let mut context = info.context;
|
||||||
|
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(),
|
||||||
|
),
|
||||||
|
);
|
||||||
|
|
||||||
|
match action {
|
||||||
|
PreToolAction::Continue => {
|
||||||
|
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 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);
|
||||||
|
batch.execution_handles.insert(call_id.clone(), handle);
|
||||||
|
batch.futures.push(Box::pin(async move {
|
||||||
|
let result = match terminal.await {
|
||||||
|
ToolExecutionTerminal::Confirmed(Ok(output)) => {
|
||||||
|
ToolResult::from_output(&call_id, output)
|
||||||
|
}
|
||||||
|
ToolExecutionTerminal::Confirmed(Err(ToolError::Cancelled(output))) => {
|
||||||
|
ToolResult::from_output_with_disposition(
|
||||||
|
&call_id,
|
||||||
|
output,
|
||||||
|
ToolResultDisposition::Cancelled,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
ToolExecutionTerminal::Confirmed(Err(ToolError::Interrupted(
|
||||||
|
output,
|
||||||
|
))) => ToolResult::from_output_with_disposition(
|
||||||
|
&call_id,
|
||||||
|
output,
|
||||||
|
ToolResultDisposition::Interrupted,
|
||||||
|
),
|
||||||
|
ToolExecutionTerminal::Confirmed(Err(error)) => {
|
||||||
|
ToolResult::error(&call_id, error.to_string())
|
||||||
|
}
|
||||||
|
ToolExecutionTerminal::OutcomeUnknown => {
|
||||||
|
ToolResult::outcome_unknown(&call_id)
|
||||||
|
}
|
||||||
|
};
|
||||||
|
(attempt_id, result)
|
||||||
|
}));
|
||||||
|
batch.dispatched_any = true;
|
||||||
|
}
|
||||||
|
PreToolAction::SyntheticResult(result) => {
|
||||||
|
if result.tool_use_id != expected_tool_use_id {
|
||||||
|
return Err(InterceptorFailure::new(
|
||||||
|
InterceptorPhase::PreToolCall,
|
||||||
|
InterceptorError::new(
|
||||||
|
InterceptorErrorCategory::ContractViolation,
|
||||||
|
"synthetic tool result changed immutable tool call identity",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.into());
|
||||||
|
}
|
||||||
|
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),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(EarlyToolAdmission::Continue)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn cancel_and_terminalize_early_batch(
|
||||||
|
&mut self,
|
||||||
|
history: &mut History<A>,
|
||||||
|
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||||
|
batch: &mut EarlyToolExecutionBatch,
|
||||||
|
) -> Option<EngineError> {
|
||||||
|
use futures::StreamExt as _;
|
||||||
|
|
||||||
|
let cancellation_request_deadline =
|
||||||
|
TokioInstant::now() + self.tool_execution_policy.cancellation_request_timeout;
|
||||||
|
let cancellation_requests = batch
|
||||||
|
.execution_handles
|
||||||
|
.iter()
|
||||||
|
.filter(|(call_id, _)| !batch.terminal_call_ids.contains(*call_id))
|
||||||
|
.map(|(call_id, handle)| {
|
||||||
|
let call_id = call_id.clone();
|
||||||
|
let handle = handle.clone();
|
||||||
|
async move {
|
||||||
|
(
|
||||||
|
call_id,
|
||||||
|
handle.cancel_before(cancellation_request_deadline).await,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let cancellation_requests: FuturesUnordered<_> = cancellation_requests.collect();
|
||||||
|
for (call_id, result) in cancellation_requests.collect::<Vec<_>>().await {
|
||||||
|
if let Err(error) = result {
|
||||||
|
warn!(%call_id, error = %error, "Tool cooperative cancellation request failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut batch_error = None;
|
||||||
|
let deadline =
|
||||||
|
TokioInstant::now() + self.tool_execution_policy.terminal_confirmation_timeout;
|
||||||
|
while !batch.futures.is_empty() {
|
||||||
|
tokio::select! {
|
||||||
|
biased;
|
||||||
|
result = batch.futures.next() => {
|
||||||
|
let (attempt_id, result) = result.expect("non-empty FuturesUnordered");
|
||||||
|
if let Err(error) = self.finalize_and_commit_tool_result(
|
||||||
|
history,
|
||||||
|
annotate,
|
||||||
|
result,
|
||||||
|
Some(&attempt_id),
|
||||||
|
&batch.call_info_map,
|
||||||
|
&mut batch.attempt_fence,
|
||||||
|
&mut batch.terminal_call_ids,
|
||||||
|
).await && batch_error.is_none() {
|
||||||
|
batch_error = Some(error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ = tokio::time::sleep_until(deadline) => break,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for (call_id, attempt_id) in &batch.started_calls {
|
||||||
|
if !batch.attempt_fence.is_terminal(call_id) {
|
||||||
|
if let Some(handle) = batch.execution_handles.get(call_id) {
|
||||||
|
handle.force_close();
|
||||||
|
}
|
||||||
|
if let Err(error) = self
|
||||||
|
.finalize_and_commit_tool_result(
|
||||||
|
history,
|
||||||
|
annotate,
|
||||||
|
ToolResult::outcome_unknown(call_id),
|
||||||
|
Some(attempt_id),
|
||||||
|
&batch.call_info_map,
|
||||||
|
&mut batch.attempt_fence,
|
||||||
|
&mut batch.terminal_call_ids,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
&& batch_error.is_none()
|
||||||
|
{
|
||||||
|
batch_error = Some(error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
batch_error
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn drain_early_tool_batch(
|
||||||
|
&mut self,
|
||||||
|
history: &mut History<A>,
|
||||||
|
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||||
|
mut batch: EarlyToolExecutionBatch,
|
||||||
|
force_cancel: bool,
|
||||||
|
initial_pause_requested: bool,
|
||||||
|
) -> Result<ToolExecutionResult, EngineError> {
|
||||||
|
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
|
||||||
|
{
|
||||||
|
let cleanup_error = self
|
||||||
|
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
|
||||||
|
.await;
|
||||||
|
return Err(cleanup_error.unwrap_or(error));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if force_cancel {
|
||||||
|
if let Some(error) = self
|
||||||
|
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
return Ok(ToolExecutionResult::Completed);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut pause_requested = initial_pause_requested;
|
||||||
|
let mut pause_deadline = initial_pause_requested
|
||||||
|
.then(|| TokioInstant::now() + self.tool_execution_policy.pause_safe_boundary_timeout);
|
||||||
|
while !batch.futures.is_empty() {
|
||||||
|
tokio::select! {
|
||||||
|
biased;
|
||||||
|
result = batch.futures.next() => {
|
||||||
|
let (attempt_id, result) = result.expect("non-empty FuturesUnordered");
|
||||||
|
if let Err(error) = self.finalize_and_commit_tool_result(
|
||||||
|
history,
|
||||||
|
annotate,
|
||||||
|
result,
|
||||||
|
Some(&attempt_id),
|
||||||
|
&batch.call_info_map,
|
||||||
|
&mut batch.attempt_fence,
|
||||||
|
&mut batch.terminal_call_ids,
|
||||||
|
).await {
|
||||||
|
let cleanup_error = self
|
||||||
|
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
|
||||||
|
.await;
|
||||||
|
return Err(cleanup_error.unwrap_or(error));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
pause = self.pause_rx.recv(), if !pause_requested => {
|
||||||
|
if pause.is_some() {
|
||||||
|
pause_requested = true;
|
||||||
|
pause_deadline = Some(
|
||||||
|
TokioInstant::now()
|
||||||
|
+ self.tool_execution_policy.pause_safe_boundary_timeout,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ = tokio::time::sleep_until(pause_deadline.unwrap_or_else(TokioInstant::now)), if pause_deadline.is_some() => {
|
||||||
|
if let Some(error) = self
|
||||||
|
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
|
||||||
|
.await {
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
return Ok(ToolExecutionResult::Paused);
|
||||||
|
}
|
||||||
|
cancel = self.cancel_rx.recv() => {
|
||||||
|
if cancel.is_some() {
|
||||||
|
info!("Tool execution cancellation requested");
|
||||||
|
}
|
||||||
|
if let Some(error) = self
|
||||||
|
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
|
||||||
|
.await {
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
return Err(EngineError::Cancelled);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(if pause_requested {
|
||||||
|
ToolExecutionResult::Paused
|
||||||
|
} else {
|
||||||
|
ToolExecutionResult::Completed
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
/// Execute tools in parallel
|
/// Execute tools in parallel
|
||||||
///
|
///
|
||||||
/// After running pre_tool_call hooks on all tools,
|
/// After running pre_tool_call hooks on all tools,
|
||||||
@@ -1231,6 +1638,13 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
for (call_index, mut tool_call) in tool_calls.into_iter().enumerate() {
|
for (call_index, mut tool_call) in tool_calls.into_iter().enumerate() {
|
||||||
let expected_tool_use_id = tool_call.id.clone();
|
let expected_tool_use_id = tool_call.id.clone();
|
||||||
let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index);
|
let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index);
|
||||||
|
if !tool_call.input.is_object() {
|
||||||
|
synthetic_results.push(ToolResult::error(
|
||||||
|
&tool_call.id,
|
||||||
|
"Tool arguments must be one complete JSON object",
|
||||||
|
));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
if let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) {
|
if let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) {
|
||||||
let invocation = self.interceptor_invocation(
|
let invocation = self.interceptor_invocation(
|
||||||
InterceptorPhase::PreToolCall,
|
InterceptorPhase::PreToolCall,
|
||||||
@@ -1934,8 +2348,30 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
self.request_trace_payload(&request),
|
self.request_trace_payload(&request),
|
||||||
);
|
);
|
||||||
let request = self.attach_transport_trace(request, current_turn, current_llm_call);
|
let request = self.attach_transport_trace(request, current_turn, current_llm_call);
|
||||||
let stream_outcome = self
|
let early_dispatch = self.tool_call_dispatch_mode
|
||||||
.stream_response(request, current_turn, current_llm_call)
|
== ToolCallDispatchMode::OnToolCallComplete
|
||||||
|
&& self.client.tool_call_completion_support()
|
||||||
|
== crate::llm_client::ToolCallCompletionSupport::PerBlock;
|
||||||
|
if self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete
|
||||||
|
&& !early_dispatch
|
||||||
|
{
|
||||||
|
self.emit_warning(
|
||||||
|
"tool-call early dispatch deferred until response completion: provider does not expose a trustworthy per-call completion boundary",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let assistant_start = history.len();
|
||||||
|
let StreamResponseOutput {
|
||||||
|
completion: stream_outcome,
|
||||||
|
mut early_tools,
|
||||||
|
} = self
|
||||||
|
.stream_response(
|
||||||
|
request,
|
||||||
|
current_turn,
|
||||||
|
current_llm_call,
|
||||||
|
history,
|
||||||
|
annotate,
|
||||||
|
early_dispatch,
|
||||||
|
)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
for cb in &self.llm_call_end_cbs {
|
for cb in &self.llm_call_end_cbs {
|
||||||
@@ -1944,6 +2380,47 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
self.llm_call_count += 1;
|
self.llm_call_count += 1;
|
||||||
|
|
||||||
if let StreamCompletion::Interrupted { reason } = stream_outcome {
|
if let StreamCompletion::Interrupted { reason } = stream_outcome {
|
||||||
|
if early_tools
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|batch| batch.dispatched_any)
|
||||||
|
{
|
||||||
|
// Once a side effect may have started, automatic stream
|
||||||
|
// continuation would allow the model to regenerate the same
|
||||||
|
// operation under a fresh call id. Close every owned attempt
|
||||||
|
// and stop this run instead.
|
||||||
|
self.timeline.abort_current_block();
|
||||||
|
self.timeline.flush_usage();
|
||||||
|
let reasoning_items = self.thinking_block_collector.take_collected();
|
||||||
|
let text_blocks = self.text_block_collector.take_collected();
|
||||||
|
let _partial_tool_calls = self.tool_call_collector.take_collected();
|
||||||
|
let assistant_items =
|
||||||
|
self.build_assistant_items(&reasoning_items, &text_blocks, &[]);
|
||||||
|
if !assistant_items.is_empty()
|
||||||
|
&& let Err(error) =
|
||||||
|
self.append_history_items(history, assistant_items, annotate)
|
||||||
|
{
|
||||||
|
if let Some(batch) = early_tools.take()
|
||||||
|
&& let Err(cleanup_error) = self
|
||||||
|
.drain_early_tool_batch(history, annotate, batch, true, false)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return Err(cleanup_error);
|
||||||
|
}
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
if let Some(batch) = early_tools.take() {
|
||||||
|
self.drain_early_tool_batch(history, annotate, batch, true, false)
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
return Err(EngineError::Client(ClientError::Api {
|
||||||
|
status: None,
|
||||||
|
code: Some("early_tool_stream_interrupted".to_string()),
|
||||||
|
message: format!(
|
||||||
|
"LLM stream interrupted after early tool execution started: {reason}"
|
||||||
|
),
|
||||||
|
retry_after: None,
|
||||||
|
}));
|
||||||
|
}
|
||||||
stream_continuations += 1;
|
stream_continuations += 1;
|
||||||
if stream_continuations > MAX_STREAM_CONTINUATIONS {
|
if stream_continuations > MAX_STREAM_CONTINUATIONS {
|
||||||
return Err(EngineError::Client(ClientError::Api {
|
return Err(EngineError::Client(ClientError::Api {
|
||||||
@@ -1989,11 +2466,61 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
// `append_history_items` so observers see each item as it lands.
|
// `append_history_items` so observers see each item as it lands.
|
||||||
let reasoning_items = self.thinking_block_collector.take_collected();
|
let reasoning_items = self.thinking_block_collector.take_collected();
|
||||||
let text_blocks = self.text_block_collector.take_collected();
|
let text_blocks = self.text_block_collector.take_collected();
|
||||||
let tool_calls = self.tool_call_collector.take_collected();
|
let remaining_tool_calls = self.tool_call_collector.take_collected();
|
||||||
let assistant_items =
|
let assistant_items =
|
||||||
self.build_assistant_items(&reasoning_items, &text_blocks, &tool_calls);
|
self.build_assistant_items(&reasoning_items, &text_blocks, &remaining_tool_calls);
|
||||||
let assistant_start = history.len();
|
if let Err(error) = self.append_history_items(history, assistant_items, annotate) {
|
||||||
self.append_history_items(history, assistant_items, annotate)?;
|
if let Some(batch) = early_tools.take()
|
||||||
|
&& let Err(cleanup_error) = self
|
||||||
|
.drain_early_tool_batch(history, annotate, batch, true, false)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return Err(cleanup_error);
|
||||||
|
}
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
|
||||||
|
if early_dispatch && !remaining_tool_calls.is_empty() {
|
||||||
|
if early_tools.is_none() {
|
||||||
|
let batch_id = format!("tool-batch-{}", self.tool_execution_batch_count);
|
||||||
|
self.tool_execution_batch_count += 1;
|
||||||
|
early_tools = Some(EarlyToolExecutionBatch::new(batch_id));
|
||||||
|
}
|
||||||
|
let admission = self
|
||||||
|
.admit_early_tool_calls(
|
||||||
|
history,
|
||||||
|
early_tools.as_mut().expect("batch initialized"),
|
||||||
|
remaining_tool_calls.clone(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let admission = match admission {
|
||||||
|
Ok(admission) => admission,
|
||||||
|
Err(error) => {
|
||||||
|
if let Some(batch) = early_tools.take()
|
||||||
|
&& let Err(cleanup_error) = self
|
||||||
|
.drain_early_tool_batch(history, annotate, batch, true, false)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return Err(cleanup_error);
|
||||||
|
}
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if admission == EarlyToolAdmission::Pause {
|
||||||
|
if let Some(batch) = early_tools.take() {
|
||||||
|
self.drain_early_tool_batch(history, annotate, batch, false, true)
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
return Ok(EngineResult::Paused);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut tool_calls = early_tools
|
||||||
|
.as_ref()
|
||||||
|
.map(|batch| batch.calls.clone())
|
||||||
|
.unwrap_or_default();
|
||||||
|
if !early_dispatch {
|
||||||
|
tool_calls = remaining_tool_calls;
|
||||||
|
}
|
||||||
|
|
||||||
let assistant_invocation = self.interceptor_invocation(
|
let assistant_invocation = self.interceptor_invocation(
|
||||||
InterceptorPhase::AssistantTurnEnd,
|
InterceptorPhase::AssistantTurnEnd,
|
||||||
@@ -2011,29 +2538,64 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
history: history.entries(),
|
history: history.entries(),
|
||||||
tool_calls: &tool_calls,
|
tool_calls: &tool_calls,
|
||||||
})
|
})
|
||||||
.await
|
.await;
|
||||||
.map_err(|error| {
|
let assistant_turn_action = match assistant_turn_action {
|
||||||
EngineError::from(InterceptorFailure::new(
|
Ok(action) => action,
|
||||||
|
Err(error) => {
|
||||||
|
let error = EngineError::from(InterceptorFailure::new(
|
||||||
InterceptorPhase::AssistantTurnEnd,
|
InterceptorPhase::AssistantTurnEnd,
|
||||||
error,
|
error,
|
||||||
))
|
));
|
||||||
})?;
|
if let Some(batch) = early_tools.take()
|
||||||
|
&& let Err(cleanup_error) = self
|
||||||
|
.drain_early_tool_batch(history, annotate, batch, true, false)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return Err(cleanup_error);
|
||||||
|
}
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
};
|
||||||
match assistant_turn_action {
|
match assistant_turn_action {
|
||||||
TurnEndAction::Finish if tool_calls.is_empty() => {
|
TurnEndAction::Finish if tool_calls.is_empty() => {
|
||||||
return Ok(EngineResult::Finished);
|
return Ok(EngineResult::Finished);
|
||||||
}
|
}
|
||||||
TurnEndAction::Finish => {}
|
TurnEndAction::Finish => {}
|
||||||
TurnEndAction::ContinueWithMessages(additional) => {
|
TurnEndAction::ContinueWithMessages(additional) => {
|
||||||
self.append_history_items(history, additional, annotate)?;
|
if let Err(error) = self.append_history_items(history, additional, annotate) {
|
||||||
|
if let Some(batch) = early_tools.take()
|
||||||
|
&& let Err(cleanup_error) = self
|
||||||
|
.drain_early_tool_batch(history, annotate, batch, true, false)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return Err(cleanup_error);
|
||||||
|
}
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
if tool_calls.is_empty() {
|
if tool_calls.is_empty() {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
TurnEndAction::Pause => {
|
TurnEndAction::Pause => {
|
||||||
|
if let Some(batch) = early_tools.take() {
|
||||||
|
self.drain_early_tool_batch(history, annotate, batch, false, true)
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
return Ok(EngineResult::Paused);
|
return Ok(EngineResult::Paused);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if let Some(batch) = early_tools.take() {
|
||||||
|
if self
|
||||||
|
.drain_early_tool_batch(history, annotate, batch, false, false)
|
||||||
|
.await?
|
||||||
|
== ToolExecutionResult::Paused
|
||||||
|
{
|
||||||
|
return Ok(EngineResult::Paused);
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
if let Some(result) = self
|
if let Some(result) = self
|
||||||
.execute_and_commit_tools(history, annotate, tool_calls)
|
.execute_and_commit_tools(history, annotate, tool_calls)
|
||||||
.await?
|
.await?
|
||||||
@@ -2224,7 +2786,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
request: Request,
|
request: Request,
|
||||||
turn: usize,
|
turn: usize,
|
||||||
llm_call: usize,
|
llm_call: usize,
|
||||||
) -> Result<StreamCompletion, EngineError> {
|
history: &mut History<A>,
|
||||||
|
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
|
||||||
|
early_dispatch: bool,
|
||||||
|
) -> Result<StreamResponseOutput, EngineError> {
|
||||||
debug!(
|
debug!(
|
||||||
item_count = request.items.len(),
|
item_count = request.items.len(),
|
||||||
tool_count = request.tools.len(),
|
tool_count = request.tools.len(),
|
||||||
@@ -2233,6 +2798,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
);
|
);
|
||||||
|
|
||||||
let mut stream = self.open_stream_with_retry(request, turn, llm_call).await?;
|
let mut stream = self.open_stream_with_retry(request, turn, llm_call).await?;
|
||||||
|
let mut early_tools: Option<EarlyToolExecutionBatch> = None;
|
||||||
|
|
||||||
let mut event_count: usize = 0;
|
let mut event_count: usize = 0;
|
||||||
loop {
|
loop {
|
||||||
@@ -2254,8 +2820,11 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
Err(err) => {
|
Err(err) => {
|
||||||
// 部分情報でも発火しておく(料金会計用)
|
// 部分情報でも発火しておく(料金会計用)
|
||||||
self.timeline.flush_usage();
|
self.timeline.flush_usage();
|
||||||
return Ok(StreamCompletion::Interrupted {
|
return Ok(StreamResponseOutput {
|
||||||
reason: err.to_string(),
|
completion: StreamCompletion::Interrupted {
|
||||||
|
reason: err.to_string(),
|
||||||
|
},
|
||||||
|
early_tools,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -2269,9 +2838,116 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
}
|
}
|
||||||
self.emit_stream_event(turn, llm_call, &event);
|
self.emit_stream_event(turn, llm_call, &event);
|
||||||
self.timeline.dispatch(&event);
|
self.timeline.dispatch(&event);
|
||||||
|
if early_dispatch {
|
||||||
|
let collected = self.tool_call_collector.take_collected();
|
||||||
|
if !collected.is_empty() {
|
||||||
|
if early_tools.is_none() {
|
||||||
|
let batch_id = format!(
|
||||||
|
"tool-batch-{}",
|
||||||
|
self.tool_execution_batch_count
|
||||||
|
);
|
||||||
|
self.tool_execution_batch_count += 1;
|
||||||
|
early_tools =
|
||||||
|
Some(EarlyToolExecutionBatch::new(batch_id));
|
||||||
|
}
|
||||||
|
let fresh_calls = {
|
||||||
|
let batch = early_tools.as_ref().expect("batch initialized");
|
||||||
|
collected
|
||||||
|
.into_iter()
|
||||||
|
.filter(|call| {
|
||||||
|
!batch.seen_call_ids.contains(&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.
|
||||||
|
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,
|
||||||
|
assistant_items,
|
||||||
|
annotate,
|
||||||
|
) {
|
||||||
|
if let Some(mut batch) = early_tools.take() {
|
||||||
|
let _ = self
|
||||||
|
.cancel_and_terminalize_early_batch(
|
||||||
|
history,
|
||||||
|
annotate,
|
||||||
|
&mut batch,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
let admission = {
|
||||||
|
let batch =
|
||||||
|
early_tools.as_mut().expect("batch initialized");
|
||||||
|
self.admit_early_tool_calls(
|
||||||
|
history,
|
||||||
|
batch,
|
||||||
|
fresh_calls,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
};
|
||||||
|
match admission {
|
||||||
|
Ok(EarlyToolAdmission::Continue) => {}
|
||||||
|
Ok(EarlyToolAdmission::Pause) => {
|
||||||
|
self.timeline.abort_current_block();
|
||||||
|
self.timeline.flush_usage();
|
||||||
|
if let Some(batch) = early_tools.take() {
|
||||||
|
self.drain_early_tool_batch(
|
||||||
|
history,
|
||||||
|
annotate,
|
||||||
|
batch,
|
||||||
|
false,
|
||||||
|
true,
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
return Err(EngineError::PauseRequested);
|
||||||
|
}
|
||||||
|
Err(error) => {
|
||||||
|
self.timeline.abort_current_block();
|
||||||
|
self.timeline.flush_usage();
|
||||||
|
if let Some(mut batch) = early_tools.take() {
|
||||||
|
let cleanup_error = self
|
||||||
|
.cancel_and_terminalize_early_batch(
|
||||||
|
history,
|
||||||
|
annotate,
|
||||||
|
&mut batch,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
return Err(cleanup_error.unwrap_or(error));
|
||||||
|
}
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
if let Event::Error(err) = &event {
|
if let Event::Error(err) = &event {
|
||||||
self.timeline.abort_current_block();
|
self.timeline.abort_current_block();
|
||||||
self.timeline.flush_usage();
|
self.timeline.flush_usage();
|
||||||
|
if let Some(mut batch) = early_tools.take() {
|
||||||
|
if let Some(error) = self
|
||||||
|
.cancel_and_terminalize_early_batch(
|
||||||
|
history,
|
||||||
|
annotate,
|
||||||
|
&mut batch,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
}
|
||||||
return Err(EngineError::Client(ClientError::Api {
|
return Err(EngineError::Client(ClientError::Api {
|
||||||
status: None,
|
status: None,
|
||||||
code: err.code.clone(),
|
code: err.code.clone(),
|
||||||
@@ -2291,6 +2967,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
// Preserve any UsageEvent already received for billing, but
|
// Preserve any UsageEvent already received for billing, but
|
||||||
// do not treat this paused request as a normal completion.
|
// do not treat this paused request as a normal completion.
|
||||||
self.timeline.flush_usage();
|
self.timeline.flush_usage();
|
||||||
|
if let Some(batch) = early_tools.take() {
|
||||||
|
self.drain_early_tool_batch(history, annotate, batch, false, true)
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
return Err(EngineError::PauseRequested);
|
return Err(EngineError::PauseRequested);
|
||||||
}
|
}
|
||||||
cancel = self.cancel_rx.recv() => {
|
cancel = self.cancel_rx.recv() => {
|
||||||
@@ -2299,6 +2979,14 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
}
|
}
|
||||||
self.timeline.abort_current_block();
|
self.timeline.abort_current_block();
|
||||||
self.timeline.flush_usage();
|
self.timeline.flush_usage();
|
||||||
|
if let Some(mut batch) = early_tools.take() {
|
||||||
|
if let Some(error) = self
|
||||||
|
.cancel_and_terminalize_early_batch(history, annotate, &mut batch)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
}
|
||||||
return Err(EngineError::Cancelled);
|
return Err(EngineError::Cancelled);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -2306,7 +2994,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
|
|||||||
// ストリーム完了時に集約済み Usage を 1 度だけ発火
|
// ストリーム完了時に集約済み Usage を 1 度だけ発火
|
||||||
self.timeline.flush_usage();
|
self.timeline.flush_usage();
|
||||||
debug!(event_count = event_count, "Stream completed");
|
debug!(event_count = event_count, "Stream completed");
|
||||||
Ok(StreamCompletion::Complete)
|
Ok(StreamResponseOutput {
|
||||||
|
completion: StreamCompletion::Complete,
|
||||||
|
early_tools,
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Execute tools and push results to history.
|
/// Execute tools and push results to history.
|
||||||
@@ -2361,6 +3052,7 @@ impl<C: LlmClient, A: Send + Sync> Engine<C, Mutable, A> {
|
|||||||
llm_call_count: 0,
|
llm_call_count: 0,
|
||||||
tool_execution_batch_count: 0,
|
tool_execution_batch_count: 0,
|
||||||
max_turns: None,
|
max_turns: None,
|
||||||
|
tool_call_dispatch_mode: ToolCallDispatchMode::default(),
|
||||||
tool_execution_policy: ToolExecutionPolicy::default(),
|
tool_execution_policy: ToolExecutionPolicy::default(),
|
||||||
turn_start_cbs: Vec::new(),
|
turn_start_cbs: Vec::new(),
|
||||||
turn_end_cbs: Vec::new(),
|
turn_end_cbs: Vec::new(),
|
||||||
@@ -2561,9 +3253,18 @@ impl<C: LlmClient, A: Send + Sync> Engine<C, Mutable, A> {
|
|||||||
self.max_turns = max_turns;
|
self.max_turns = max_turns;
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Apply configuration (reserved for future extensions)
|
/// Select when provider-confirmed tool calls may begin execution.
|
||||||
#[allow(dead_code)]
|
pub fn set_tool_call_dispatch_mode(&mut self, mode: ToolCallDispatchMode) {
|
||||||
pub fn config(self, _config: EngineConfig) -> Self {
|
self.tool_call_dispatch_mode = mode;
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn tool_call_dispatch_mode(&self) -> ToolCallDispatchMode {
|
||||||
|
self.tool_call_dispatch_mode
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Apply configuration.
|
||||||
|
pub fn config(mut self, config: EngineConfig) -> Self {
|
||||||
|
self.tool_call_dispatch_mode = config.tool_call_dispatch;
|
||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -2640,6 +3341,7 @@ impl<C: LlmClient, A: Send + Sync> Engine<C, Mutable, A> {
|
|||||||
llm_call_count: self.llm_call_count,
|
llm_call_count: self.llm_call_count,
|
||||||
tool_execution_batch_count: self.tool_execution_batch_count,
|
tool_execution_batch_count: self.tool_execution_batch_count,
|
||||||
max_turns: self.max_turns,
|
max_turns: self.max_turns,
|
||||||
|
tool_call_dispatch_mode: self.tool_call_dispatch_mode,
|
||||||
tool_execution_policy: self.tool_execution_policy,
|
tool_execution_policy: self.tool_execution_policy,
|
||||||
turn_start_cbs: self.turn_start_cbs,
|
turn_start_cbs: self.turn_start_cbs,
|
||||||
turn_end_cbs: self.turn_end_cbs,
|
turn_end_cbs: self.turn_end_cbs,
|
||||||
@@ -2823,6 +3525,7 @@ impl<C: LlmClient, A: Send + Sync> Engine<C, Locked, A> {
|
|||||||
llm_call_count: self.llm_call_count,
|
llm_call_count: self.llm_call_count,
|
||||||
tool_execution_batch_count: self.tool_execution_batch_count,
|
tool_execution_batch_count: self.tool_execution_batch_count,
|
||||||
max_turns: self.max_turns,
|
max_turns: self.max_turns,
|
||||||
|
tool_call_dispatch_mode: self.tool_call_dispatch_mode,
|
||||||
tool_execution_policy: self.tool_execution_policy,
|
tool_execution_policy: self.tool_execution_policy,
|
||||||
turn_start_cbs: self.turn_start_cbs,
|
turn_start_cbs: self.turn_start_cbs,
|
||||||
turn_end_cbs: self.turn_end_cbs,
|
turn_end_cbs: self.turn_end_cbs,
|
||||||
|
|||||||
@@ -22,7 +22,7 @@ pub use agen_macros::{description, tool, tool_registry};
|
|||||||
pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope};
|
pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope};
|
||||||
pub use engine::{
|
pub use engine::{
|
||||||
Engine, EngineConfig, EngineError, EngineResult, EngineRunExit, EngineRunOutput,
|
Engine, EngineConfig, EngineError, EngineResult, EngineRunExit, EngineRunOutput,
|
||||||
LlmRetryNotice, RunInterruptionReason, ToolRegistryError,
|
LlmRetryNotice, RunInterruptionReason, ToolCallDispatchMode, ToolRegistryError,
|
||||||
};
|
};
|
||||||
pub use handler::ToolUseBlockStart;
|
pub use handler::ToolUseBlockStart;
|
||||||
pub use history::{History, HistoryEntry};
|
pub use history::{History, HistoryEntry};
|
||||||
|
|||||||
@@ -38,6 +38,16 @@ impl std::fmt::Display for ConfigWarning {
|
|||||||
|
|
||||||
pub type ResponseStream = Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>;
|
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
|
/// LLMクライアントのtrait
|
||||||
///
|
///
|
||||||
/// 各プロバイダはこのtraitを実装し、統一されたインターフェースを提供する。
|
/// 各プロバイダはこのtraitを実装し、統一されたインターフェースを提供する。
|
||||||
@@ -59,6 +69,11 @@ pub trait LlmClient: Send + Sync {
|
|||||||
/// compaction) without access to the original construction parameters.
|
/// compaction) without access to the original construction parameters.
|
||||||
fn clone_boxed(&self) -> Box<dyn LlmClient>;
|
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
|
/// # Arguments
|
||||||
@@ -92,6 +107,10 @@ impl LlmClient for Box<dyn LlmClient> {
|
|||||||
(**self).clone_boxed()
|
(**self).clone_boxed()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
|
||||||
|
(**self).tool_call_completion_support()
|
||||||
|
}
|
||||||
|
|
||||||
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
|
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
|
||||||
(**self).validate_config(config)
|
(**self).validate_config(config)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ use crate::llm_client::{
|
|||||||
ClientError,
|
ClientError,
|
||||||
auth::AuthRequirement,
|
auth::AuthRequirement,
|
||||||
capability::ModelCapability,
|
capability::ModelCapability,
|
||||||
|
client::ToolCallCompletionSupport,
|
||||||
event::{BlockType, Event, ReasoningBlockData},
|
event::{BlockType, Event, ReasoningBlockData},
|
||||||
scheme::Scheme,
|
scheme::Scheme,
|
||||||
types::Request,
|
types::Request,
|
||||||
@@ -104,4 +105,8 @@ impl Scheme for AnthropicScheme {
|
|||||||
fn default_capability(&self) -> ModelCapability {
|
fn default_capability(&self) -> ModelCapability {
|
||||||
super::capability::default_capability()
|
super::capability::default_capability()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
|
||||||
|
ToolCallCompletionSupport::PerBlock
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ use serde_json::Value;
|
|||||||
|
|
||||||
use super::auth::AuthRequirement;
|
use super::auth::AuthRequirement;
|
||||||
use super::capability::ModelCapability;
|
use super::capability::ModelCapability;
|
||||||
use super::client::ConfigWarning;
|
use super::client::{ConfigWarning, ToolCallCompletionSupport};
|
||||||
use super::error::ClientError;
|
use super::error::ClientError;
|
||||||
use super::event::Event;
|
use super::event::Event;
|
||||||
use super::types::{Request, RequestConfig};
|
use super::types::{Request, RequestConfig};
|
||||||
@@ -81,6 +81,11 @@ pub trait Scheme: Clone + Send + Sync + 'static {
|
|||||||
/// 高レベルの client 構築層の責務で、scheme はここには関与しない。
|
/// 高レベルの client 構築層の責務で、scheme はここには関与しない。
|
||||||
fn default_capability(&self) -> ModelCapability;
|
fn default_capability(&self) -> ModelCapability;
|
||||||
|
|
||||||
|
/// Strongest safe completion boundary for streamed tool calls.
|
||||||
|
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
|
||||||
|
ToolCallCompletionSupport::ResponseComplete
|
||||||
|
}
|
||||||
|
|
||||||
/// scheme 側でサポートしていない `RequestConfig` フィールドを
|
/// scheme 側でサポートしていない `RequestConfig` フィールドを
|
||||||
/// 警告として返す(例: OpenAI Chat は `top_k` 非対応)。
|
/// 警告として返す(例: OpenAI Chat は `top_k` 非対応)。
|
||||||
/// デフォルトは空 Vec。
|
/// デフォルトは空 Vec。
|
||||||
@@ -89,3 +94,29 @@ pub trait Scheme: Clone + Send + Sync + 'static {
|
|||||||
Vec::new()
|
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
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ use crate::llm_client::{
|
|||||||
ClientError,
|
ClientError,
|
||||||
auth::AuthRequirement,
|
auth::AuthRequirement,
|
||||||
capability::ModelCapability,
|
capability::ModelCapability,
|
||||||
client::ConfigWarning,
|
client::{ConfigWarning, ToolCallCompletionSupport},
|
||||||
event::Event,
|
event::Event,
|
||||||
scheme::Scheme,
|
scheme::Scheme,
|
||||||
types::{Request, RequestConfig},
|
types::{Request, RequestConfig},
|
||||||
@@ -56,6 +56,10 @@ impl Scheme for OpenAIResponsesScheme {
|
|||||||
super::capability::default_capability()
|
super::capability::default_capability()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
|
||||||
|
ToolCallCompletionSupport::PerBlock
|
||||||
|
}
|
||||||
|
|
||||||
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
|
fn validate_config(&self, config: &RequestConfig) -> Vec<ConfigWarning> {
|
||||||
let mut warnings = Vec::new();
|
let mut warnings = Vec::new();
|
||||||
// Some compatible backends reject `max_output_tokens` with HTTP 400.
|
// Some compatible backends reject `max_output_tokens` with HTTP 400.
|
||||||
|
|||||||
@@ -510,6 +510,10 @@ impl<S: Scheme + Clone + 'static> LlmClient for HttpTransport<S> {
|
|||||||
self.scheme.validate_config(config)
|
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> {
|
async fn stream(&self, request: Request) -> Result<ResponseStream, ClientError> {
|
||||||
let total_started = Instant::now();
|
let total_started = Instant::now();
|
||||||
let path = self.scheme.path(&self.model_id);
|
let path = self.scheme.path(&self.model_id);
|
||||||
|
|||||||
@@ -256,8 +256,8 @@ where
|
|||||||
H: Handler<ToolUseBlockKind>,
|
H: Handler<ToolUseBlockKind>,
|
||||||
{
|
{
|
||||||
handler: H,
|
handler: H,
|
||||||
scope: Option<H::Scope>,
|
scopes: HashMap<usize, H::Scope>,
|
||||||
current_tool: Option<(String, String)>, // (id, name)
|
tools: HashMap<usize, (String, String)>, // index -> (id, name)
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<H> ToolUseBlockHandlerWrapper<H>
|
impl<H> ToolUseBlockHandlerWrapper<H>
|
||||||
@@ -267,8 +267,8 @@ where
|
|||||||
fn new(handler: H) -> Self {
|
fn new(handler: H) -> Self {
|
||||||
Self {
|
Self {
|
||||||
handler,
|
handler,
|
||||||
scope: None,
|
scopes: HashMap::new(),
|
||||||
current_tool: None,
|
tools: HashMap::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -279,60 +279,55 @@ where
|
|||||||
H::Scope: Send + Sync,
|
H::Scope: Send + Sync,
|
||||||
{
|
{
|
||||||
fn dispatch_start(&mut self, start: &BlockStart) {
|
fn dispatch_start(&mut self, start: &BlockStart) {
|
||||||
if let Some(scope) = &mut self.scope {
|
if let BlockMetadata::ToolUse { id, name } = &start.metadata {
|
||||||
if let BlockMetadata::ToolUse { id, name } = &start.metadata {
|
let scope = self.scopes.entry(start.index).or_default();
|
||||||
self.current_tool = Some((id.clone(), name.clone()));
|
self.tools.insert(start.index, (id.clone(), name.clone()));
|
||||||
self.handler.on_event(
|
self.handler.on_event(
|
||||||
scope,
|
scope,
|
||||||
&ToolUseBlockEvent::Start(ToolUseBlockStart {
|
&ToolUseBlockEvent::Start(ToolUseBlockStart {
|
||||||
index: start.index,
|
index: start.index,
|
||||||
id: id.clone(),
|
id: id.clone(),
|
||||||
name: name.clone(),
|
name: name.clone(),
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn dispatch_delta(&mut self, delta: &BlockDelta) {
|
fn dispatch_delta(&mut self, delta: &BlockDelta) {
|
||||||
if let Some(scope) = &mut self.scope {
|
if let DeltaContent::InputJson(json) = &delta.delta {
|
||||||
if let DeltaContent::InputJson(json) = &delta.delta {
|
let scope = self.scopes.entry(delta.index).or_default();
|
||||||
self.handler
|
self.handler
|
||||||
.on_event(scope, &ToolUseBlockEvent::InputJsonDelta(json.clone()));
|
.on_event(scope, &ToolUseBlockEvent::InputJsonDelta(json.clone()));
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn dispatch_stop(&mut self, stop: &BlockStop) {
|
fn dispatch_stop(&mut self, stop: &BlockStop) {
|
||||||
if let Some(scope) = &mut self.scope {
|
if let (Some(mut scope), Some((id, name))) = (
|
||||||
if let Some((id, name)) = self.current_tool.take() {
|
self.scopes.remove(&stop.index),
|
||||||
self.handler.on_event(
|
self.tools.remove(&stop.index),
|
||||||
scope,
|
) {
|
||||||
&ToolUseBlockEvent::Stop(ToolUseBlockStop {
|
self.handler.on_event(
|
||||||
index: stop.index,
|
&mut scope,
|
||||||
id,
|
&ToolUseBlockEvent::Stop(ToolUseBlockStop {
|
||||||
name,
|
index: stop.index,
|
||||||
}),
|
id,
|
||||||
);
|
name,
|
||||||
}
|
}),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn dispatch_abort(&mut self, _abort: &BlockAbort) {
|
fn dispatch_abort(&mut self, abort: &BlockAbort) {
|
||||||
self.current_tool = None;
|
self.scopes.remove(&abort.index);
|
||||||
|
self.tools.remove(&abort.index);
|
||||||
}
|
}
|
||||||
|
|
||||||
fn start_scope(&mut self) {
|
fn start_scope(&mut self) {}
|
||||||
self.scope = Some(H::Scope::default());
|
|
||||||
}
|
|
||||||
|
|
||||||
fn end_scope(&mut self) {
|
fn end_scope(&mut self) {}
|
||||||
self.scope = None;
|
|
||||||
self.current_tool = None;
|
|
||||||
}
|
|
||||||
|
|
||||||
fn has_scope(&self) -> bool {
|
fn has_scope(&self) -> bool {
|
||||||
self.scope.is_some()
|
!self.scopes.is_empty()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -5,7 +5,6 @@
|
|||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
handler::{Handler, ToolUseBlockEvent, ToolUseBlockKind},
|
handler::{Handler, ToolUseBlockEvent, ToolUseBlockKind},
|
||||||
llm_client::types::parse_tool_arguments,
|
|
||||||
tool::ToolCall,
|
tool::ToolCall,
|
||||||
};
|
};
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
@@ -85,7 +84,18 @@ impl Handler<ToolUseBlockKind> for ToolCallCollector {
|
|||||||
// ブロック完了時にToolCallを確定
|
// ブロック完了時にToolCallを確定
|
||||||
if let (Some(id), Some(name)) = (scope.current_id.take(), scope.current_name.take())
|
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 tool_call = ToolCall { id, name, input };
|
||||||
|
|
||||||
@@ -144,6 +154,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]
|
#[test]
|
||||||
fn test_collect_multiple_tool_calls() {
|
fn test_collect_multiple_tool_calls() {
|
||||||
let collector = ToolCallCollector::new();
|
let collector = ToolCallCollector::new();
|
||||||
|
|||||||
@@ -11,12 +11,16 @@ use agen::interceptor::{
|
|||||||
PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo,
|
PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo,
|
||||||
};
|
};
|
||||||
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
|
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
|
||||||
|
use agen::llm_client::{
|
||||||
|
ClientError, LlmClient, Request, ResponseStream, ToolCallCompletionSupport,
|
||||||
|
};
|
||||||
use agen::tool::{
|
use agen::tool::{
|
||||||
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
|
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
|
||||||
ToolResultDisposition,
|
ToolResultDisposition,
|
||||||
};
|
};
|
||||||
use agen::{
|
use agen::{
|
||||||
Engine, EngineError, EngineRunExit, History, Item, RunInterruptionReason, ToolExecutionPolicy,
|
Engine, EngineError, EngineRunExit, History, Item, RunInterruptionReason, ToolCallDispatchMode,
|
||||||
|
ToolExecutionPolicy,
|
||||||
};
|
};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
|
||||||
@@ -251,10 +255,476 @@ impl Tool for ContextRecordingTool {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct ControlledStreamClient {
|
||||||
|
receiver: Arc<Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<Result<Event, ClientError>>>>>,
|
||||||
|
completion_support: ToolCallCompletionSupport,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ControlledStreamClient {
|
||||||
|
fn new(
|
||||||
|
completion_support: ToolCallCompletionSupport,
|
||||||
|
) -> (
|
||||||
|
Self,
|
||||||
|
tokio::sync::mpsc::UnboundedSender<Result<Event, ClientError>>,
|
||||||
|
) {
|
||||||
|
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
|
||||||
|
(
|
||||||
|
Self {
|
||||||
|
receiver: Arc::new(Mutex::new(Some(rx))),
|
||||||
|
completion_support,
|
||||||
|
},
|
||||||
|
tx,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl LlmClient for ControlledStreamClient {
|
||||||
|
async fn stream(&self, _request: Request) -> Result<ResponseStream, ClientError> {
|
||||||
|
let mut rx = self
|
||||||
|
.receiver
|
||||||
|
.lock()
|
||||||
|
.unwrap()
|
||||||
|
.take()
|
||||||
|
.ok_or_else(|| ClientError::Api {
|
||||||
|
status: Some(500),
|
||||||
|
code: Some("controlled_stream_exhausted".to_string()),
|
||||||
|
message: "controlled stream already consumed".to_string(),
|
||||||
|
retry_after: None,
|
||||||
|
})?;
|
||||||
|
Ok(Box::pin(futures::stream::poll_fn(move |cx| {
|
||||||
|
rx.poll_recv(cx)
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn clone_boxed(&self) -> Box<dyn LlmClient> {
|
||||||
|
Box::new(self.clone())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn tool_call_completion_support(&self) -> ToolCallCompletionSupport {
|
||||||
|
self.completion_support
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct BarrierTool {
|
||||||
|
name: &'static str,
|
||||||
|
starts: Arc<AtomicUsize>,
|
||||||
|
started: Arc<tokio::sync::Notify>,
|
||||||
|
release: Arc<tokio::sync::Notify>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BarrierTool {
|
||||||
|
fn new(name: &'static str) -> Self {
|
||||||
|
Self {
|
||||||
|
name,
|
||||||
|
starts: Arc::new(AtomicUsize::new(0)),
|
||||||
|
started: Arc::new(tokio::sync::Notify::new()),
|
||||||
|
release: Arc::new(tokio::sync::Notify::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn definition(&self) -> ToolDefinition {
|
||||||
|
let tool = self.clone();
|
||||||
|
Arc::new(move || {
|
||||||
|
(
|
||||||
|
ToolMeta::new(tool.name).input_schema(serde_json::json!({"type":"object"})),
|
||||||
|
Arc::new(tool.clone()) as Arc<dyn Tool>,
|
||||||
|
)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for BarrierTool {
|
||||||
|
async fn execute(
|
||||||
|
&self,
|
||||||
|
_input_json: &str,
|
||||||
|
_ctx: ToolExecutionContext,
|
||||||
|
) -> Result<ToolOutput, ToolError> {
|
||||||
|
self.starts.fetch_add(1, Ordering::SeqCst);
|
||||||
|
self.started.notify_one();
|
||||||
|
self.release.notified().await;
|
||||||
|
Ok(format!("{} complete", self.name).into())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// =============================================================================
|
// =============================================================================
|
||||||
// Tests
|
// Tests
|
||||||
// =============================================================================
|
// =============================================================================
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn default_dispatch_waits_for_response_completion() {
|
||||||
|
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
|
||||||
|
let tool = BarrierTool::new("barrier_default");
|
||||||
|
let probe = tool.clone();
|
||||||
|
let mut engine = Engine::new(client);
|
||||||
|
engine.set_max_turns(Some(1));
|
||||||
|
engine.register_tool(tool.definition());
|
||||||
|
|
||||||
|
let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||||
|
engine.on_stream_event(move |_, _, event| {
|
||||||
|
if matches!(event, Event::Ping(_)) {
|
||||||
|
let _ = observed_tx.send(());
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let run = tokio::spawn(async move {
|
||||||
|
let mut history = History::new();
|
||||||
|
let output = engine.run(&mut history, "run").await;
|
||||||
|
(output, history)
|
||||||
|
});
|
||||||
|
|
||||||
|
tx.send(Ok(Event::tool_use_start(
|
||||||
|
0,
|
||||||
|
"call_default",
|
||||||
|
"barrier_default",
|
||||||
|
)))
|
||||||
|
.unwrap();
|
||||||
|
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
|
||||||
|
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
|
||||||
|
tx.send(Ok(Event::ping())).unwrap();
|
||||||
|
observed_rx.recv().await.unwrap();
|
||||||
|
assert_eq!(probe.starts.load(Ordering::SeqCst), 0);
|
||||||
|
|
||||||
|
tx.send(Ok(Event::Status(StatusEvent {
|
||||||
|
status: ResponseStatus::Completed,
|
||||||
|
})))
|
||||||
|
.unwrap();
|
||||||
|
drop(tx);
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), probe.started.notified())
|
||||||
|
.await
|
||||||
|
.expect("tool starts after response completion");
|
||||||
|
probe.release.notify_one();
|
||||||
|
let _ = run.await.unwrap();
|
||||||
|
assert_eq!(probe.starts.load(Ordering::SeqCst), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_result() {
|
||||||
|
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
|
||||||
|
let first = BarrierTool::new("barrier_a");
|
||||||
|
let second = BarrierTool::new("barrier_b");
|
||||||
|
let first_probe = first.clone();
|
||||||
|
let second_probe = second.clone();
|
||||||
|
let mut engine = Engine::new(client);
|
||||||
|
engine.set_max_turns(Some(1));
|
||||||
|
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
|
||||||
|
engine.register_tool(first.definition());
|
||||||
|
engine.register_tool(second.definition());
|
||||||
|
|
||||||
|
let run = tokio::spawn(async move {
|
||||||
|
let mut history = History::new();
|
||||||
|
let output = engine.run(&mut history, "run both").await;
|
||||||
|
(output, history)
|
||||||
|
});
|
||||||
|
|
||||||
|
tx.send(Ok(Event::tool_use_start(0, "call_a", "barrier_a")))
|
||||||
|
.unwrap();
|
||||||
|
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
|
||||||
|
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), first_probe.started.notified())
|
||||||
|
.await
|
||||||
|
.expect("first tool starts while stream is open");
|
||||||
|
|
||||||
|
tx.send(Ok(Event::tool_use_start(1, "call_b", "barrier_b")))
|
||||||
|
.unwrap();
|
||||||
|
tx.send(Ok(Event::tool_input_delta(1, r#"{}"#))).unwrap();
|
||||||
|
tx.send(Ok(Event::tool_use_stop(1))).unwrap();
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), second_probe.started.notified())
|
||||||
|
.await
|
||||||
|
.expect("second tool starts before first completes");
|
||||||
|
assert_eq!(first_probe.starts.load(Ordering::SeqCst), 1);
|
||||||
|
assert_eq!(second_probe.starts.load(Ordering::SeqCst), 1);
|
||||||
|
|
||||||
|
tx.send(Ok(Event::Status(StatusEvent {
|
||||||
|
status: ResponseStatus::Completed,
|
||||||
|
})))
|
||||||
|
.unwrap();
|
||||||
|
drop(tx);
|
||||||
|
first_probe.release.notify_one();
|
||||||
|
second_probe.release.notify_one();
|
||||||
|
let (_output, history) = run.await.unwrap();
|
||||||
|
|
||||||
|
let calls = history
|
||||||
|
.iter()
|
||||||
|
.filter(|entry| matches!(entry.item, Item::ToolCall { .. }))
|
||||||
|
.count();
|
||||||
|
let results = history
|
||||||
|
.iter()
|
||||||
|
.filter(|entry| matches!(entry.item, Item::ToolResult { .. }))
|
||||||
|
.count();
|
||||||
|
assert_eq!((calls, results), (2, 2));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn later_history_failure_terminalizes_already_started_early_sibling() {
|
||||||
|
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
|
||||||
|
let first = BarrierTool::new("history_first");
|
||||||
|
let second = BarrierTool::new("history_second");
|
||||||
|
let first_probe = first.clone();
|
||||||
|
let second_probe = second.clone();
|
||||||
|
let mut engine = Engine::new(client);
|
||||||
|
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
|
||||||
|
engine.set_tool_execution_policy(ToolExecutionPolicy {
|
||||||
|
pause_safe_boundary_timeout: Duration::from_millis(1),
|
||||||
|
cancellation_request_timeout: Duration::from_millis(1),
|
||||||
|
terminal_confirmation_timeout: Duration::from_millis(1),
|
||||||
|
});
|
||||||
|
engine.register_tool(first.definition());
|
||||||
|
engine.register_tool(second.definition());
|
||||||
|
engine.on_history_append(|item| match item {
|
||||||
|
Item::ToolCall { call_id, .. } if call_id == "call_history_second" => {
|
||||||
|
Err("durable append rejected".to_string())
|
||||||
|
}
|
||||||
|
_ => Ok(()),
|
||||||
|
});
|
||||||
|
|
||||||
|
let run = tokio::spawn(async move {
|
||||||
|
let mut history = History::new();
|
||||||
|
let output = engine.run(&mut history, "run").await;
|
||||||
|
(output, history)
|
||||||
|
});
|
||||||
|
tx.send(Ok(Event::tool_use_start(
|
||||||
|
0,
|
||||||
|
"call_history_first",
|
||||||
|
"history_first",
|
||||||
|
)))
|
||||||
|
.unwrap();
|
||||||
|
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
|
||||||
|
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), first_probe.started.notified())
|
||||||
|
.await
|
||||||
|
.expect("first side effect starts");
|
||||||
|
tx.send(Ok(Event::tool_use_start(
|
||||||
|
1,
|
||||||
|
"call_history_second",
|
||||||
|
"history_second",
|
||||||
|
)))
|
||||||
|
.unwrap();
|
||||||
|
tx.send(Ok(Event::tool_input_delta(1, r#"{}"#))).unwrap();
|
||||||
|
tx.send(Ok(Event::tool_use_stop(1))).unwrap();
|
||||||
|
drop(tx);
|
||||||
|
|
||||||
|
let (output, history) = run.await.unwrap();
|
||||||
|
assert!(matches!(
|
||||||
|
output.result,
|
||||||
|
EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(
|
||||||
|
EngineError::HistoryAppend(ref message)
|
||||||
|
)) if message == "durable append rejected"
|
||||||
|
));
|
||||||
|
assert_eq!(first_probe.starts.load(Ordering::SeqCst), 1);
|
||||||
|
assert_eq!(second_probe.starts.load(Ordering::SeqCst), 0);
|
||||||
|
assert!(history.iter().any(|entry| matches!(
|
||||||
|
&entry.item,
|
||||||
|
Item::ToolResult {
|
||||||
|
call_id,
|
||||||
|
disposition: ToolResultDisposition::OutcomeUnknown,
|
||||||
|
..
|
||||||
|
} if call_id == "call_history_first"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn malformed_completed_arguments_commit_error_without_executing_tool() {
|
||||||
|
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
|
||||||
|
let tool = SlowTool::new("must_not_run", 0);
|
||||||
|
let probe = tool.clone();
|
||||||
|
let mut engine = Engine::new(client);
|
||||||
|
engine.set_max_turns(Some(1));
|
||||||
|
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
|
||||||
|
engine.register_tool(tool.definition());
|
||||||
|
let run = tokio::spawn(async move {
|
||||||
|
let mut history = History::new();
|
||||||
|
let output = engine.run(&mut history, "run").await;
|
||||||
|
(output, history)
|
||||||
|
});
|
||||||
|
|
||||||
|
tx.send(Ok(Event::tool_use_start(
|
||||||
|
0,
|
||||||
|
"call_malformed",
|
||||||
|
"must_not_run",
|
||||||
|
)))
|
||||||
|
.unwrap();
|
||||||
|
tx.send(Ok(Event::tool_input_delta(0, r#"{"broken":"#)))
|
||||||
|
.unwrap();
|
||||||
|
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
|
||||||
|
tx.send(Ok(Event::Status(StatusEvent {
|
||||||
|
status: ResponseStatus::Completed,
|
||||||
|
})))
|
||||||
|
.unwrap();
|
||||||
|
drop(tx);
|
||||||
|
|
||||||
|
let (_output, history) = run.await.unwrap();
|
||||||
|
assert_eq!(probe.call_count(), 0);
|
||||||
|
assert!(history.iter().any(|entry| matches!(
|
||||||
|
&entry.item,
|
||||||
|
Item::ToolResult {
|
||||||
|
call_id,
|
||||||
|
summary,
|
||||||
|
..
|
||||||
|
} if call_id == "call_malformed" && summary.contains("complete JSON object")
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn pause_during_early_execution_drains_result_and_resume_does_not_rerun() {
|
||||||
|
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
|
||||||
|
let tool = BarrierTool::new("barrier_pause");
|
||||||
|
let probe = tool.clone();
|
||||||
|
let mut engine = Engine::new(client);
|
||||||
|
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
|
||||||
|
engine.register_tool(tool.definition());
|
||||||
|
let pause = engine.pause_sender();
|
||||||
|
|
||||||
|
let run = tokio::spawn(async move {
|
||||||
|
let mut history = History::new();
|
||||||
|
let output = engine.run(&mut history, "run").await;
|
||||||
|
(output, history)
|
||||||
|
});
|
||||||
|
tx.send(Ok(Event::tool_use_start(0, "call_pause", "barrier_pause")))
|
||||||
|
.unwrap();
|
||||||
|
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
|
||||||
|
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), probe.started.notified())
|
||||||
|
.await
|
||||||
|
.expect("tool starts before pause");
|
||||||
|
pause.send(()).await.unwrap();
|
||||||
|
probe.release.notify_one();
|
||||||
|
drop(tx);
|
||||||
|
|
||||||
|
let (output, mut history) = run.await.unwrap();
|
||||||
|
assert!(matches!(output.result, EngineRunExit::Paused));
|
||||||
|
assert!(history.iter().any(|entry| matches!(
|
||||||
|
&entry.item,
|
||||||
|
Item::ToolResult { call_id, .. } if call_id == "call_pause"
|
||||||
|
)));
|
||||||
|
let mut engine = output.engine;
|
||||||
|
let _ = engine.resume(&mut history).await;
|
||||||
|
assert_eq!(probe.starts.load(Ordering::SeqCst), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn interrupted_stream_after_early_start_terminalizes_without_continuation_or_reexecution() {
|
||||||
|
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
|
||||||
|
let tool = BarrierTool::new("barrier_interrupted");
|
||||||
|
let probe = tool.clone();
|
||||||
|
let mut engine = Engine::new(client);
|
||||||
|
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
|
||||||
|
engine.register_tool(tool.definition());
|
||||||
|
|
||||||
|
let run = tokio::spawn(async move {
|
||||||
|
let mut history = History::new();
|
||||||
|
let output = engine.run(&mut history, "run").await;
|
||||||
|
(output, history)
|
||||||
|
});
|
||||||
|
tx.send(Ok(Event::tool_use_start(
|
||||||
|
0,
|
||||||
|
"call_interrupted",
|
||||||
|
"barrier_interrupted",
|
||||||
|
)))
|
||||||
|
.unwrap();
|
||||||
|
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
|
||||||
|
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), probe.started.notified())
|
||||||
|
.await
|
||||||
|
.expect("tool starts before stream interruption");
|
||||||
|
tx.send(Err(ClientError::Api {
|
||||||
|
status: None,
|
||||||
|
code: Some("connection_lost".to_string()),
|
||||||
|
message: "connection lost".to_string(),
|
||||||
|
retry_after: None,
|
||||||
|
}))
|
||||||
|
.unwrap();
|
||||||
|
probe.release.notify_one();
|
||||||
|
drop(tx);
|
||||||
|
|
||||||
|
let (output, mut history) = run.await.unwrap();
|
||||||
|
assert!(matches!(
|
||||||
|
output.result,
|
||||||
|
EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(EngineError::Client(
|
||||||
|
ClientError::Api {
|
||||||
|
code: Some(ref code),
|
||||||
|
..
|
||||||
|
}
|
||||||
|
))) if code == "early_tool_stream_interrupted"
|
||||||
|
));
|
||||||
|
assert_eq!(probe.starts.load(Ordering::SeqCst), 1);
|
||||||
|
assert_eq!(
|
||||||
|
history
|
||||||
|
.iter()
|
||||||
|
.filter(|entry| matches!(&entry.item, Item::ToolCall { call_id, .. } if call_id == "call_interrupted"))
|
||||||
|
.count(),
|
||||||
|
1
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
history
|
||||||
|
.iter()
|
||||||
|
.filter(|entry| matches!(&entry.item, Item::ToolResult { call_id, .. } if call_id == "call_interrupted"))
|
||||||
|
.count(),
|
||||||
|
1
|
||||||
|
);
|
||||||
|
|
||||||
|
let mut engine = output.engine;
|
||||||
|
let _ = engine.resume(&mut history).await;
|
||||||
|
assert_eq!(
|
||||||
|
probe.starts.load(Ordering::SeqCst),
|
||||||
|
1,
|
||||||
|
"terminalized early calls must not be resumed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unsupported_provider_visibly_defers_opt_in_dispatch() {
|
||||||
|
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::ResponseComplete);
|
||||||
|
let tool = BarrierTool::new("barrier_fallback");
|
||||||
|
let probe = tool.clone();
|
||||||
|
let mut engine = Engine::new(client);
|
||||||
|
engine.set_max_turns(Some(1));
|
||||||
|
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete);
|
||||||
|
engine.register_tool(tool.definition());
|
||||||
|
let warnings = Arc::new(Mutex::new(Vec::new()));
|
||||||
|
let warning_probe = warnings.clone();
|
||||||
|
engine.on_warning(move |warning| warning_probe.lock().unwrap().push(warning.to_string()));
|
||||||
|
let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||||
|
engine.on_stream_event(move |_, _, event| {
|
||||||
|
if matches!(event, Event::Ping(_)) {
|
||||||
|
let _ = observed_tx.send(());
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let run = tokio::spawn(async move {
|
||||||
|
let mut history = History::new();
|
||||||
|
let output = engine.run(&mut history, "run").await;
|
||||||
|
(output, history)
|
||||||
|
});
|
||||||
|
tx.send(Ok(Event::tool_use_start(
|
||||||
|
0,
|
||||||
|
"call_fallback",
|
||||||
|
"barrier_fallback",
|
||||||
|
)))
|
||||||
|
.unwrap();
|
||||||
|
tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap();
|
||||||
|
tx.send(Ok(Event::tool_use_stop(0))).unwrap();
|
||||||
|
tx.send(Ok(Event::ping())).unwrap();
|
||||||
|
observed_rx.recv().await.unwrap();
|
||||||
|
assert_eq!(probe.starts.load(Ordering::SeqCst), 0);
|
||||||
|
assert!(warnings.lock().unwrap().iter().any(|warning| {
|
||||||
|
warning.contains("provider does not expose a trustworthy per-call completion boundary")
|
||||||
|
}));
|
||||||
|
|
||||||
|
tx.send(Ok(Event::Status(StatusEvent {
|
||||||
|
status: ResponseStatus::Completed,
|
||||||
|
})))
|
||||||
|
.unwrap();
|
||||||
|
drop(tx);
|
||||||
|
tokio::time::timeout(Duration::from_secs(1), probe.started.notified())
|
||||||
|
.await
|
||||||
|
.expect("fallback starts tool only after response completion");
|
||||||
|
probe.release.notify_one();
|
||||||
|
let _ = run.await.unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
/// Verify that multiple tools are executed in parallel
|
/// Verify that multiple tools are executed in parallel
|
||||||
///
|
///
|
||||||
/// If each tool takes 100ms, sequential execution would take 300ms+,
|
/// If each tool takes 100ms, sequential execution would take 300ms+,
|
||||||
|
|||||||
@@ -542,6 +542,8 @@ pub struct EngineManifestConfig {
|
|||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub reasoning: Option<ReasoningControl>,
|
pub reasoning: Option<ReasoningControl>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
|
pub tool_call_dispatch: Option<crate::ToolCallDispatchMode>,
|
||||||
|
#[serde(default)]
|
||||||
pub tool_output: ToolOutputLimitsPartial,
|
pub tool_output: ToolOutputLimitsPartial,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub file_upload: FileUploadLimitsPartial,
|
pub file_upload: FileUploadLimitsPartial,
|
||||||
@@ -868,6 +870,7 @@ impl EngineManifestConfig {
|
|||||||
top_k: upper.top_k.or(self.top_k),
|
top_k: upper.top_k.or(self.top_k),
|
||||||
stop_sequences: upper.stop_sequences.or(self.stop_sequences),
|
stop_sequences: upper.stop_sequences.or(self.stop_sequences),
|
||||||
reasoning: upper.reasoning.or(self.reasoning),
|
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),
|
tool_output: self.tool_output.merge(upper.tool_output),
|
||||||
file_upload: self.file_upload.merge(upper.file_upload),
|
file_upload: self.file_upload.merge(upper.file_upload),
|
||||||
}
|
}
|
||||||
@@ -1171,6 +1174,7 @@ impl TryFrom<WorkerManifestConfig> for WorkerManifest {
|
|||||||
top_k: cfg.engine.top_k,
|
top_k: cfg.engine.top_k,
|
||||||
stop_sequences: cfg.engine.stop_sequences.unwrap_or_default(),
|
stop_sequences: cfg.engine.stop_sequences.unwrap_or_default(),
|
||||||
reasoning: cfg.engine.reasoning,
|
reasoning: cfg.engine.reasoning,
|
||||||
|
tool_call_dispatch: cfg.engine.tool_call_dispatch.unwrap_or_default(),
|
||||||
tool_output: ToolOutputLimits {
|
tool_output: ToolOutputLimits {
|
||||||
default_max_bytes: cfg
|
default_max_bytes: cfg
|
||||||
.engine
|
.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]
|
#[test]
|
||||||
fn resolve_minimal_succeeds() {
|
fn resolve_minimal_succeeds() {
|
||||||
let manifest: WorkerManifest = minimal_valid().try_into().unwrap();
|
let manifest: WorkerManifest = minimal_valid().try_into().unwrap();
|
||||||
|
|||||||
@@ -587,6 +587,20 @@ pub struct WorkerMeta {
|
|||||||
pub name: String,
|
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.
|
/// Worker-level configuration embedded in the manifest.
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct EngineManifest {
|
pub struct EngineManifest {
|
||||||
@@ -614,6 +628,9 @@ pub struct EngineManifest {
|
|||||||
pub stop_sequences: Vec<String>,
|
pub stop_sequences: Vec<String>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub reasoning: Option<ReasoningControl>,
|
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
|
/// Byte-size caps applied to tool `content` before it reaches the
|
||||||
/// conversation history. The section is optional in TOML — when
|
/// conversation history. The section is optional in TOML — when
|
||||||
/// omitted, `ToolOutputLimits::default()` (64 KiB default cap, no
|
/// omitted, `ToolOutputLimits::default()` (64 KiB default cap, no
|
||||||
|
|||||||
@@ -965,6 +965,7 @@ fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfi
|
|||||||
stop_sequences: (!manifest.engine.stop_sequences.is_empty())
|
stop_sequences: (!manifest.engine.stop_sequences.is_empty())
|
||||||
.then_some(manifest.engine.stop_sequences.clone()),
|
.then_some(manifest.engine.stop_sequences.clone()),
|
||||||
reasoning: manifest.engine.reasoning.clone(),
|
reasoning: manifest.engine.reasoning.clone(),
|
||||||
|
tool_call_dispatch: Some(manifest.engine.tool_call_dispatch),
|
||||||
tool_output: ToolOutputLimitsPartial {
|
tool_output: ToolOutputLimitsPartial {
|
||||||
default_max_bytes: Some(manifest.engine.tool_output.default_max_bytes),
|
default_max_bytes: Some(manifest.engine.tool_output.default_max_bytes),
|
||||||
per_tool: manifest.engine.tool_output.per_tool.clone(),
|
per_tool: manifest.engine.tool_output.per_tool.clone(),
|
||||||
@@ -2226,7 +2227,8 @@ max_tokens = 3333
|
|||||||
Some("reviewer"),
|
Some("reviewer"),
|
||||||
&[("reviewer", "reviewer.toml", REVIEWER_PROFILE)],
|
&[("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 scope = vec![abs_rule(&delegated, Permission::Read)];
|
||||||
|
|
||||||
let profile_config = child_config_from_profile(
|
let profile_config = child_config_from_profile(
|
||||||
@@ -2248,6 +2250,10 @@ max_tokens = 3333
|
|||||||
SpawnProfileSelector::Inherit,
|
SpawnProfileSelector::Inherit,
|
||||||
);
|
);
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
inherit_config.engine.tool_call_dispatch,
|
||||||
|
Some(manifest::ToolCallDispatchMode::OnToolCallComplete)
|
||||||
|
);
|
||||||
for config in [profile_config, inherit_config] {
|
for config in [profile_config, inherit_config] {
|
||||||
assert_eq!(config.scope.allow, scope);
|
assert_eq!(config.scope.allow, scope);
|
||||||
assert!(config.scope.deny.is_empty());
|
assert!(config.scope.deny.is_empty());
|
||||||
|
|||||||
@@ -7138,6 +7138,12 @@ pub fn apply_worker_manifest<C: LlmClient + 'static, A: Send + Sync>(
|
|||||||
) {
|
) {
|
||||||
worker.set_request_config(request_config_from_engine_manifest(wm));
|
worker.set_request_config(request_config_from_engine_manifest(wm));
|
||||||
worker.set_max_turns(wm.max_turns.map(|n| n.get()));
|
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.
|
// Worker owns the lifecycle strategy for already-started tool operations.
|
||||||
// The provider must first accept cooperative cancellation, then confirm a
|
// The provider must first accept cooperative cancellation, then confirm a
|
||||||
// terminal result before this bounded deadline; Agen handles only the
|
// terminal result before this bounded deadline; Agen handles only the
|
||||||
@@ -8721,6 +8727,7 @@ mod build_summary_prompt_tests {
|
|||||||
top_k: Some(40),
|
top_k: Some(40),
|
||||||
stop_sequences: vec!["\n\n".into(), "</stop>".into()],
|
stop_sequences: vec!["\n\n".into(), "</stop>".into()],
|
||||||
reasoning: None,
|
reasoning: None,
|
||||||
|
tool_call_dispatch: manifest::ToolCallDispatchMode::OnToolCallComplete,
|
||||||
tool_output: manifest::ToolOutputLimits::default(),
|
tool_output: manifest::ToolOutputLimits::default(),
|
||||||
file_upload: manifest::FileUploadLimits::default(),
|
file_upload: manifest::FileUploadLimits::default(),
|
||||||
};
|
};
|
||||||
@@ -8732,6 +8739,13 @@ mod build_summary_prompt_tests {
|
|||||||
assert_eq!(config.top_p, Some(0.9));
|
assert_eq!(config.top_p, Some(0.9));
|
||||||
assert_eq!(config.top_k, Some(40));
|
assert_eq!(config.top_k, Some(40));
|
||||||
assert_eq!(config.stop_sequences, vec!["\n\n", "</stop>"]);
|
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]
|
#[test]
|
||||||
|
|||||||
@@ -121,6 +121,14 @@ ref = "anthropic/claude-sonnet-4-6"
|
|||||||
# reasoning = "medium"
|
# reasoning = "medium"
|
||||||
# reasoning = -1
|
# reasoning = -1
|
||||||
|
|
||||||
|
# 任意。デフォルト: "after_response"。
|
||||||
|
# "on_tool_call_complete" は provider が個別 ToolCall の完成を保証する場合だけ、
|
||||||
|
# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider では
|
||||||
|
# 応答完了まで明示的に遅延する。これは直列実行モードではなく、先行 call の完了を
|
||||||
|
# 待たず後続 call も開始するため、完了順・DB 反映順・採番順は保証しない。
|
||||||
|
# 早期開始後に stream が切断された場合は自動 continuation せず、安全に terminalize して停止する。
|
||||||
|
# tool_call_dispatch = "on_tool_call_complete"
|
||||||
|
|
||||||
# 任意。tool 実行 content の byte 長キャップ。
|
# 任意。tool 実行 content の byte 長キャップ。
|
||||||
# セクション省略時は default_max_bytes = 64 * 1024、per_tool 空。
|
# セクション省略時は default_max_bytes = 64 * 1024、per_tool 空。
|
||||||
# [worker.tool_output]
|
# [worker.tool_output]
|
||||||
|
|||||||
Reference in New Issue
Block a user