tool: add execution context

This commit is contained in:
2026-06-09 19:31:11 +09:00
parent b21fab82fc
commit d8aed7befe
39 changed files with 1212 additions and 259 deletions
+33 -13
View File
@@ -151,7 +151,11 @@ struct SearchSessionLogTool {
#[async_trait]
impl Tool for SearchSessionLogTool {
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> {
let params: SearchSessionParams = serde_json::from_str(input_json).map_err(|e| {
ToolError::InvalidArgument(format!("invalid search_session_log input: {e}"))
})?;
@@ -206,7 +210,11 @@ struct ReadSessionItemsTool {
#[async_trait]
impl Tool for ReadSessionItemsTool {
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> {
let params: ReadSessionParams = serde_json::from_str(input_json).map_err(|e| {
ToolError::InvalidArgument(format!("invalid read_session_items input: {e}"))
})?;
@@ -368,7 +376,11 @@ struct MarkReadRequiredTool {
#[async_trait]
impl Tool for MarkReadRequiredTool {
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> {
let params: MarkParams = serde_json::from_str(input_json).map_err(|e| {
ToolError::InvalidArgument(format!("invalid mark_read_required input: {e}"))
})?;
@@ -425,7 +437,11 @@ struct AddReferenceTool {
#[async_trait]
impl Tool for AddReferenceTool {
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> {
let params: ReferenceParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid add_reference input: {e}")))?;
let mut guard = self.ctx.lock().expect("compact worker context poisoned");
@@ -449,7 +465,11 @@ struct WriteSummaryTool {
#[async_trait]
impl Tool for WriteSummaryTool {
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> {
let params: SummaryParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid write_summary input: {e}")))?;
let mut guard = self.ctx.lock().expect("compact worker context poisoned");
@@ -749,7 +769,7 @@ mod tests {
ctx: ctx.clone(),
});
let input = serde_json::json!({ "file_path": path.to_str().unwrap() }).to_string();
let out = tool.execute(&input).await.unwrap();
let out = tool.execute(&input, Default::default()).await.unwrap();
assert!(out.summary.starts_with("Marked"));
let guard = ctx.lock().unwrap();
@@ -770,7 +790,7 @@ mod tests {
ctx: ctx.clone(),
});
let input = serde_json::json!({ "file_path": path.to_str().unwrap() }).to_string();
let res = tool.execute(&input).await;
let res = tool.execute(&input, Default::default()).await;
assert!(matches!(res, Err(ToolError::ExecutionFailed(_))));
let guard = ctx.lock().unwrap();
@@ -784,11 +804,11 @@ mod tests {
let tool: Arc<dyn Tool> = Arc::new(WriteSummaryTool { ctx: ctx.clone() });
let first = serde_json::json!({ "text": "first" }).to_string();
let out1 = tool.execute(&first).await.unwrap();
let out1 = tool.execute(&first, Default::default()).await.unwrap();
assert!(out1.summary.contains("recorded"));
let second = serde_json::json!({ "text": "second" }).to_string();
let out2 = tool.execute(&second).await.unwrap();
let out2 = tool.execute(&second, Default::default()).await.unwrap();
assert!(out2.summary.contains("replaced"));
assert_eq!(ctx.lock().unwrap().summary.as_deref(), Some("second"));
@@ -801,8 +821,8 @@ mod tests {
let p = "/abs/path.rs";
let input = serde_json::json!({ "file_path": p }).to_string();
tool.execute(&input).await.unwrap();
tool.execute(&input).await.unwrap();
tool.execute(&input, Default::default()).await.unwrap();
tool.execute(&input, Default::default()).await.unwrap();
let guard = ctx.lock().unwrap();
assert_eq!(guard.references.len(), 1);
@@ -823,7 +843,7 @@ mod tests {
state: Arc::new(SessionLogToolState { items }),
});
let input = serde_json::json!({ "query": "compact", "limit": 10 }).to_string();
let out = tool.execute(&input).await.unwrap();
let out = tool.execute(&input, Default::default()).await.unwrap();
let content = out.content.unwrap();
assert!(content.contains("investigate compact failure"));
@@ -842,7 +862,7 @@ mod tests {
state: Arc::new(SessionLogToolState { items }),
});
let input = serde_json::json!({ "offset": 0, "limit": 1, "mode": "full" }).to_string();
let out = tool.execute(&input).await.unwrap();
let out = tool.execute(&input, Default::default()).await.unwrap();
let content = out.content.unwrap();
assert!(content.contains("raw trace detail"));
+16 -4
View File
@@ -752,7 +752,11 @@ impl<St> Tool for ListPodsTool<St>
where
St: PodMetadataStore + Clone + Send + Sync + 'static,
{
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> {
let items = self
.discovery
.list_visible()
@@ -775,7 +779,11 @@ impl<St> Tool for RestorePodTool<St>
where
St: PodMetadataStore + Clone + Send + Sync + 'static,
{
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> {
let input: PodNameInput = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid RestorePod input: {e}")))?;
let result = self
@@ -847,7 +855,11 @@ impl<St> Tool for SendToPeerPodTool<St>
where
St: PodMetadataStore + Clone + Send + Sync + 'static,
{
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> {
let input: SendToPeerPodInput = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid SendToPeerPod input: {e}")))?;
let detail = self
@@ -1392,7 +1404,7 @@ mod tests {
let (_, tool) = send_to_peer_pod_tool(discovery)();
let output = tool
.execute(r#"{"name":"target","message":"hello"}"#)
.execute(r#"{"name":"target","message":"hello"}"#, Default::default())
.await
.unwrap();
assert_eq!(output.summary, "sent peer message to `target`");
+5 -1
View File
@@ -1292,7 +1292,11 @@ mod tests {
#[async_trait]
impl Tool for DummyTool {
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(ToolOutput::from("ok".to_string()))
}
}
@@ -73,7 +73,11 @@ step: leave the task as-is, summarize the problem to the user, and end the turn.
#[async_trait]
impl Tool for TaskCreateTool {
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> {
let params: TaskCreateParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid TaskCreate input: {e}")))?;
let created = self.store.create(params.subject, params.description);
@@ -93,7 +97,11 @@ impl Tool for TaskCreateTool {
#[async_trait]
impl Tool for TaskListTool {
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> {
let _: TaskListParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid TaskList input: {e}")))?;
let tasks = self.store.list();
@@ -106,7 +114,11 @@ impl Tool for TaskListTool {
#[async_trait]
impl Tool for TaskGetTool {
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> {
let params: TaskGetParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid TaskGet input: {e}")))?;
let task = self.store.get(params.taskid).ok_or_else(|| {
@@ -122,7 +134,11 @@ impl Tool for TaskGetTool {
#[async_trait]
impl Tool for TaskUpdateTool {
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> {
let params: TaskUpdateParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid TaskUpdate input: {e}")))?;
let updated = self
@@ -241,14 +257,20 @@ mod tests {
let update = tool(task_update_tool(store.clone()));
let out = create
.execute(r#"{"subject":"implement","description":"write code"}"#)
.execute(
r#"{"subject":"implement","description":"write code"}"#,
Default::default(),
)
.await
.unwrap();
assert!(out.summary.contains("Created task 1"));
assert_eq!(store.get(1).unwrap().status, TaskStatus::Pending);
let out = update
.execute(r#"{"taskid":1,"status":"inprogress","subject":"implement tasks"}"#)
.execute(
r#"{"taskid":1,"status":"inprogress","subject":"implement tasks"}"#,
Default::default(),
)
.await
.unwrap();
assert!(out.summary.contains("Updated task 1"));
@@ -256,11 +278,14 @@ mod tests {
assert_eq!(task.status, TaskStatus::Inprogress);
assert_eq!(task.subject, "implement tasks");
let out = get.execute(r#"{"taskid":1}"#).await.unwrap();
let out = get
.execute(r#"{"taskid":1}"#, Default::default())
.await
.unwrap();
assert!(out.summary.contains("Task 1 (inprogress)"));
assert!(out.content.unwrap().contains("implement tasks"));
let out = list.execute("{}").await.unwrap();
let out = list.execute("{}", Default::default()).await.unwrap();
assert!(out.summary.contains("1 task(s)"));
let content = out.content.unwrap();
assert!(content.contains("\"taskid\": 1"));
@@ -273,11 +298,14 @@ mod tests {
store.create("s".into(), "d".into());
let update = tool(task_update_tool(store));
let err = update.execute(r#"{"taskid":1}"#).await.unwrap_err();
let err = update
.execute(r#"{"taskid":1}"#, Default::default())
.await
.unwrap_err();
assert!(err.to_string().contains("at least one"));
let err = update
.execute(r#"{"taskid":99,"status":"deleted"}"#)
.execute(r#"{"taskid":99,"status":"deleted"}"#, Default::default())
.await
.unwrap_err();
assert!(err.to_string().contains("taskid 99 not found"));
+2
View File
@@ -491,6 +491,7 @@ mod tests {
},
meta,
tool,
context: llm_worker::tool::ToolExecutionContext::new("call-id", "test-batch", 0),
}
}
@@ -898,6 +899,7 @@ mod tests {
),
meta: info.meta,
tool: info.tool,
context: info.context,
};
let action = interceptor.post_tool_call(&mut result_info).await;
+15 -3
View File
@@ -62,7 +62,11 @@ struct SendToPodTool {
#[async_trait]
impl Tool for SendToPodTool {
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> {
let input: SendToPodInput = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid SendToPod input: {e}")))?;
let record = self
@@ -123,7 +127,11 @@ struct ReadPodOutputTool {
#[async_trait]
impl Tool for ReadPodOutputTool {
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> {
let input: NameInput = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid ReadPodOutput input: {e}")))?;
let record = self
@@ -197,7 +205,11 @@ struct StopPodTool {
#[async_trait]
impl Tool for StopPodTool {
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> {
let input: NameInput = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid StopPod input: {e}")))?;
let record = self
+5 -1
View File
@@ -298,7 +298,11 @@ impl SpawnPodTool {
#[async_trait]
impl Tool for SpawnPodTool {
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> {
let input: SpawnPodInput = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid SpawnPod input: {e}")))?;
+5 -1
View File
@@ -1351,7 +1351,11 @@ struct HangingTool;
#[async_trait]
impl Tool for HangingTool {
async fn execute(&self, _input: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input: &str,
_ctx: llm_worker::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
std::future::pending::<()>().await;
unreachable!()
}
+10 -10
View File
@@ -262,7 +262,7 @@ async fn send_to_pod_delivers_run_method() {
let def = send_to_pod_tool(registry);
let (_meta, tool) = def();
let input = json!({ "name": "child", "message": "hello there" }).to_string();
let output: ToolOutput = tool.execute(&input).await.unwrap();
let output: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(
output.summary.contains("child"),
"summary: {}",
@@ -285,7 +285,7 @@ async fn send_to_pod_errors_on_unknown_pod() {
let def = send_to_pod_tool(registry);
let (_meta, tool) = def();
let input = json!({ "name": "nope", "message": "hi" }).to_string();
let err = tool.execute(&input).await.unwrap_err();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
assert!(err.to_string().contains("no spawned pod"), "{err}");
}
@@ -307,7 +307,7 @@ async fn send_to_pod_errors_when_pod_already_running() {
let def = send_to_pod_tool(registry);
let (_meta, tool) = def();
let input = json!({ "name": "child", "message": "hi" }).to_string();
let err = tool.execute(&input).await.unwrap_err();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
assert!(
err.to_string().contains("already running"),
"expected AlreadyRunning wording: {err}"
@@ -341,13 +341,13 @@ async fn read_pod_output_returns_new_assistant_text_then_empty_on_second_call()
let (_meta, tool) = def();
let input = json!({ "name": "child" }).to_string();
let first: ToolOutput = tool.execute(&input).await.unwrap();
let first: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
let body = first.content.expect("first read should have content");
assert!(body.contains("hi back"), "body: {body}");
assert!(body.contains("still working"), "body: {body}");
// Cursor now points past all items — second call returns no new text.
let second: ToolOutput = tool.execute(&input).await.unwrap();
let second: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(
second.content.is_none(),
"unexpected content: {:?}",
@@ -371,7 +371,7 @@ async fn read_pod_output_reports_stopped_on_dead_socket() {
let def = read_pod_output_tool(registry);
let (_meta, tool) = def();
let input = json!({ "name": "child" }).to_string();
let output: ToolOutput = tool.execute(&input).await.unwrap();
let output: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(output.summary.contains("stopped"), "{}", output.summary);
}
@@ -452,7 +452,7 @@ async fn stop_pod_sends_shutdown_and_releases_scope() {
let def = stop_pod_tool(registry.clone());
let (_meta, tool) = def();
let input = json!({ "name": "child" }).to_string();
let output: ToolOutput = tool.execute(&input).await.unwrap();
let output: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(output.summary.contains("stopped"), "{}", output.summary);
// The child got a Shutdown.
@@ -497,7 +497,7 @@ async fn stop_pod_succeeds_even_when_child_unreachable() {
let def = stop_pod_tool(registry.clone());
let (_meta, tool) = def();
let input = json!({ "name": "child" }).to_string();
let output: ToolOutput = tool.execute(&input).await.unwrap();
let output: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(output.summary.contains("stopped"), "{}", output.summary);
// Registry no longer knows about the child.
@@ -545,7 +545,7 @@ async fn restored_registry_uses_pod_state_without_runtime_file() {
let def = send_to_pod_tool(restored.clone());
let (_meta, tool) = def();
let input = json!({ "name": "child", "message": "after restart" }).to_string();
tool.execute(&input).await.unwrap();
tool.execute(&input, Default::default()).await.unwrap();
match received.recv().await.expect("expected Run") {
Method::Run { input } => match input.as_slice() {
[protocol::Segment::Text { content }] => assert_eq!(content, "after restart"),
@@ -556,7 +556,7 @@ async fn restored_registry_uses_pod_state_without_runtime_file() {
let def = stop_pod_tool(restored.clone());
let (_meta, tool) = def();
tool.execute(&json!({ "name": "child" }).to_string())
tool.execute(&json!({ "name": "child" }).to_string(), Default::default())
.await
.unwrap();
assert!(matches!(
+5 -1
View File
@@ -79,7 +79,11 @@ struct BigContentTool {
#[async_trait]
impl Tool for BigContentTool {
async fn execute(&self, _input: &str) -> Result<ToolOutput, ToolError> {
async fn execute(
&self,
_input: &str,
_ctx: llm_worker::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
Ok(ToolOutput {
summary: self.summary.into(),
content: Some(self.content.clone()),
+7 -7
View File
@@ -312,7 +312,7 @@ async fn spawn_pod_launches_runtime_in_workspace_and_passes_tool_cwd() {
})
.to_string();
tool.execute(&input).await.unwrap();
tool.execute(&input, Default::default()).await.unwrap();
assert!(matches!(received.await.unwrap(), Some(Method::Run { .. })));
let invocation = read_recorded_runtime_invocation(&output_path).await;
assert_eq!(invocation[0], allow_root.path().to_str().unwrap());
@@ -373,7 +373,7 @@ async fn spawn_pod_omitted_cwd_preserves_spawner_pwd() {
})
.to_string();
tool.execute(&input).await.unwrap();
tool.execute(&input, Default::default()).await.unwrap();
assert!(matches!(received.await.unwrap(), Some(Method::Run { .. })));
let invocation = read_recorded_runtime_invocation(&output_path).await;
assert_eq!(invocation[0], allow_root.path().to_str().unwrap());
@@ -433,7 +433,7 @@ async fn spawn_pod_delegates_scope_and_sends_run() {
.is_writable(&allow_root.path().join("a.txt"))
);
let output: ToolOutput = tool.execute(&input).await.unwrap();
let output: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(
output.summary.contains("child"),
"summary: {}",
@@ -519,7 +519,7 @@ async fn spawn_pod_requires_explicit_delegation_even_with_direct_scope() {
})
.to_string();
let err = tool.execute(&input).await.unwrap_err();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
match err {
ToolError::InvalidArgument(message) => {
assert!(message.contains("no delegation scope grant"), "{message}");
@@ -587,7 +587,7 @@ async fn spawn_pod_rejects_child_non_recursive_scope_under_parent_non_recursive_
})
.to_string();
let err = tool.execute(&input).await.unwrap_err();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
match err {
ToolError::InvalidArgument(message) => {
assert!(
@@ -639,7 +639,7 @@ async fn spawn_pod_rejects_scope_outside_spawner() {
})
.to_string();
let err = tool.execute(&input).await.unwrap_err();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
match err {
ToolError::InvalidArgument(msg) => {
assert!(
@@ -712,7 +712,7 @@ async fn spawn_pod_rolls_back_reservation_when_socket_never_appears() {
})
.to_string();
let err = tool.execute(&input).await.unwrap_err();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
match err {
ToolError::ExecutionFailed(msg) => {
assert!(