use async_trait::async_trait;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::any::Any;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use crate::event::Source;
use crate::llm::ToolDefinition;
use crate::{AgentLlmConfig, LlmProvider};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ForbiddenCmd {
pub command: String,
pub reason: String,
#[serde(default)]
pub reconsider: bool,
}
impl ForbiddenCmd {
pub fn validate(&self) -> std::result::Result<(), Vec<String>> {
let mut errors = Vec::new();
if self.command.is_empty() {
errors.push("forbidden command name cannot be empty".to_string());
}
if regex::Regex::new(&self.command).is_err() {
errors.push(format!("invalid regex for name: '{}'", self.command));
}
if errors.is_empty() {
Ok(())
} else {
Err(errors)
}
}
}
pub trait CommandAllowList: Send + Sync {
fn is_allowed(&self, command: &str) -> bool;
fn add_session_pattern(&self, pattern: &str);
}
pub trait InterruptPayload: Send + Sync + std::fmt::Debug {
fn as_any(&self) -> &dyn Any;
fn display_output(&self) -> String;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum DiffLine {
Added { line: u64, content: String },
Removed { line: u64, content: String },
Context { line: u64, content: String },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileDiff {
pub path: String,
pub diff: Vec<DiffLine>,
}
#[derive(Clone)]
pub enum ToolResult {
Success {
tool_call_id: String,
output: String,
sources: Vec<Source>,
diffs: Vec<FileDiff>,
working_dir: Option<PathBuf>,
},
Failure {
tool_call_id: String,
output: String,
},
Interrupt {
tool_call_id: String,
payload: Arc<dyn InterruptPayload>,
},
}
impl std::fmt::Debug for ToolResult {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Success {
tool_call_id,
output,
..
} => f
.debug_struct("Success")
.field("tool_call_id", tool_call_id)
.field("output", output)
.finish(),
Self::Failure {
tool_call_id,
output,
} => f
.debug_struct("Failure")
.field("tool_call_id", tool_call_id)
.field("output", output)
.finish(),
Self::Interrupt {
tool_call_id,
payload,
} => f
.debug_struct("Interrupt")
.field("tool_call_id", tool_call_id)
.field("output", &payload.display_output())
.finish(),
}
}
}
impl ToolResult {
pub fn success(tool_call_id: &str, output: impl Into<String>) -> Self {
Self::Success {
tool_call_id: tool_call_id.into(),
output: output.into(),
sources: Vec::new(),
diffs: Vec::new(),
working_dir: None,
}
}
pub fn failure(tool_call_id: &str, output: impl Into<String>) -> Self {
Self::Failure {
tool_call_id: tool_call_id.into(),
output: output.into(),
}
}
pub fn success_with_diffs(
tool_call_id: &str,
output: impl Into<String>,
diffs: Vec<FileDiff>,
) -> Self {
Self::Success {
tool_call_id: tool_call_id.into(),
output: output.into(),
sources: Vec::new(),
diffs,
working_dir: None,
}
}
pub fn interrupt(tool_call_id: &str, payload: Arc<dyn InterruptPayload>) -> Self {
Self::Interrupt {
tool_call_id: tool_call_id.into(),
payload,
}
}
pub fn with_sources(mut self, sources: Vec<Source>) -> Self {
if let Self::Success {
sources: ref mut s, ..
} = self
{
*s = sources;
}
self
}
pub fn with_working_dir(mut self, path: PathBuf) -> Self {
if let Self::Success { working_dir, .. } = &mut self {
*working_dir = Some(path);
}
self
}
pub fn working_dir(&self) -> Option<&Path> {
match self {
Self::Success { working_dir, .. } => working_dir.as_deref(),
_ => None,
}
}
pub fn tool_call_id(&self) -> &str {
match self {
Self::Success { tool_call_id, .. }
| Self::Failure { tool_call_id, .. }
| Self::Interrupt { tool_call_id, .. } => tool_call_id,
}
}
pub fn output(&self) -> String {
match self {
Self::Success { output, .. } | Self::Failure { output, .. } => output.clone(),
Self::Interrupt { payload, .. } => payload.display_output(),
}
}
pub fn payload(&self) -> Option<&Arc<dyn InterruptPayload>> {
match self {
Self::Interrupt { payload, .. } => Some(payload),
_ => None,
}
}
pub fn is_success(&self) -> bool {
matches!(self, Self::Success { .. })
}
pub fn is_interrupt(&self) -> bool {
matches!(self, Self::Interrupt { .. })
}
pub fn sources(&self) -> &[Source] {
match self {
Self::Success { sources, .. } => sources,
_ => &[],
}
}
pub fn diffs(&self) -> &[FileDiff] {
match self {
Self::Success { diffs, .. } => diffs,
_ => &[],
}
}
pub fn truncate_output(&mut self, max_bytes: usize, max_lines: usize) {
let output = match self {
Self::Success { output, .. } | Self::Failure { output, .. } => output,
Self::Interrupt { .. } => return,
};
if output.len() <= max_bytes && max_lines == usize::MAX {
return;
}
let byte_cut = output.floor_char_boundary(max_bytes);
let line_cut = output
.match_indices('\n')
.nth(max_lines.saturating_sub(1))
.map(|(i, _)| i)
.unwrap_or(output.len());
let cut = byte_cut.min(line_cut);
if cut < output.len() {
output.truncate(cut);
output.push_str("\n[truncated]");
}
}
}
pub fn parse_args<T: DeserializeOwned>(tool_name: &str, args: Value) -> anyhow::Result<T> {
serde_json::from_value(args)
.map_err(|e| anyhow::anyhow!("Tool '{}': Invalid arguments: {}", tool_name, e))
}
#[async_trait]
pub trait Tool: Send + Sync {
fn name(&self) -> &str;
fn description(&self) -> &str;
fn parameters(&self) -> Value;
async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult>;
fn timeout_secs(&self) -> Option<u64> {
None
}
fn max_output_bytes(&self) -> usize {
256 * 1024
}
fn max_output_lines(&self) -> usize {
10_000
}
fn concurrent(&self) -> bool {
true
}
fn supports_background(&self) -> bool {
false
}
fn humanize(&self, args: &Value) -> String {
let Some(obj) = args.as_object() else {
return self.name().to_string();
};
let keys: Vec<&str> = obj.keys().map(|k| k.as_str()).collect();
if keys.is_empty() {
return self.name().to_string();
}
let summary = keys.join(", ");
if summary.len() > 80 {
let truncated: String = summary.chars().take(77).collect();
format!("{truncated}…")
} else {
summary
}
}
}
#[derive(Clone, Debug, Default)]
pub struct BackgroundRequest {
requested: CancellationToken,
}
impl BackgroundRequest {
pub fn request(&self) {
self.requested.cancel();
}
pub async fn requested(&self) {
self.requested.cancelled().await;
}
}
pub struct ToolContext<'a> {
pub tool_call_id: &'a str,
pub args: Value,
pub llm: Option<&'a AgentLlmConfig>,
pub provider: Option<Arc<dyn LlmProvider>>,
pub working_dir: Option<&'a PathBuf>,
cancel_token: &'a CancellationToken,
progress_tx: Option<tokio::sync::mpsc::UnboundedSender<String>>,
background_request: Option<&'a BackgroundRequest>,
}
impl<'a> ToolContext<'a> {
pub fn new(
tool_call_id: &'a str,
args: Value,
working_dir: Option<&'a PathBuf>,
cancel_token: &'a CancellationToken,
) -> Self {
Self {
tool_call_id,
args,
llm: None,
provider: None,
working_dir,
cancel_token,
progress_tx: None,
background_request: None,
}
}
pub fn with_llm(mut self, llm: &'a AgentLlmConfig) -> Self {
self.llm = Some(llm);
self
}
pub fn with_provider(mut self, provider: Arc<dyn LlmProvider>) -> Self {
self.provider = Some(provider);
self
}
pub fn with_progress(mut self, tx: tokio::sync::mpsc::UnboundedSender<String>) -> Self {
self.progress_tx = Some(tx);
self
}
pub fn with_background_request(mut self, request: &'a BackgroundRequest) -> Self {
self.background_request = Some(request);
self
}
pub fn background_request(&self) -> Option<&BackgroundRequest> {
self.background_request
}
pub fn progress_sender(&self) -> Option<tokio::sync::mpsc::UnboundedSender<String>> {
self.progress_tx.clone()
}
pub fn progress(&self, line: &str) {
if let Some(ref tx) = self.progress_tx {
let _ = tx.send(line.to_string());
}
}
pub fn parse_args<T: DeserializeOwned>(&self, tool_name: &str) -> anyhow::Result<T> {
T::deserialize(&self.args)
.map_err(|e| anyhow::anyhow!("Tool '{}': Invalid arguments: {}", tool_name, e))
}
pub fn cancel_token(&self) -> &CancellationToken {
self.cancel_token
}
pub fn child_token(&self) -> CancellationToken {
self.cancel_token.child_token()
}
}
#[derive(Clone)]
pub struct ToolRegistry {
tools: HashMap<String, Arc<dyn Tool>>,
aliases: HashMap<String, String>,
truncate_on_overflow: bool,
}
impl ToolRegistry {
pub fn new() -> Self {
Self {
tools: HashMap::new(),
aliases: HashMap::new(),
truncate_on_overflow: false,
}
}
pub fn set_truncate_on_overflow(&mut self, enabled: bool) {
self.truncate_on_overflow = enabled;
}
pub fn register(&mut self, tool: Arc<dyn Tool>) {
tracing::debug!(tool = tool.name(), "registered tool");
self.tools.insert(tool.name().to_string(), tool);
}
pub fn remove(&mut self, name: &str) -> bool {
let removed = self.tools.remove(name).is_some();
if removed {
let tools = &self.tools;
self.aliases.retain(|_, target| tools.contains_key(target));
}
removed
}
pub fn strip_prefixes(&mut self, prefixes: &[&str]) {
let to_remove: Vec<String> = self
.tools
.keys()
.filter(|name| prefixes.iter().any(|p| name.starts_with(p)))
.cloned()
.collect();
for name in to_remove {
self.tools.remove(&name);
}
let tools = &self.tools;
self.aliases.retain(|_, target| tools.contains_key(target));
}
pub fn alias(&mut self, alias: &str, target: &str) {
tracing::debug!(alias, target, "registered tool alias");
self.aliases.insert(alias.to_string(), target.to_string());
}
pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
let result = self.tools.get(name).or_else(|| {
if let Some(target) = self.aliases.get(name) {
tracing::debug!(
alias = name,
target = target.as_str(),
"resolved tool alias"
);
self.tools.get(target)
} else {
None
}
});
if result.is_none() {
tracing::warn!(tool = name, "tool not found in registry");
}
result
}
pub fn contains(&self, name: &str) -> bool {
self.resolve(name).is_some()
}
fn resolve(&self, name: &str) -> Option<&Arc<dyn Tool>> {
self.tools.get(name).or_else(|| {
self.aliases
.get(name)
.and_then(|target| self.tools.get(target))
})
}
pub fn retain(&mut self, mut f: impl FnMut(&str) -> bool) {
self.tools.retain(|name, _| f(name));
let tools = &self.tools;
self.aliases.retain(|_, target| tools.contains_key(target));
}
pub fn list(&self) -> Vec<&str> {
let mut names: Vec<&str> = self.tools.keys().map(|s| s.as_str()).collect();
names.sort();
names
}
pub fn list_with_descriptions(&self) -> Vec<(&str, &str)> {
let mut items: Vec<(&str, &str)> = self
.tools
.values()
.map(|t| (t.name(), t.description()))
.collect();
items.sort_by_key(|(name, _)| *name);
items
}
pub async fn execute(
&self,
call: &crate::ToolCall,
working_dir: Option<&PathBuf>,
cancel_token: &CancellationToken,
) -> ToolResult {
self.execute_inner(call, working_dir, cancel_token, None)
.await
}
pub async fn execute_with_progress(
&self,
call: &crate::ToolCall,
working_dir: Option<&PathBuf>,
cancel_token: &CancellationToken,
progress_tx: tokio::sync::mpsc::UnboundedSender<String>,
) -> ToolResult {
self.execute_inner(call, working_dir, cancel_token, Some(progress_tx))
.await
}
async fn execute_inner(
&self,
call: &crate::ToolCall,
working_dir: Option<&PathBuf>,
cancel_token: &CancellationToken,
progress_tx: Option<tokio::sync::mpsc::UnboundedSender<String>>,
) -> ToolResult {
let mut ctx = ToolContext::new(&call.id, call.arguments.clone(), working_dir, cancel_token);
if let Some(tx) = progress_tx {
ctx = ctx.with_progress(tx);
}
self.execute_in_context(call, ctx).await
}
pub async fn execute_in_context(
&self,
call: &crate::ToolCall,
ctx: ToolContext<'_>,
) -> ToolResult {
let Some(tool) = self.get(&call.name) else {
return ToolResult::failure(&call.id, format!("Unknown tool: {}", call.name));
};
let cancel_token = ctx.cancel_token;
let child_token = cancel_token.child_token();
let ctx = ToolContext {
cancel_token: &child_token,
..ctx
};
let fut = tool.execute(ctx);
let timeout_secs = tool.timeout_secs();
let run = async {
match timeout_secs {
Some(secs) => tokio::time::timeout(std::time::Duration::from_secs(secs), fut)
.await
.map_err(|_| secs),
None => Ok(fut.await),
}
};
let outcome = tokio::select! {
biased;
_ = cancel_token.cancelled() => {
child_token.cancel();
return ToolResult::failure(
&call.id,
format!("Tool '{}' cancelled", call.name),
);
}
r = run => match r {
Ok(result) => result,
Err(secs) => {
child_token.cancel();
return ToolResult::failure(
&call.id,
format!("Tool '{}' timed out after {secs}s", call.name),
);
}
},
};
let result = match outcome {
Ok(result) => result,
Err(e) => ToolResult::failure(&call.id, format!("Tool error: {e:#}")),
};
enforce_output_limits(
&call.id,
result,
tool.max_output_bytes(),
tool.max_output_lines(),
self.truncate_on_overflow,
)
}
pub fn is_concurrent(&self, call: &crate::ToolCall) -> bool {
self.resolve(&call.name)
.is_none_or(|tool| tool.concurrent())
}
pub fn supports_background(&self, call: &crate::ToolCall) -> bool {
self.resolve(&call.name)
.is_some_and(|tool| tool.supports_background())
}
pub fn humanize(&self, call: &crate::ToolCall) -> String {
self.get(&call.name)
.map(|t| t.humanize(&call.arguments))
.unwrap_or_else(|| call.name.clone())
}
pub fn definitions(&self) -> Vec<ToolDefinition> {
let mut defs: Vec<ToolDefinition> = self
.tools
.values()
.map(|t| ToolDefinition {
name: t.name().to_string(),
description: t.description().to_string(),
parameters: t.parameters(),
})
.collect();
defs.sort_by(|a, b| a.name.cmp(&b.name));
tracing::debug!(count = defs.len(), "generated tool definitions");
defs
}
}
impl Default for ToolRegistry {
fn default() -> Self {
Self::new()
}
}
fn enforce_output_limits(
tool_call_id: &str,
mut result: ToolResult,
max_bytes: usize,
max_lines: usize,
truncate: bool,
) -> ToolResult {
let output = match &result {
ToolResult::Success { output, .. } => output,
ToolResult::Failure { .. } => {
result.truncate_output(max_bytes, max_lines);
return result;
}
ToolResult::Interrupt { .. } => return result,
};
let bytes = output.len();
let lines = output.lines().count();
let over_bytes = bytes > max_bytes;
let over_lines = lines > max_lines;
if !over_bytes && !over_lines {
return result;
}
if truncate {
result.truncate_output(max_bytes, max_lines);
return result;
}
if over_bytes {
ToolResult::failure(
tool_call_id,
format!(
"Output too large ({bytes} bytes, limit {max_bytes}). \
Use more specific arguments to reduce output size."
),
)
} else {
ToolResult::failure(
tool_call_id,
format!(
"Output too many lines ({lines}, limit {max_lines}). \
Use more specific arguments to reduce output size."
),
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn failed_output_is_truncated_not_passed_through() {
let output = "x\n".repeat(100);
for truncate in [false, true] {
let result = enforce_output_limits(
"call",
ToolResult::failure("call", output.clone()),
50,
usize::MAX,
truncate,
);
let ToolResult::Failure { output, .. } = result else {
panic!("a failure stays a failure");
};
assert!(output.len() < 70, "{output}");
assert!(output.ends_with("[truncated]"), "{output}");
}
}
struct DummyTool(&'static str);
#[async_trait]
impl Tool for DummyTool {
fn name(&self) -> &str {
self.0
}
fn description(&self) -> &str {
"dummy"
}
fn parameters(&self) -> Value {
serde_json::json!({})
}
async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
Ok(ToolResult::success(ctx.tool_call_id, "ok"))
}
}
#[test]
fn test_registry_empty() {
let registry = ToolRegistry::new();
assert_eq!(registry.list().len(), 0);
}
#[test]
fn test_registry_definitions_empty() {
let registry = ToolRegistry::new();
assert!(registry.definitions().is_empty());
}
#[test]
fn test_registry_lookup() {
let mut registry = ToolRegistry::new();
registry.register(Arc::new(DummyTool("exec")));
registry.register(Arc::new(DummyTool("slack_react")));
registry.register(Arc::new(DummyTool("browser")));
assert_eq!(registry.definitions().len(), 3);
assert!(registry.get("slack_react").is_some());
}
#[test]
fn test_registry_alias_contains_and_remove_cleanup() {
let mut registry = ToolRegistry::new();
registry.register(Arc::new(DummyTool("exec")));
registry.alias("bash_exec", "exec");
assert!(registry.contains("exec"));
assert!(registry.contains("bash_exec"));
assert!(registry.get("bash_exec").is_some());
assert!(registry.remove("exec"));
assert!(!registry.contains("exec"));
assert!(!registry.contains("bash_exec"));
assert!(registry.get("bash_exec").is_none());
}
struct SlowTimeoutTool;
#[async_trait]
impl Tool for SlowTimeoutTool {
fn name(&self) -> &str {
"slow"
}
fn description(&self) -> &str {
"sleeps forever"
}
fn parameters(&self) -> Value {
serde_json::json!({})
}
fn timeout_secs(&self) -> Option<u64> {
Some(600)
}
async fn execute(&self, _ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
Ok(ToolResult::success("unused", "should not return"))
}
}
#[tokio::test]
async fn registry_cancel_short_circuits_tool_with_timeout_secs() {
let mut registry = ToolRegistry::new();
registry.register(Arc::new(SlowTimeoutTool));
let cancel = CancellationToken::new();
let cancel_clone = cancel.clone();
let call = crate::ToolCall {
id: "c1".into(),
name: "slow".into(),
arguments: serde_json::json!({}),
};
let registry = Arc::new(registry);
let registry2 = registry.clone();
let exec = tokio::spawn(async move { registry2.execute(&call, None, &cancel).await });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
cancel_clone.cancel();
let result = tokio::time::timeout(std::time::Duration::from_secs(2), exec)
.await
.expect("registry did not return within 2s of cancellation")
.unwrap();
assert!(!result.is_success(), "expected failure on cancel");
assert!(
result.output().contains("cancelled"),
"expected 'cancelled' in output, got: {}",
result.output()
);
}
}