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,
crates/llm-worker/tests/callback_test.rs
+10 -2
View File
@@ -218,7 +218,11 @@ struct FixedOutputTool {
#[async_trait]
impl Tool for FixedOutputTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input_json: &str,
_ctx: llm_worker::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
Ok(self.output.clone())
}
}
@@ -289,7 +293,11 @@ struct ErroringTool {
#[async_trait]
impl Tool for ErroringTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input_json: &str,
_ctx: llm_worker::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
Err(ToolError::ExecutionFailed(self.message.clone()))
}
}
crates/llm-worker/tests/parallel_execution_test.rs
+257 -8
View File
@@ -2,8 +2,8 @@
//!
//! Verify that Worker executes multiple tools in parallel.
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use async_trait::async_trait;
@@ -12,7 +12,9 @@ use llm_worker::interceptor::{
Interceptor, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo,
};
use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent};
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput, ToolResult};
use llm_worker::tool::{
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
};
mod common;
use common::MockLlmClient;
@@ -59,13 +61,54 @@ impl SlowTool {
#[async_trait]
impl Tool for SlowTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input_json: &str,
_ctx: llm_worker::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
self.call_count.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(self.delay_ms)).await;
Ok(format!("Completed after {}ms", self.delay_ms).into())
}
}
#[derive(Clone)]
struct ContextRecordingTool {
name: String,
contexts: Arc<Mutex<Vec<ToolExecutionContext>>>,
}
impl ContextRecordingTool {
fn new(name: impl Into<String>, contexts: Arc<Mutex<Vec<ToolExecutionContext>>>) -> Self {
Self {
name: name.into(),
contexts,
}
}
fn definition(&self) -> ToolDefinition {
let tool = self.clone();
Arc::new(move || {
let meta = ToolMeta::new(&tool.name)
.description("Records tool execution context")
.input_schema(serde_json::json!({"type": "object"}));
(meta, Arc::new(tool.clone()) as Arc<dyn Tool>)
})
}
}
#[async_trait]
impl Tool for ContextRecordingTool {
async fn execute(
&self,
_input_json: &str,
ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
self.contexts.lock().unwrap().push(ctx);
Ok("recorded".to_string().into())
}
}
// =============================================================================
// Tests
// =============================================================================
@@ -92,10 +135,18 @@ async fn test_parallel_tool_execution() {
}),
];
let client = MockLlmClient::new(events);
let client = MockLlmClient::with_responses(vec![
events,
vec![
Event::text_block_start(0),
Event::text_delta(0, "Done"),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
]);
let mut worker = Worker::new(client);
// Each tool waits 100ms
let tool1 = SlowTool::new("slow_tool_1", 100);
let tool2 = SlowTool::new("slow_tool_2", 100);
let tool3 = SlowTool::new("slow_tool_3", 100);
@@ -129,7 +180,201 @@ async fn test_parallel_tool_execution() {
println!("Parallel execution completed in {:?}", elapsed);
}
/// Hook: pre_tool_call - verify that skipped tools are not executed
#[tokio::test]
async fn test_tool_execution_context_order_and_batch_id() {
let client = MockLlmClient::with_responses(vec![
vec![
Event::tool_use_start(0, "call_a", "record_a"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::tool_use_start(1, "call_b", "record_b"),
Event::tool_input_delta(1, r#"{}"#),
Event::tool_use_stop(1),
Event::tool_use_start(2, "call_c", "record_c"),
Event::tool_input_delta(2, r#"{}"#),
Event::tool_use_stop(2),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
vec![
Event::text_block_start(0),
Event::text_delta(0, "Done"),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
]);
let mut worker = Worker::new(client);
let contexts = Arc::new(Mutex::new(Vec::new()));
worker.register_tool(ContextRecordingTool::new("record_a", contexts.clone()).definition());
worker.register_tool(ContextRecordingTool::new("record_b", contexts.clone()).definition());
worker.register_tool(ContextRecordingTool::new("record_c", contexts.clone()).definition());
let _ = worker.run("record contexts").await;
let mut contexts = contexts.lock().unwrap().clone();
contexts.sort_by_key(|ctx| ctx.call_index);
assert_eq!(contexts.len(), 3);
assert_eq!(contexts[0].call_id, "call_a");
assert_eq!(contexts[0].call_index, 0);
assert_eq!(contexts[1].call_id, "call_b");
assert_eq!(contexts[1].call_index, 1);
assert_eq!(contexts[2].call_id, "call_c");
assert_eq!(contexts[2].call_index, 2);
assert_eq!(contexts[0].batch_id, contexts[1].batch_id);
assert_eq!(contexts[1].batch_id, contexts[2].batch_id);
}
#[tokio::test]
async fn test_tool_execution_context_batch_id_changes_between_batches() {
let client = MockLlmClient::with_responses(vec![
vec![
Event::tool_use_start(0, "call_first", "record"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
vec![
Event::tool_use_start(0, "call_second", "record"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
vec![
Event::text_block_start(0),
Event::text_delta(0, "Done"),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
]);
let mut worker = Worker::new(client);
let contexts = Arc::new(Mutex::new(Vec::new()));
worker.register_tool(ContextRecordingTool::new("record", contexts.clone()).definition());
let _ = worker.run("record batches").await;
let contexts = contexts.lock().unwrap().clone();
assert_eq!(contexts.len(), 2);
assert_eq!(contexts[0].call_id, "call_first");
assert_eq!(contexts[0].call_index, 0);
assert_eq!(contexts[1].call_id, "call_second");
assert_eq!(contexts[1].call_index, 0);
assert_ne!(contexts[0].batch_id, contexts[1].batch_id);
}
#[tokio::test]
async fn test_tool_execution_context_for_skipped_and_synthetic_paths() {
let client = MockLlmClient::with_responses(vec![
vec![
Event::tool_use_start(0, "call_run", "record"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::tool_use_start(1, "call_skip", "skip_tool"),
Event::tool_input_delta(1, r#"{}"#),
Event::tool_use_stop(1),
Event::tool_use_start(2, "call_synth", "synthetic_tool"),
Event::tool_input_delta(2, r#"{}"#),
Event::tool_use_stop(2),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
vec![
Event::text_block_start(0),
Event::text_delta(0, "Done"),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
]);
let mut worker = Worker::new(client);
let executed_contexts = Arc::new(Mutex::new(Vec::new()));
let pre_contexts = Arc::new(Mutex::new(Vec::new()));
let post_contexts = Arc::new(Mutex::new(Vec::new()));
worker
.register_tool(ContextRecordingTool::new("record", executed_contexts.clone()).definition());
worker.register_tool(
ContextRecordingTool::new("skip_tool", executed_contexts.clone()).definition(),
);
worker.register_tool(
ContextRecordingTool::new("synthetic_tool", executed_contexts.clone()).definition(),
);
struct ContextPolicy {
pre_contexts: Arc<Mutex<Vec<ToolExecutionContext>>>,
post_contexts: Arc<Mutex<Vec<ToolExecutionContext>>>,
}
#[async_trait]
impl Interceptor for ContextPolicy {
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction {
self.pre_contexts.lock().unwrap().push(info.context.clone());
match info.call.name.as_str() {
"skip_tool" => PreToolAction::Skip,
"synthetic_tool" => PreToolAction::SyntheticResult(ToolResult::from_output(
&info.call.id,
ToolOutput::from("synthetic result".to_string()),
)),
_ => PreToolAction::Continue,
}
}
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction {
self.post_contexts
.lock()
.unwrap()
.push(info.context.clone());
PostToolAction::Continue
}
}
worker.set_interceptor(ContextPolicy {
pre_contexts: pre_contexts.clone(),
post_contexts: post_contexts.clone(),
});
let _ = worker.run("record skipped and synthetic contexts").await;
let mut pre_contexts = pre_contexts.lock().unwrap().clone();
pre_contexts.sort_by_key(|ctx| ctx.call_index);
assert_eq!(pre_contexts.len(), 3);
assert_eq!(pre_contexts[0].call_id, "call_run");
assert_eq!(pre_contexts[0].call_index, 0);
assert_eq!(pre_contexts[1].call_id, "call_skip");
assert_eq!(pre_contexts[1].call_index, 1);
assert_eq!(pre_contexts[2].call_id, "call_synth");
assert_eq!(pre_contexts[2].call_index, 2);
assert_eq!(pre_contexts[0].batch_id, pre_contexts[1].batch_id);
assert_eq!(pre_contexts[1].batch_id, pre_contexts[2].batch_id);
let executed_contexts = executed_contexts.lock().unwrap().clone();
assert_eq!(executed_contexts.len(), 1);
assert_eq!(executed_contexts[0].call_id, "call_run");
assert_eq!(executed_contexts[0].call_index, 0);
let mut post_contexts = post_contexts.lock().unwrap().clone();
post_contexts.sort_by_key(|ctx| ctx.call_index);
assert_eq!(post_contexts.len(), 2);
assert_eq!(post_contexts[0].call_id, "call_run");
assert_eq!(post_contexts[0].call_index, 0);
assert_eq!(post_contexts[1].call_id, "call_synth");
assert_eq!(post_contexts[1].call_index, 2);
assert_eq!(post_contexts[0].batch_id, post_contexts[1].batch_id);
}
#[tokio::test]
async fn test_before_tool_call_skip() {
let events = vec![
@@ -220,7 +465,11 @@ async fn test_post_tool_call_modification() {
#[async_trait]
impl Tool for SimpleTool {
async fn execute(&self, _: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_: &str,
_ctx: llm_worker::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
Ok("Original Result".to_string().into())
}
}
crates/llm-worker/tests/tool_macro_test.rs
+43 -9
View File
@@ -9,6 +9,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use schemars;
use serde;
use llm_worker::ToolExecutionContext;
use llm_worker_macros::tool_registry;
// =============================================================================
@@ -42,6 +43,15 @@ impl SimpleContext {
async fn get_prefix(&self) -> String {
self.prefix.clone()
}
/// Tool that observes execution context
#[tool]
async fn context_echo(&self, ctx: ToolExecutionContext, message: String) -> String {
format!(
"{}:{}:{}:{}",
ctx.batch_id, ctx.call_index, ctx.call_id, message
)
}
}
#[tokio::test]
@@ -74,7 +84,9 @@ async fn test_basic_tool_generation() {
);
// Execution test
let result = tool.execute(r#"{"message": "World"}"#).await;
let result = tool
.execute(r#"{"message": "World"}"#, Default::default())
.await;
assert!(result.is_ok(), "Should execute successfully");
let output = result.unwrap();
assert!(
@@ -97,7 +109,9 @@ async fn test_multiple_arguments() {
assert_eq!(meta.name, "add");
let result = tool.execute(r#"{"a": 10, "b": 20}"#).await;
let result = tool
.execute(r#"{"a": 10, "b": 20}"#, Default::default())
.await;
assert!(result.is_ok());
let output = result.unwrap();
assert!(
@@ -118,7 +132,7 @@ async fn test_no_arguments() {
assert_eq!(meta.name, "get_prefix");
// Call with empty JSON object
let result = tool.execute(r#"{}"#).await;
let result = tool.execute(r#"{}"#, Default::default()).await;
assert!(result.is_ok());
let output = result.unwrap();
assert!(
@@ -137,7 +151,9 @@ async fn test_invalid_arguments() {
let (_, tool) = ctx.greet_definition()();
// Invalid JSON
let result = tool.execute(r#"{"wrong_field": "value"}"#).await;
let result = tool
.execute(r#"{"wrong_field": "value"}"#, Default::default())
.await;
assert!(result.is_err(), "Should fail with invalid arguments");
}
@@ -175,7 +191,7 @@ async fn test_result_return_type_success() {
let ctx = FallibleContext;
let (_, tool) = ctx.validate_definition()();
let result = tool.execute(r#"{"value": 42}"#).await;
let result = tool.execute(r#"{"value": 42}"#, Default::default()).await;
assert!(result.is_ok(), "Should succeed for positive value");
let output = result.unwrap();
assert!(
@@ -190,7 +206,7 @@ async fn test_result_return_type_error() {
let ctx = FallibleContext;
let (_, tool) = ctx.validate_definition()();
let result = tool.execute(r#"{"value": -1}"#).await;
let result = tool.execute(r#"{"value": -1}"#, Default::default()).await;
assert!(result.is_err(), "Should fail for negative value");
let err = result.unwrap_err();
@@ -228,9 +244,9 @@ async fn test_sync_method() {
let (_, tool) = ctx.increment_definition()();
// Execute 3 times
let result1 = tool.execute(r#"{}"#).await;
let result2 = tool.execute(r#"{}"#).await;
let result3 = tool.execute(r#"{}"#).await;
let result1 = tool.execute(r#"{}"#, Default::default()).await;
let result2 = tool.execute(r#"{}"#, Default::default()).await;
let result3 = tool.execute(r#"{}"#, Default::default()).await;
assert!(result1.is_ok());
assert!(result2.is_ok());
@@ -240,6 +256,24 @@ async fn test_sync_method() {
assert_eq!(ctx.counter.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_tool_macro_passes_execution_context() {
let ctx = SimpleContext {
prefix: "Test".to_string(),
};
let (_, tool) = ctx.context_echo_definition()();
let output = tool
.execute(
r#"{"message":"hello"}"#,
ToolExecutionContext::new("call-ctx", "batch-ctx", 7),
)
.await
.unwrap();
assert_eq!(output.summary, "\"batch-ctx:7:call-ctx:hello\"");
}
// =============================================================================
// Test: ToolMeta Immutability
// =============================================================================
crates/llm-worker/tests/worker_fixtures.rs
+5 -1
View File
@@ -58,7 +58,11 @@ impl MockWeatherTool {
#[async_trait]
impl Tool for MockWeatherTool {
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
input_json: &str,
_ctx: llm_worker::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
self.call_count.fetch_add(1, Ordering::SeqCst);
// Parse input
crates/llm-worker/tests/worker_state_test.rs
+5 -1
View File
@@ -136,7 +136,11 @@ impl CountingTool {
#[async_trait]
impl Tool for CountingTool {
async fn execute(&self, _input_json: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input_json: &str,
_ctx: llm_worker::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(format!("{}-ok", self.name).into())
}