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

tool: add execution context

Browse Source
This commit is contained in:
Hare
2026-06-09 19:31:11 +09:00
parent b21fab82fc
commit d8aed7befe
39 changed files with 1212 additions and 259 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/llm-worker/src/interceptor.rs
+5 -1
View File
@@ -10,7 +10,7 @@ use std::sync::Arc;
use async_trait::async_trait;
use crate::Item;
use crate::tool::{Tool, ToolCall, ToolMeta, ToolResult};
use crate::tool::{Tool, ToolCall, ToolExecutionContext, ToolMeta, ToolResult};
// =============================================================================
// Action Enums
@@ -107,6 +107,8 @@ pub struct ToolCallInfo {
pub meta: ToolMeta,
/// Tool instance (for state access).
pub tool: Arc<dyn Tool>,
/// Response-local execution context for this call.
pub context: ToolExecutionContext,
}
/// Context for post-tool-call decisions.
@@ -119,6 +121,8 @@ pub struct ToolResultInfo {
pub meta: ToolMeta,
/// Tool instance (for state access).
pub tool: Arc<dyn Tool>,
/// Response-local execution context for this call.
pub context: ToolExecutionContext,
}
// =============================================================================
crates/llm-worker/src/lib.rs
+1 -1
View File
@@ -57,7 +57,7 @@ pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope};
pub use handler::ToolUseBlockStart;
pub use interceptor::Interceptor;
pub use message::{ContentPart, Item, Message, Role};
pub use tool::{ToolCall, ToolOutputLimits, ToolResult};
pub use tool::{ToolCall, ToolExecutionContext, ToolOutputLimits, ToolResult};
pub use usage_record::UsageRecord;
pub use worker::{
LlmRetryNotice, RunOutput, ToolRegistryError, Worker, WorkerConfig, WorkerError, WorkerResult,
crates/llm-worker/src/tool.rs
+47 -4
View File
@@ -189,6 +189,44 @@ impl ToolMeta {
/// ```
pub type ToolDefinition = Arc<dyn Fn() -> (ToolMeta, Arc<dyn Tool>) + Send + Sync>;
/// Per-call context supplied by the worker when executing a tool call.
///
/// The context identifies a tool call within one assistant response's tool-call
/// batch without imposing any scheduling policy on the worker. Tool
/// implementations may use it for response-local ordering, diagnostics, or
/// correlation, but it is intentionally not a handle to worker state, history,
/// or session mutation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolExecutionContext {
/// Provider/tool-call id for the call being executed.
pub call_id: String,
/// Worker-local identity shared by all tool calls from one execution batch.
pub batch_id: String,
/// Zero-based order of this call in the model-returned tool-call list.
pub call_index: usize,
}
impl ToolExecutionContext {
pub fn new(call_id: impl Into<String>, batch_id: impl Into<String>, call_index: usize) -> Self {
Self {
call_id: call_id.into(),
batch_id: batch_id.into(),
call_index,
}
}
/// Context for direct, non-worker calls in unit tests and low-level callers.
pub fn direct() -> Self {
Self::new("direct", "direct", 0)
}
}
impl Default for ToolExecutionContext {
fn default() -> Self {
Self::direct()
}
}
// =============================================================================
// Tool trait
// =============================================================================
@@ -219,16 +257,16 @@ pub type ToolDefinition = Arc<dyn Fn() -> (ToolMeta, Arc<dyn Tool>) + Send + Syn
/// # Manual Implementation
///
/// ```ignore
/// use llm_worker::tool::{Tool, ToolError, ToolMeta, ToolDefinition};
/// use llm_worker::tool::{Tool, ToolError, ToolExecutionContext, ToolMeta, ToolDefinition, ToolOutput};
/// use std::sync::Arc;
///
/// struct MyTool { counter: std::sync::atomic::AtomicUsize }
///
/// #[async_trait::async_trait]
/// impl Tool for MyTool {
/// async fn execute(&self, input: &str) -> Result<String, ToolError> {
/// async fn execute(&self, input: &str, ctx: ToolExecutionContext) -> Result<ToolOutput, ToolError> {
/// self.counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
/// Ok("result".to_string())
/// Ok(format!("call {}: {}", ctx.call_index, input).into())
/// }
/// }
///
@@ -247,11 +285,16 @@ pub trait Tool: Send + Sync {
///
/// # Arguments
/// * `input_json` - JSON-formatted arguments generated by LLM
/// * `ctx` - response-local call identity and ordering context
///
/// # Returns
/// A [`ToolOutput`] with summary and optional detailed content.
/// For simple cases, use `From<String>`: `Ok("done".to_string().into())`
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError>;
async fn execute(
&self,
input_json: &str,
ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError>;
}
// =============================================================================
crates/llm-worker/src/tool_server.rs
+50 -14
View File
@@ -4,7 +4,9 @@ use std::sync::{Arc, Mutex};
use thiserror::Error;
use crate::llm_client::ToolDefinition as LlmToolDefinition;
use crate::tool::{Tool, ToolDefinition as WorkerToolDefinition, ToolMeta, ToolOutput};
use crate::tool::{
Tool, ToolDefinition as WorkerToolDefinition, ToolExecutionContext, ToolMeta, ToolOutput,
};
type ToolMap = HashMap<String, (ToolMeta, Arc<dyn Tool>)>;
@@ -117,6 +119,7 @@ impl ToolServerHandle {
&self,
name: &str,
input_json: &str,
ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolServerError> {
let tool = {
let guard = self.tools.lock().unwrap_or_else(|e| e.into_inner());
@@ -125,7 +128,7 @@ impl ToolServerHandle {
.ok_or_else(|| ToolServerError::ToolNotFound(name.to_string()))?;
Arc::clone(tool)
};
tool.execute(input_json)
tool.execute(input_json, ctx)
.await
.map_err(|e| ToolServerError::ToolExecution(e.to_string()))
}
@@ -187,7 +190,11 @@ mod tests {
#[async_trait]
impl Tool for EchoTool {
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
input_json: &str,
_ctx: crate::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
Ok(input_json.to_string().into())
}
}
@@ -236,12 +243,15 @@ mod tests {
handle.register_tool(def("echo"));
handle.flush_pending();
let out = handle.call_tool("echo", r#"{"x":1}"#).await.expect("call");
let out = handle
.call_tool("echo", r#"{"x":1}"#, Default::default())
.await
.expect("call");
assert_eq!(out.summary, r#"{"x":1}"#);
assert!(out.content.is_none());
let err = handle
.call_tool("missing", "{}")
.call_tool("missing", "{}", Default::default())
.await
.expect_err("missing tool");
assert_eq!(err, ToolServerError::ToolNotFound("missing".to_string()));
@@ -298,7 +308,11 @@ mod tests {
#[async_trait]
impl Tool for FixedTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input_json: &str,
_ctx: crate::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
Ok("replaced".to_string().into())
}
}
@@ -327,7 +341,11 @@ mod tests {
#[async_trait]
impl Tool for ConstTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input_json: &str,
_ctx: crate::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
Ok("const".to_string().into())
}
}
@@ -342,7 +360,10 @@ mod tests {
});
handle.replace(replacement).expect("replace");
let out = handle.call_tool("echo", "{}").await.expect("call");
let out = handle
.call_tool("echo", "{}", Default::default())
.await
.expect("call");
assert_eq!(out.summary, "const");
}
@@ -360,7 +381,11 @@ mod tests {
#[async_trait]
impl Tool for GatedTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input_json: &str,
_ctx: crate::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
self.started.notify_one();
self.finish.notified().await;
Ok("done".to_string().into())
@@ -384,7 +409,7 @@ mod tests {
handle.flush_pending();
let h = handle.clone();
let call = tokio::spawn(async move { h.call_tool("slow", "{}").await });
let call = tokio::spawn(async move { h.call_tool("slow", "{}", Default::default()).await });
// Wait until the tool is actually executing.
started.notified().await;
@@ -413,7 +438,11 @@ mod tests {
#[async_trait]
impl Tool for OldTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input_json: &str,
_ctx: crate::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
self.started.notify_one();
self.finish.notified().await;
Ok("old".to_string().into())
@@ -437,7 +466,7 @@ mod tests {
handle.flush_pending();
let h = handle.clone();
let call = tokio::spawn(async move { h.call_tool("t", "{}").await });
let call = tokio::spawn(async move { h.call_tool("t", "{}", Default::default()).await });
// Wait until the old tool is mid-execution.
started.notified().await;
@@ -447,7 +476,11 @@ mod tests {
#[async_trait]
impl Tool for NewTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input_json: &str,
_ctx: crate::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
Ok("new".to_string().into())
}
}
@@ -469,7 +502,10 @@ mod tests {
assert_eq!(result.expect("call").summary, "old");
// New calls use the replacement.
let out = handle.call_tool("t", "{}").await.expect("call");
let out = handle
.call_tool("t", "{}", Default::default())
.await
.expect("call");
assert_eq!(out.summary, "new");
}
crates/llm-worker/src/worker.rs
+39 -12
View File
@@ -26,8 +26,8 @@ use crate::{
timeline::event::{ErrorEvent, StatusEvent, UsageEvent},
timeline::{TextBlockCollector, ThinkingBlockCollector, Timeline, ToolCallCollector},
tool::{
ToolCall, ToolDefinition as WorkerToolDefinition, ToolError, ToolOutputLimits, ToolResult,
truncate_content,
ToolCall, ToolDefinition as WorkerToolDefinition, ToolError, ToolExecutionContext,
ToolOutputLimits, ToolResult, truncate_content,
},
tool_server::{ToolServer, ToolServerHandle},
};
@@ -187,6 +187,10 @@ pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
/// LlmCall count (per-Worker running counter, monotonic). Unlike
/// `turn_count` this never collapses retries.
llm_call_count: usize,
/// Tool execution batch count (per-Worker running counter, monotonic).
/// Each batch corresponds to one collected assistant tool-call set or one
/// resumed pending tool-call set.
tool_execution_batch_count: usize,
/// Maximum number of AgentTurns (None = unlimited)
max_turns: Option<u32>,
/// AgentTurn-start callbacks (1:1 with LlmCall today)
@@ -912,19 +916,23 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
) -> Result<ToolExecutionResult, WorkerError> {
use futures::future::join_all;
// Map from tool call ID to (ToolCall, Meta, Tool)
// Map from tool call ID to (ToolCall, Meta, Tool, Context)
// Retained because it's needed for PostToolCall hooks
let mut call_info_map = HashMap::new();
let mut synthetic_results = Vec::new();
let batch_id = format!("tool-batch-{}", self.tool_execution_batch_count);
self.tool_execution_batch_count += 1;
// Phase 1: Apply pre_tool_call interceptor (determine skip/abort/synthetic result)
let mut approved_calls = Vec::new();
for mut tool_call in tool_calls {
for (call_index, mut tool_call) in tool_calls.into_iter().enumerate() {
let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index);
if let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) {
let mut info = ToolCallInfo {
call: tool_call.clone(),
meta,
tool,
context,
};
match self.interceptor.pre_tool_call(&mut info).await {
@@ -934,9 +942,11 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
PreToolAction::SyntheticResult(result) => {
let tool_call = info.call;
let mut context = info.context;
context.call_id = tool_call.id.clone();
call_info_map.insert(
tool_call.id.clone(),
(tool_call, info.meta.clone(), info.tool.clone()),
(tool_call, info.meta.clone(), info.tool.clone(), context),
);
synthetic_results.push(result);
continue;
@@ -953,26 +963,37 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
// Reflect changes made by interceptor
tool_call = info.call;
let mut context = info.context;
context.call_id = tool_call.id.clone();
call_info_map.insert(
tool_call.id.clone(),
(tool_call.clone(), info.meta.clone(), info.tool.clone()),
(
tool_call.clone(),
info.meta.clone(),
info.tool.clone(),
context.clone(),
),
);
approved_calls.push(tool_call);
approved_calls.push((tool_call, context));
} else {
// Unknown tools go into approved list as-is (will error at execution)
approved_calls.push(tool_call);
let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index);
approved_calls.push((tool_call, context));
}
}
// Phase 2: Execute approved tools in parallel (cancellable)
let futures: Vec<_> = approved_calls
.into_iter()
.map(|tool_call| {
.map(|(tool_call, context)| {
let tool_server = self.tool_server.clone();
async move {
let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default();
match tool_server.call_tool(&tool_call.name, &input_json).await {
match tool_server
.call_tool(&tool_call.name, &input_json, context)
.await
{
Ok(output) => ToolResult::from_output(&tool_call.id, output),
Err(e) => ToolResult::error(&tool_call.id, e.to_string()),
}
@@ -996,12 +1017,15 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
// Phase 3: Apply post_tool_call interceptor
for tool_result in &mut results {
if let Some((tool_call, meta, tool)) = call_info_map.get(&tool_result.tool_use_id) {
if let Some((tool_call, meta, tool, context)) =
call_info_map.get(&tool_result.tool_use_id)
{
let mut info = ToolResultInfo {
call: tool_call.clone(),
result: tool_result.clone(),
meta: meta.clone(),
tool: tool.clone(),
context: context.clone(),
};
match self.interceptor.post_tool_call(&mut info).await {
@@ -1026,7 +1050,7 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
let Some(content) = tool_result.content.as_mut() else {
continue;
};
let Some((tool_call, _, _)) = call_info_map.get(&tool_result.tool_use_id) else {
let Some((tool_call, _, _, _)) = call_info_map.get(&tool_result.tool_use_id) else {
continue;
};
let limit = limits.limit_for(&tool_call.name);
@@ -1628,6 +1652,7 @@ impl<C: LlmClient> Worker<C, Mutable> {
locked_prefix_len: 0,
turn_count: 0,
llm_call_count: 0,
tool_execution_batch_count: 0,
max_turns: None,
turn_start_cbs: Vec::new(),
turn_end_cbs: Vec::new(),
@@ -1892,6 +1917,7 @@ impl<C: LlmClient> Worker<C, Mutable> {
locked_prefix_len,
turn_count: self.turn_count,
llm_call_count: self.llm_call_count,
tool_execution_batch_count: self.tool_execution_batch_count,
max_turns: self.max_turns,
turn_start_cbs: self.turn_start_cbs,
turn_end_cbs: self.turn_end_cbs,
@@ -1984,6 +2010,7 @@ impl<C: LlmClient> Worker<C, Locked> {
locked_prefix_len: 0,
turn_count: self.turn_count,
llm_call_count: self.llm_call_count,
tool_execution_batch_count: self.tool_execution_batch_count,
max_turns: self.max_turns,
turn_start_cbs: self.turn_start_cbs,
turn_end_cbs: self.turn_end_cbs,