scriptoriumadmin/fabro
1
0
Fork 0
mirror of https://github.com/fabro-sh/fabro.git synced 2026-09-06 08:18:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Replace abort_flag with CancellationToken for abort-aware process cancellation

Browse source 
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<AtomicBool> 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 <noreply@anthropic.com>
This commit is contained in:
Bryan Helmkamp 2026-02-23 18:10:10 -05:00
parent 028353c82e
commit 2a2a610654
13 changed files with 176 additions and 114 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
Cargo.lock generated
View file

@ -23,6 +23,7 @@ dependencies = [
"terminal",
"thiserror 2.0.18",
"tokio",
"tokio-util",
"uuid",
]

1
crates/agent/Cargo.toml
View file

@ -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 }

4
crates/agent/README.md
View file

@ -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

5
crates/agent/src/cli.rs
View file

@ -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

35
crates/agent/src/docker_env.rs
View file

@ -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<String, String>>,
cancel_token: Option<CancellationToken>,
) -> Result<ExecResult, String> {
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<String, String>>,
cancel_token: Option<CancellationToken>,
) -> Result<ExecResult, String> {
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<String> = 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<String> = 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();

4
crates/agent/src/execution_env.rs
View file

@ -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<String, String>>,
cancel_token: Option<CancellationToken>,
) -> Result<ExecResult, String>;
async fn grep(
&self,
@ -92,7 +94,7 @@ mod tests {
#[tokio::test]
async fn mock_env_exec_command() {
let env: Arc<dyn ExecutionEnvironment> = 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);
}

83
crates/agent/src/local_env.rs
View file

@ -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<String, String>>,
cancel_token: Option<CancellationToken>,
) -> Result<ExecResult, String> {
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();

2
crates/agent/src/profiles/openai.rs
View file

@ -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")

78
crates/agent/src/session.rs
View file

@ -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<dyn ExecutionEnvironment>,
steering_queue: Arc<Mutex<VecDeque<String>>>,
followup_queue: Arc<Mutex<VecDeque<String>>>,
abort_flag: Arc<AtomicBool>,
cancel_token: CancellationToken,
project_docs: Vec<String>,
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<Mutex<VecDeque<String>>> {
self.followup_queue.clone()
}
pub fn abort_flag_handle(&self) -> Arc<AtomicBool> {
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<ToolResult> {
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<dyn ExecutionEnvironment>,
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()) })
}),
});

18
crates/agent/src/subagent.rs
View file

@ -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<dyn Fn() -> Session + Send + Sync>;
@ -24,7 +24,7 @@ pub struct SubAgent {
depth: usize,
task: Option<tokio::task::JoinHandle<Result<SubAgentResult, AgentError>>>,
followup_queue: Arc<Mutex<VecDeque<String>>>,
abort_flag: Arc<AtomicBool>,
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")?;

7
crates/agent/src/test_support.rs
View file

@ -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<String, String>>,
_cancel_token: Option<CancellationToken>,
) -> Result<ExecResult, String> {
*self
.captured_timeout
@ -232,6 +234,7 @@ impl ExecutionEnvironment for MutableMockExecutionEnvironment {
_timeout_ms: u64,
_working_dir: Option<&str>,
_env_vars: Option<&std::collections::HashMap<String, String>>,
_cancel_token: Option<CancellationToken>,
) -> Result<ExecResult, String> {
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()) })
}),
}

10
crates/agent/src/tool_registry.rs
View file

@ -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<dyn ExecutionEnvironment>,
CancellationToken,
) -> Pin<Box<dyn Future<Output = Result<String, String>> + 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<dyn ExecutionEnvironment> = 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");
}

42
crates/agent/src/tools.rs
View file

@ -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"));