From 2a2a61065481f969ee057b35dba5503429d3cd1d Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Mon, 23 Feb 2026 18:10:10 -0500 Subject: [PATCH] Replace abort_flag with CancellationToken for abort-aware process cancellation Thread CancellationToken into tool executors and exec_command so that running processes are killed (SIGTERM -> 2s -> SIGKILL) when abort fires, rather than only checking the flag between LLM calls. Key changes: - ToolExecutor type gains CancellationToken parameter - ExecutionEnvironment::exec_command gains cancel_token param - LocalExecutionEnvironment uses tokio::select! (completion vs timeout vs cancellation) with extracted sigterm_then_kill helper - DockerExecutionEnvironment uses same select! pattern - Session replaces Arc with CancellationToken, passes child_token() per tool call - Shell tool forwards cancel token to exec_command - CLI SIGINT handler calls cancel_token.cancel() Co-Authored-By: Claude Opus 4.6 --- Cargo.lock | 1 + crates/agent/Cargo.toml | 1 + crates/agent/README.md | 4 +- crates/agent/src/cli.rs | 5 +- crates/agent/src/docker_env.rs | 35 ++++++++---- crates/agent/src/execution_env.rs | 4 +- crates/agent/src/local_env.rs | 83 ++++++++++++++++------------- crates/agent/src/profiles/openai.rs | 2 +- crates/agent/src/session.rs | 78 ++++++++++++++++++--------- crates/agent/src/subagent.rs | 18 +++---- crates/agent/src/test_support.rs | 7 ++- crates/agent/src/tool_registry.rs | 10 ++-- crates/agent/src/tools.rs | 42 +++++++++------ 13 files changed, 176 insertions(+), 114 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 87c73adaf..92e8f5cd8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -23,6 +23,7 @@ dependencies = [ "terminal", "thiserror 2.0.18", "tokio", + "tokio-util", "uuid", ] diff --git a/crates/agent/Cargo.toml b/crates/agent/Cargo.toml index b3e2efbaa..cd74f61a5 100644 --- a/crates/agent/Cargo.toml +++ b/crates/agent/Cargo.toml @@ -35,6 +35,7 @@ futures.workspace = true async-trait.workspace = true jsonschema.workspace = true chrono.workspace = true +tokio-util.workspace = true glob = "0.3" bollard = { workspace = true, optional = true } tar = { workspace = true, optional = true } diff --git a/crates/agent/README.md b/crates/agent/README.md index 33657f713..36a04a0da 100644 --- a/crates/agent/README.md +++ b/crates/agent/README.md @@ -173,9 +173,9 @@ session.follow_up("Now run the test suite to verify".into()); Cancel a running session from another thread: ```rust -let abort_flag = session.abort_flag_handle(); +let cancel_token = session.cancel_token(); // From another task: -abort_flag.store(true, std::sync::atomic::Ordering::SeqCst); +cancel_token.cancel(); ``` ### Custom Tools diff --git a/crates/agent/src/cli.rs b/crates/agent/src/cli.rs index 6acbdc147..db13d0ecb 100644 --- a/crates/agent/src/cli.rs +++ b/crates/agent/src/cli.rs @@ -6,7 +6,6 @@ use clap::{Parser, ValueEnum}; use llm::client::Client; use std::io::{IsTerminal, Write}; use std::path::PathBuf; -use std::sync::atomic::Ordering; use std::sync::{Arc, Mutex}; use terminal::Styles; @@ -337,10 +336,10 @@ pub async fn run() -> anyhow::Result<()> { let mut session = Session::new(client, profile, env, config); // SIGINT handler - let abort_flag = session.abort_flag_handle(); + let cancel_token = session.cancel_token(); tokio::spawn(async move { tokio::signal::ctrl_c().await.ok(); - abort_flag.store(true, Ordering::SeqCst); + cancel_token.cancel(); }); // Subscribe to events for real-time tool status on stderr diff --git a/crates/agent/src/docker_env.rs b/crates/agent/src/docker_env.rs index 593928ae7..1908f06ad 100644 --- a/crates/agent/src/docker_env.rs +++ b/crates/agent/src/docker_env.rs @@ -1,5 +1,6 @@ use crate::execution_env::{format_lines_numbered, DirEntry, ExecResult, ExecutionEnvironment, GrepOptions}; use async_trait::async_trait; +use tokio_util::sync::CancellationToken; use bollard::container::{ Config, CreateContainerOptions, RemoveContainerOptions, StartContainerOptions, StopContainerOptions, UploadToContainerOptions, @@ -153,13 +154,14 @@ impl DockerExecutionEnvironment { Ok((stdout, stderr, exit_code)) } - /// Runs a shell command inside the container with timeout support. + /// Runs a shell command inside the container with timeout and cancellation support. async fn docker_exec_shell( &self, command: &str, timeout_ms: u64, working_dir: Option<&str>, env_vars: Option<&HashMap>, + cancel_token: Option, ) -> Result { let start = Instant::now(); @@ -180,10 +182,10 @@ impl DockerExecutionEnvironment { ]; let timeout_duration = std::time::Duration::from_millis(timeout_ms); - let exec_future = self.docker_exec(cmd, Some(&effective_dir), env); + let token = cancel_token.unwrap_or_default(); - match tokio::time::timeout(timeout_duration, exec_future).await { - Ok(result) => { + tokio::select! { + result = self.docker_exec(cmd, Some(&effective_dir), env) => { let (stdout, stderr, exit_code) = result?; let duration_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX); Ok(ExecResult { @@ -194,7 +196,7 @@ impl DockerExecutionEnvironment { duration_ms, }) } - Err(_) => { + () = tokio::time::sleep(timeout_duration) => { let duration_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX); Ok(ExecResult { stdout: String::new(), @@ -204,6 +206,16 @@ impl DockerExecutionEnvironment { duration_ms, }) } + () = token.cancelled() => { + let duration_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX); + Ok(ExecResult { + stdout: String::new(), + stderr: "Command cancelled".to_string(), + exit_code: -1, + timed_out: true, + duration_ms, + }) + } } } @@ -350,9 +362,10 @@ impl ExecutionEnvironment for DockerExecutionEnvironment { timeout_ms: u64, working_dir: Option<&str>, env_vars: Option<&HashMap>, + cancel_token: Option, ) -> Result { let dir = working_dir.map(|d| self.resolve_container_path(d)); - self.docker_exec_shell(command, timeout_ms, dir.as_deref(), env_vars) + self.docker_exec_shell(command, timeout_ms, dir.as_deref(), env_vars, cancel_token) .await } @@ -583,7 +596,7 @@ impl ExecutionEnvironment for DockerExecutionEnvironment { // Run through shell so that quoting works correctly let result = self - .docker_exec_shell(&command, 30_000, None, None) + .docker_exec_shell(&command, 30_000, None, None, None) .await?; let results: Vec = result @@ -613,7 +626,7 @@ impl ExecutionEnvironment for DockerExecutionEnvironment { ); let result = self - .docker_exec_shell(&script, 30_000, None, None) + .docker_exec_shell(&script, 30_000, None, None, None) .await?; let results: Vec = result @@ -679,7 +692,7 @@ mod tests { assert!(env.os_version().starts_with("linux ")); // exec_command - let result = env.exec_command("echo hello", 5000, None, None).await.unwrap(); + let result = env.exec_command("echo hello", 5000, None, None, None).await.unwrap(); assert_eq!(result.stdout.trim(), "hello"); assert_eq!(result.exit_code, 0); assert!(!result.timed_out); @@ -728,7 +741,7 @@ mod tests { let env = DockerExecutionEnvironment::new(config).unwrap(); env.initialize().await.unwrap(); - let result = env.exec_command("sleep 60", 1000, None, None).await.unwrap(); + let result = env.exec_command("sleep 60", 1000, None, None, None).await.unwrap(); assert!(result.timed_out); assert_eq!(result.exit_code, -1); @@ -751,7 +764,7 @@ mod tests { env.write_file("special.txt", content).await.unwrap(); // Read raw content back via cat to verify exact match - let result = env.exec_command("cat /workspace/special.txt", 5000, None, None).await.unwrap(); + let result = env.exec_command("cat /workspace/special.txt", 5000, None, None, None).await.unwrap(); assert_eq!(result.stdout, content); env.cleanup().await.unwrap(); diff --git a/crates/agent/src/execution_env.rs b/crates/agent/src/execution_env.rs index f80601253..28168197e 100644 --- a/crates/agent/src/execution_env.rs +++ b/crates/agent/src/execution_env.rs @@ -1,5 +1,6 @@ use async_trait::async_trait; use std::fmt::Write; +use tokio_util::sync::CancellationToken; /// Formats file content with line numbers for display. /// @@ -55,6 +56,7 @@ pub trait ExecutionEnvironment: Send + Sync { timeout_ms: u64, working_dir: Option<&str>, env_vars: Option<&std::collections::HashMap>, + cancel_token: Option, ) -> Result; async fn grep( &self, @@ -92,7 +94,7 @@ mod tests { #[tokio::test] async fn mock_env_exec_command() { let env: Arc = Arc::new(MockExecutionEnvironment::default()); - let result = env.exec_command("echo", 5000, None, None).await.unwrap(); + let result = env.exec_command("echo", 5000, None, None, None).await.unwrap(); assert_eq!(result.exit_code, 0); assert!(!result.timed_out); } diff --git a/crates/agent/src/local_env.rs b/crates/agent/src/local_env.rs index 90919c3d4..c12504cbe 100644 --- a/crates/agent/src/local_env.rs +++ b/crates/agent/src/local_env.rs @@ -4,6 +4,7 @@ use std::path::{Path, PathBuf}; use std::time::Instant; use tokio::io::AsyncReadExt; use tokio::process::Command; +use tokio_util::sync::CancellationToken; pub struct LocalExecutionEnvironment { working_directory: PathBuf, @@ -131,6 +132,7 @@ impl ExecutionEnvironment for LocalExecutionEnvironment { timeout_ms: u64, working_dir: Option<&str>, env_vars: Option<&std::collections::HashMap>, + cancel_token: Option, ) -> Result { let start = Instant::now(); @@ -171,43 +173,22 @@ impl ExecutionEnvironment for LocalExecutionEnvironment { .map_err(|e| format!("Failed to spawn command: {e}"))?; let timeout_duration = std::time::Duration::from_millis(timeout_ms); + let token = cancel_token.unwrap_or_default(); - let (timed_out, exit_code) = - if let Ok(status_result) = tokio::time::timeout(timeout_duration, child.wait()).await { - let status = - status_result.map_err(|e| format!("Failed to wait for process: {e}"))?; + let (timed_out, exit_code) = tokio::select! { + status_result = child.wait() => { + let status = status_result.map_err(|e| format!("Failed to wait for process: {e}"))?; (false, status.code().unwrap_or(-1)) - } else { - // SIGTERM the process group first, then SIGKILL after 2 seconds - #[cfg(unix)] - if let Some(pid) = child.id() { - #[allow(clippy::cast_possible_wrap)] - unsafe { - // Negative pid sends signal to the entire process group - libc::kill(-(pid as i32), libc::SIGTERM); - } - // Wait 2 seconds for graceful shutdown - if tokio::time::timeout( - std::time::Duration::from_secs(2), - child.wait(), - ) - .await - .is_err() - { - let _ = child.kill().await; - let _ = child.wait().await; - } - } else { - let _ = child.kill().await; - let _ = child.wait().await; - } - #[cfg(not(unix))] - { - let _ = child.kill().await; - let _ = child.wait().await; - } + } + () = tokio::time::sleep(timeout_duration) => { + sigterm_then_kill(&mut child).await; (true, -1) - }; + } + () = token.cancelled() => { + sigterm_then_kill(&mut child).await; + (true, -1) + } + }; let duration_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX); @@ -372,6 +353,32 @@ impl ExecutionEnvironment for LocalExecutionEnvironment { } } +/// Send SIGTERM to the process group, wait 2s for graceful shutdown, then SIGKILL. +async fn sigterm_then_kill(child: &mut tokio::process::Child) { + #[cfg(unix)] + if let Some(pid) = child.id() { + #[allow(clippy::cast_possible_wrap)] + unsafe { + libc::kill(-(pid as i32), libc::SIGTERM); + } + if tokio::time::timeout(std::time::Duration::from_secs(2), child.wait()) + .await + .is_err() + { + let _ = child.kill().await; + let _ = child.wait().await; + } + } else { + let _ = child.kill().await; + let _ = child.wait().await; + } + #[cfg(not(unix))] + { + let _ = child.kill().await; + let _ = child.wait().await; + } +} + #[cfg(test)] mod tests { use super::*; @@ -473,7 +480,7 @@ mod tests { let dir = temp_dir(); let env = LocalExecutionEnvironment::new(dir.clone()); let result = env - .exec_command("echo hello", 5000, None, None) + .exec_command("echo hello", 5000, None, None, None) .await .unwrap(); @@ -489,7 +496,7 @@ mod tests { let dir = temp_dir(); let env = LocalExecutionEnvironment::new(dir.clone()); let result = env - .exec_command("exit 42", 5000, None, None) + .exec_command("exit 42", 5000, None, None, None) .await .unwrap(); @@ -503,7 +510,7 @@ mod tests { let dir = temp_dir(); let env = LocalExecutionEnvironment::new(dir.clone()); let result = env - .exec_command("sleep 10", 200, None, None) + .exec_command("sleep 10", 200, None, None, None) .await .unwrap(); @@ -517,7 +524,7 @@ mod tests { let dir = temp_dir(); let env = LocalExecutionEnvironment::new(dir.clone()); let result = env - .exec_command("echo err >&2", 5000, None, None) + .exec_command("echo err >&2", 5000, None, None, None) .await .unwrap(); diff --git a/crates/agent/src/profiles/openai.rs b/crates/agent/src/profiles/openai.rs index c5c53efc2..c87b76868 100644 --- a/crates/agent/src/profiles/openai.rs +++ b/crates/agent/src/profiles/openai.rs @@ -402,7 +402,7 @@ fn make_apply_patch_tool() -> RegisteredTool { "required": ["patch"] }), }, - executor: Arc::new(|args, env| { + executor: Arc::new(|args, env, _cancel| { Box::pin(async move { let patch_text = args .get("patch") diff --git a/crates/agent/src/session.rs b/crates/agent/src/session.rs index 167d3aa6f..eb0643afe 100644 --- a/crates/agent/src/session.rs +++ b/crates/agent/src/session.rs @@ -11,7 +11,6 @@ use crate::tool_registry::ToolRegistry; use crate::truncation::truncate_tool_output; use crate::types::{EventData, EventKind, SessionState, Turn}; use std::collections::VecDeque; -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::time::SystemTime; use futures::StreamExt; @@ -19,6 +18,7 @@ use llm::client::Client; use llm::error::{ProviderErrorKind, SdkError}; use llm::generate::StreamAccumulator; use llm::types::{Message, Request, StreamEvent, ToolChoice, ToolResult}; +use tokio_util::sync::CancellationToken; pub struct Session { id: String, @@ -31,7 +31,7 @@ pub struct Session { execution_env: Arc, steering_queue: Arc>>, followup_queue: Arc>>, - abort_flag: Arc, + cancel_token: CancellationToken, project_docs: Vec, env_context: EnvContext, } @@ -55,7 +55,7 @@ impl Session { execution_env, steering_queue: Arc::new(Mutex::new(VecDeque::new())), followup_queue: Arc::new(Mutex::new(VecDeque::new())), - abort_flag: Arc::new(AtomicBool::new(false)), + cancel_token: CancellationToken::new(), project_docs: Vec::new(), env_context: EnvContext::default(), } @@ -91,7 +91,7 @@ impl Session { // Detect git info via execution environment let git_branch = self .execution_env - .exec_command("git rev-parse --abbrev-ref HEAD", 5000, None, None) + .exec_command("git rev-parse --abbrev-ref HEAD", 5000, None, None, None) .await .ok() .filter(|r| r.exit_code == 0) @@ -101,7 +101,7 @@ impl Session { let git_status_short = if is_git_repo { self.execution_env - .exec_command("git status --short", 5000, None, None) + .exec_command("git status --short", 5000, None, None, None) .await .ok() .filter(|r| r.exit_code == 0) @@ -113,7 +113,7 @@ impl Session { let git_recent_commits = if is_git_repo { self.execution_env - .exec_command("git log --oneline -10", 5000, None, None) + .exec_command("git log --oneline -10", 5000, None, None, None) .await .ok() .filter(|r| r.exit_code == 0) @@ -157,15 +157,15 @@ impl Session { } pub fn abort(&self) { - self.abort_flag.store(true, Ordering::SeqCst); + self.cancel_token.cancel(); } pub fn followup_queue_handle(&self) -> Arc>> { self.followup_queue.clone() } - pub fn abort_flag_handle(&self) -> Arc { - self.abort_flag.clone() + pub fn cancel_token(&self) -> CancellationToken { + self.cancel_token.clone() } pub fn close(&mut self) { @@ -253,8 +253,8 @@ impl Session { break; } - // Check abort flag - if self.abort_flag.load(Ordering::SeqCst) { + // Check cancellation + if self.cancel_token.is_cancelled() { self.close(); return Err(AgentError::Aborted); } @@ -315,15 +315,15 @@ impl Session { } } - // Check abort flag between chunks - if self.abort_flag.load(Ordering::SeqCst) { + // Check cancellation between chunks + if self.cancel_token.is_cancelled() { break; } } // If aborted during streaming, drop the stream to cancel the HTTP // connection, then close the session before returning. - if self.abort_flag.load(Ordering::SeqCst) { + if self.cancel_token.is_cancelled() { drop(event_stream); self.close(); return Err(AgentError::Aborted); @@ -385,6 +385,16 @@ impl Session { // Execute tool calls (parallel or sequential based on provider) let results = self.execute_tool_calls(&tool_calls).await; + // Check cancellation after tool execution + if self.cancel_token.is_cancelled() { + self.history.push(Turn::ToolResults { + results, + timestamp: SystemTime::now(), + }); + self.close(); + return Err(AgentError::Aborted); + } + // Record tool results turn self.history.push(Turn::ToolResults { results, @@ -479,6 +489,17 @@ impl Session { ) -> Vec { let mut results = Vec::new(); for tc in tool_calls { + if self.cancel_token.is_cancelled() { + results.push(ToolResult { + tool_call_id: tc.id.clone(), + content: serde_json::json!("Cancelled"), + is_error: true, + image_data: None, + image_media_type: None, + }); + continue; + } + self.event_emitter.emit( EventKind::ToolCallStart, self.id.clone(), @@ -496,6 +517,7 @@ impl Session { self.provider_profile.tool_registry(), self.execution_env.clone(), self.config.tool_approval.as_ref(), + self.cancel_token.child_token(), ) .await; @@ -533,6 +555,7 @@ impl Session { let profile = self.provider_profile.clone(); let session_id = self.id.clone(); let config = self.config.clone(); + let cancel_token = self.cancel_token.clone(); let futures: Vec<_> = tool_calls .iter() @@ -542,6 +565,7 @@ impl Session { let profile = profile.clone(); let session_id = session_id.clone(); let config = config.clone(); + let cancel_token = cancel_token.clone(); let tc = tc.clone(); async move { emitter.emit( @@ -561,6 +585,7 @@ impl Session { profile.tool_registry(), env, config.tool_approval.as_ref(), + cancel_token.child_token(), ) .await; @@ -654,6 +679,7 @@ async fn execute_one_tool( registry: &ToolRegistry, env: Arc, tool_approval: Option<&ToolApprovalFn>, + cancel_token: CancellationToken, ) -> ToolResult { if let Some(approval_fn) = tool_approval { if let Err(denial_message) = approval_fn(tool_name, arguments) { @@ -681,7 +707,7 @@ async fn execute_one_tool( }; } - match (registered_tool.executor)(arguments.clone(), env).await { + match (registered_tool.executor)(arguments.clone(), env, cancel_token).await { Ok(output) => ToolResult { tool_call_id: tool_call_id.to_string(), content: serde_json::json!(output), @@ -1090,20 +1116,20 @@ mod tests { #[tokio::test] async fn abort_transitions_to_closed() { - let abort_flag = Arc::new(AtomicBool::new(false)); - let abort_flag_for_tool = abort_flag.clone(); + let cancel_token = CancellationToken::new(); + let cancel_token_for_tool = cancel_token.clone(); - // Tool that sets the abort flag when executed + // Tool that cancels the token when executed let abort_tool = RegisteredTool { definition: ToolDefinition { name: "set_abort".into(), description: "Sets abort flag".into(), parameters: serde_json::json!({"type": "object"}), }, - executor: Arc::new(move |_args, _env| { - let flag = abort_flag_for_tool.clone(); + executor: Arc::new(move |_args, _env, _cancel| { + let token = cancel_token_for_tool.clone(); Box::pin(async move { - flag.store(true, Ordering::SeqCst); + token.cancel(); Ok("done".to_string()) }) }), @@ -1127,8 +1153,8 @@ mod tests { }; let mut session = Session::new(client, profile, env, config); - // Wire the session's abort_flag to our shared one - session.abort_flag = abort_flag; + // Wire the session's cancel_token to our shared one + session.cancel_token = cancel_token; let result = session.process_input("Do something").await; @@ -1137,7 +1163,7 @@ mod tests { assert_eq!(session.state(), SessionState::Closed); // Should have processed: User + Assistant(tool_call) + ToolResults = 3 turns - // The tool set the abort flag, so the loop breaks before the next LLM call + // The tool cancelled the token, so the loop breaks before the next LLM call let turns = session.history().turns(); assert_eq!(turns.len(), 3); assert!(matches!(&turns[0], Turn::User { .. })); @@ -1371,7 +1397,7 @@ mod tests { "required": ["text"] }), }, - executor: Arc::new(|_args, _env| { + executor: Arc::new(|_args, _env, _cancel| { Box::pin(async move { Ok("should not reach".to_string()) }) }), }); @@ -1412,7 +1438,7 @@ mod tests { "required": ["text"] }), }, - executor: Arc::new(|_args, _env| { + executor: Arc::new(|_args, _env, _cancel| { Box::pin(async move { Ok("tool executed".to_string()) }) }), }); diff --git a/crates/agent/src/subagent.rs b/crates/agent/src/subagent.rs index 934e1d06c..479d41118 100644 --- a/crates/agent/src/subagent.rs +++ b/crates/agent/src/subagent.rs @@ -4,9 +4,9 @@ use crate::tool_registry::RegisteredTool; use crate::tools::required_str; use crate::types::Turn; use std::collections::{HashMap, VecDeque}; -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use llm::types::ToolDefinition; +use tokio_util::sync::CancellationToken; pub type SessionFactory = Arc Session + Send + Sync>; @@ -24,7 +24,7 @@ pub struct SubAgent { depth: usize, task: Option>>, followup_queue: Arc>>, - abort_flag: Arc, + cancel_token: CancellationToken, } #[cfg(test)] @@ -62,7 +62,7 @@ impl SubAgentManager { let agent_id = uuid::Uuid::new_v4().to_string(); let followup_queue = session.followup_queue_handle(); - let abort_flag = session.abort_flag_handle(); + let cancel_token = session.cancel_token(); let task = tokio::spawn(async move { session.process_input(&task_prompt).await?; @@ -85,7 +85,7 @@ impl SubAgentManager { depth, task: Some(task), followup_queue, - abort_flag, + cancel_token, }, ); @@ -138,7 +138,7 @@ impl SubAgentManager { AgentError::InvalidState(format!("No agent found with id: {agent_id}")) })?; - agent.abort_flag.store(true, Ordering::SeqCst); + agent.cancel_token.cancel(); if let Some(join_handle) = agent.task { join_handle.abort(); @@ -185,7 +185,7 @@ pub fn make_spawn_agent_tool( "required": ["task"] }), }, - executor: Arc::new(move |args, _env| { + executor: Arc::new(move |args, _env, _cancel| { let manager = manager.clone(); let session_factory = session_factory.clone(); Box::pin(async move { @@ -232,7 +232,7 @@ pub fn make_send_input_tool( "required": ["agent_id", "message"] }), }, - executor: Arc::new(move |args, _env| { + executor: Arc::new(move |args, _env, _cancel| { let manager = manager.clone(); Box::pin(async move { let agent_id = required_str(&args, "agent_id")?; @@ -265,7 +265,7 @@ pub fn make_wait_tool( "required": ["agent_id"] }), }, - executor: Arc::new(move |args, _env| { + executor: Arc::new(move |args, _env, _cancel| { let manager = manager.clone(); Box::pin(async move { let agent_id = required_str(&args, "agent_id")?; @@ -300,7 +300,7 @@ pub fn make_close_agent_tool( "required": ["agent_id"] }), }, - executor: Arc::new(move |args, _env| { + executor: Arc::new(move |args, _env, _cancel| { let manager = manager.clone(); Box::pin(async move { let agent_id = required_str(&args, "agent_id")?; diff --git a/crates/agent/src/test_support.rs b/crates/agent/src/test_support.rs index c9a403e22..64406ee07 100644 --- a/crates/agent/src/test_support.rs +++ b/crates/agent/src/test_support.rs @@ -12,6 +12,7 @@ use llm::client::Client; use llm::error::SdkError; use llm::provider::{ProviderAdapter, StreamEventStream}; use llm::types::{ContentPart, FinishReason, Message, Request, Response, StreamEvent, Usage}; +use tokio_util::sync::CancellationToken; // --- MockExecutionEnvironment --- @@ -120,6 +121,7 @@ impl ExecutionEnvironment for MockExecutionEnvironment { timeout_ms: u64, _working_dir: Option<&str>, _env_vars: Option<&std::collections::HashMap>, + _cancel_token: Option, ) -> Result { *self .captured_timeout @@ -232,6 +234,7 @@ impl ExecutionEnvironment for MutableMockExecutionEnvironment { _timeout_ms: u64, _working_dir: Option<&str>, _env_vars: Option<&std::collections::HashMap>, + _cancel_token: Option, ) -> Result { Ok(ExecResult { stdout: String::new(), @@ -542,7 +545,7 @@ pub(crate) fn make_echo_tool() -> crate::tool_registry::RegisteredTool { description: "Echoes the input".into(), parameters: serde_json::json!({"type": "object", "properties": {"text": {"type": "string"}}}), }, - executor: Arc::new(|args, _env| { + executor: Arc::new(|args, _env, _cancel| { Box::pin(async move { let text = args .get("text") @@ -562,7 +565,7 @@ pub(crate) fn make_error_tool() -> crate::tool_registry::RegisteredTool { description: "Always fails".into(), parameters: serde_json::json!({"type": "object"}), }, - executor: Arc::new(|_args, _env| { + executor: Arc::new(|_args, _env, _cancel| { Box::pin(async move { Err("tool execution failed".to_string()) }) }), } diff --git a/crates/agent/src/tool_registry.rs b/crates/agent/src/tool_registry.rs index b944e83dd..848a292d5 100644 --- a/crates/agent/src/tool_registry.rs +++ b/crates/agent/src/tool_registry.rs @@ -4,11 +4,13 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; use llm::types::ToolDefinition; +use tokio_util::sync::CancellationToken; pub type ToolExecutor = Arc< dyn Fn( serde_json::Value, Arc, + CancellationToken, ) -> Pin> + Send>> + Send + Sync, @@ -72,7 +74,7 @@ mod tests { description: format!("Tool {name}"), parameters: serde_json::json!({"type": "object"}), }, - executor: Arc::new(|_args, _env| Box::pin(async { Ok("ok".into()) })), + executor: Arc::new(|_args, _env, _cancel| Box::pin(async { Ok("ok".into()) })), } } @@ -116,7 +118,7 @@ mod tests { description: "version 1".into(), parameters: serde_json::json!({}), }, - executor: Arc::new(|_args, _env| Box::pin(async { Ok("v1".into()) })), + executor: Arc::new(|_args, _env, _cancel| Box::pin(async { Ok("v1".into()) })), }); registry.register(RegisteredTool { definition: ToolDefinition { @@ -124,7 +126,7 @@ mod tests { description: "version 2".into(), parameters: serde_json::json!({}), }, - executor: Arc::new(|_args, _env| Box::pin(async { Ok("v2".into()) })), + executor: Arc::new(|_args, _env, _cancel| Box::pin(async { Ok("v2".into()) })), }); let tool = registry.get("tool_a").unwrap(); @@ -167,7 +169,7 @@ mod tests { use crate::test_support::MockExecutionEnvironment; let env: Arc = Arc::new(MockExecutionEnvironment::default()); - let result = (tool.executor)(serde_json::json!({}), env).await; + let result = (tool.executor)(serde_json::json!({}), env, CancellationToken::new()).await; assert_eq!(result.unwrap(), "ok"); } diff --git a/crates/agent/src/tools.rs b/crates/agent/src/tools.rs index 276fecbd1..45a3a8ee2 100644 --- a/crates/agent/src/tools.rs +++ b/crates/agent/src/tools.rs @@ -27,7 +27,7 @@ pub fn make_read_file_tool() -> RegisteredTool { "required": ["file_path"] }), }, - executor: Arc::new(|args, env| { + executor: Arc::new(|args, env, _cancel| { Box::pin(async move { let file_path = required_str(&args, "file_path")?; let offset = args.get("offset").and_then(serde_json::Value::as_u64); @@ -60,7 +60,7 @@ pub fn make_write_file_tool() -> RegisteredTool { "required": ["file_path", "content"] }), }, - executor: Arc::new(|args, env| { + executor: Arc::new(|args, env, _cancel| { Box::pin(async move { let file_path = required_str(&args, "file_path")?; let content = required_str(&args, "content")?; @@ -89,7 +89,7 @@ pub fn make_edit_file_tool() -> RegisteredTool { "required": ["file_path", "old_string", "new_string"] }), }, - executor: Arc::new(|args, env| { + executor: Arc::new(|args, env, _cancel| { Box::pin(async move { let file_path = required_str(&args, "file_path")?; let old_string = required_str(&args, "old_string")?; @@ -157,7 +157,7 @@ pub fn make_shell_tool_with_config(config: &SessionConfig) -> RegisteredTool { "required": ["command"] }), }, - executor: Arc::new(move |args, env| { + executor: Arc::new(move |args, env, cancel| { Box::pin(async move { let command = required_str(&args, "command")?; let timeout_ms = args @@ -167,7 +167,7 @@ pub fn make_shell_tool_with_config(config: &SessionConfig) -> RegisteredTool { .min(max_timeout); let result = env - .exec_command(command, timeout_ms, None, None) + .exec_command(command, timeout_ms, None, None, Some(cancel)) .await?; let mut output = String::new(); @@ -203,7 +203,7 @@ pub fn make_grep_tool() -> RegisteredTool { "required": ["pattern"] }), }, - executor: Arc::new(|args, env| { + executor: Arc::new(|args, env, _cancel| { Box::pin(async move { let pattern = required_str(&args, "pattern")?; let path = args @@ -249,7 +249,7 @@ pub fn make_glob_tool() -> RegisteredTool { "required": ["pattern"] }), }, - executor: Arc::new(|args, env| { + executor: Arc::new(|args, env, _cancel| { Box::pin(async move { let pattern = required_str(&args, "pattern")?; let path = args @@ -281,7 +281,7 @@ pub(crate) fn make_read_many_files_tool() -> RegisteredTool { "required": ["paths"] }), }, - executor: Arc::new(|args, env| { + executor: Arc::new(|args, env, _cancel| { Box::pin(async move { let paths = args["paths"] .as_array() @@ -322,7 +322,7 @@ pub(crate) fn make_list_dir_tool() -> RegisteredTool { "required": ["path"] }), }, - executor: Arc::new(|args, env| { + executor: Arc::new(|args, env, _cancel| { Box::pin(async move { let path = required_str(&args, "path")?; #[allow(clippy::cast_possible_truncation)] @@ -363,7 +363,7 @@ pub(crate) fn make_web_search_tool() -> RegisteredTool { "required": ["query"] }), }, - executor: Arc::new(|_args, _env| { + executor: Arc::new(|_args, _env, _cancel| { Box::pin(async move { Ok("Web search is not configured. This is a placeholder tool.".to_string()) }) @@ -385,7 +385,7 @@ pub(crate) fn make_web_fetch_tool() -> RegisteredTool { "required": ["url"] }), }, - executor: Arc::new(|_args, _env| { + executor: Arc::new(|_args, _env, _cancel| { Box::pin(async move { Ok("Web fetch is not configured. This is a placeholder tool.".to_string()) }) @@ -399,6 +399,7 @@ mod tests { use crate::execution_env::*; use crate::test_support::MockExecutionEnvironment; use std::collections::HashMap; + use tokio_util::sync::CancellationToken; #[tokio::test] async fn read_file_returns_content() { @@ -410,7 +411,7 @@ mod tests { apply_read_offset_limit: true, ..Default::default() }); - let result = (tool.executor)(serde_json::json!({"file_path": "/test.txt"}), env).await; + let result = (tool.executor)(serde_json::json!({"file_path": "/test.txt"}), env, CancellationToken::new()).await; assert_eq!(result.unwrap(), " 1 | hello\n 2 | world"); } @@ -430,6 +431,7 @@ mod tests { let result = (tool.executor)( serde_json::json!({"file_path": "/test.txt", "offset": 2, "limit": 2}), env, + CancellationToken::new(), ) .await; assert_eq!(result.unwrap(), " 2 | line2\n 3 | line3"); @@ -443,6 +445,7 @@ mod tests { let result = (tool.executor)( serde_json::json!({"file_path": "/out.txt", "content": "hello"}), env_clone, + CancellationToken::new(), ) .await; assert_eq!(result.unwrap(), "Successfully wrote to /out.txt"); @@ -469,6 +472,7 @@ mod tests { "new_string": "goodbye" }), env_clone, + CancellationToken::new(), ) .await; assert_eq!(result.unwrap(), "Successfully edited /f.txt"); @@ -493,6 +497,7 @@ mod tests { "new_string": "replacement" }), env, + CancellationToken::new(), ) .await; assert_eq!(result.unwrap_err(), "old_string not found in file"); @@ -514,6 +519,7 @@ mod tests { "new_string": "cc" }), env, + CancellationToken::new(), ) .await; let err = result.unwrap_err(); @@ -539,6 +545,7 @@ mod tests { "replace_all": true }), env_clone, + CancellationToken::new(), ) .await; assert_eq!(result.unwrap(), "Successfully edited /f.txt"); @@ -560,7 +567,7 @@ mod tests { }, ..Default::default() }); - let result = (tool.executor)(serde_json::json!({"command": "echo hello"}), env).await; + let result = (tool.executor)(serde_json::json!({"command": "echo hello"}), env, CancellationToken::new()).await; let output = result.unwrap(); assert!(output.contains("Exit code: 0")); assert!(output.contains("hello")); @@ -574,6 +581,7 @@ mod tests { let _result = (tool.executor)( serde_json::json!({"command": "sleep 1", "timeout_ms": 5000}), env_clone, + CancellationToken::new(), ) .await; assert_eq!(*env.captured_timeout.lock().unwrap(), Some(5000)); @@ -592,7 +600,7 @@ mod tests { }, ..Default::default() }); - let result = (tool.executor)(serde_json::json!({"command": "false"}), env).await; + let result = (tool.executor)(serde_json::json!({"command": "false"}), env, CancellationToken::new()).await; let output = result.unwrap(); assert!(output.contains("Exit code: 1")); assert!(output.contains("error")); @@ -611,7 +619,7 @@ mod tests { }, ..Default::default() }); - let result = (tool.executor)(serde_json::json!({"command": "sleep 100"}), env).await; + let result = (tool.executor)(serde_json::json!({"command": "sleep 100"}), env, CancellationToken::new()).await; let output = result.unwrap(); assert!(output.starts_with("Command timed out.\n")); } @@ -623,7 +631,7 @@ mod tests { grep_results: vec!["src/main.rs:10:fn main()".into(), "src/lib.rs:5:pub fn".into()], ..Default::default() }); - let result = (tool.executor)(serde_json::json!({"pattern": "fn"}), env).await; + let result = (tool.executor)(serde_json::json!({"pattern": "fn"}), env, CancellationToken::new()).await; let output = result.unwrap(); assert!(output.contains("src/main.rs:10:fn main()")); assert!(output.contains("src/lib.rs:5:pub fn")); @@ -636,7 +644,7 @@ mod tests { glob_results: vec!["src/main.rs".into(), "src/lib.rs".into()], ..Default::default() }); - let result = (tool.executor)(serde_json::json!({"pattern": "src/**/*.rs"}), env).await; + let result = (tool.executor)(serde_json::json!({"pattern": "src/**/*.rs"}), env, CancellationToken::new()).await; let output = result.unwrap(); assert!(output.contains("src/main.rs")); assert!(output.contains("src/lib.rs"));