use async_trait::async_trait;
use futures_util::future::BoxFuture;
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 std::time::Duration;
use tempfile::TempDir;
use tokio::fs::OpenOptions;
use tokio::io::AsyncWriteExt;
use tokio::sync::OnceCell;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
use crate::event::Source;
use crate::llm::ToolDefinition;
use crate::{AgentLlmConfig, ContentPart, 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);
}
#[async_trait]
pub trait ToolGate: Send + Sync {
async fn check(&self, call: &crate::ToolCall) -> GateDecision;
}
pub enum GateDecision {
Allow,
Deny(String),
Ask(GateAsk),
}
pub struct GateAsk {
pub answer: BoxFuture<'static, GateAnswer>,
pub timeout: Duration,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GateAnswer {
Allow,
Deny(String),
}
impl std::fmt::Debug for GateDecision {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Allow => f.write_str("Allow"),
Self::Deny(reason) => f.debug_tuple("Deny").field(reason).finish(),
Self::Ask(ask) => f
.debug_struct("Ask")
.field("timeout", &ask.timeout)
.finish(),
}
}
}
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,
content: Vec<ContentPart>,
sources: Vec<Source>,
diffs: Vec<FileDiff>,
working_dir: Option<PathBuf>,
},
Failure {
tool_call_id: String,
content: Vec<ContentPart>,
},
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, .. } => f
.debug_struct("Success")
.field("tool_call_id", tool_call_id)
.field("output", &self.output())
.finish(),
Self::Failure { tool_call_id, .. } => f
.debug_struct("Failure")
.field("tool_call_id", tool_call_id)
.field("output", &self.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_with_content(
tool_call_id,
vec![ContentPart::Text {
text: output.into(),
}],
)
}
pub fn success_with_content(tool_call_id: &str, content: Vec<ContentPart>) -> Self {
Self::Success {
tool_call_id: tool_call_id.into(),
content,
sources: Vec::new(),
diffs: Vec::new(),
working_dir: None,
}
}
pub fn with_parts(mut self, parts: Vec<ContentPart>) -> Self {
if let Self::Success { content, .. } = &mut self {
content.extend(parts);
}
self
}
pub fn content(&self) -> &[ContentPart] {
match self {
Self::Success { content, .. } | Self::Failure { content, .. } => content,
Self::Interrupt { .. } => &[],
}
}
pub fn append_text(&mut self, text: &str) {
match self {
Self::Success { content, .. } | Self::Failure { content, .. } => {
match content.last_mut() {
Some(ContentPart::Text { text: last }) => last.push_str(text),
_ => content.push(ContentPart::Text {
text: text.to_owned(),
}),
}
}
Self::Interrupt { .. } => {}
}
}
pub fn failure(tool_call_id: &str, output: impl Into<String>) -> Self {
Self::failure_with_content(
tool_call_id,
vec![ContentPart::Text {
text: output.into(),
}],
)
}
pub fn failure_with_content(tool_call_id: &str, content: Vec<ContentPart>) -> Self {
Self::Failure {
tool_call_id: tool_call_id.into(),
content,
}
}
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(),
content: vec![ContentPart::Text {
text: 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 { content, .. } | Self::Failure { content, .. } => content
.iter()
.filter_map(|part| match part {
ContentPart::Text { text } => Some(text.as_str()),
_ => None,
})
.collect(),
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 = self.output();
let content = match self {
Self::Success { content, .. } | Self::Failure { content, .. } => content,
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 = if max_lines == 0 {
0
} else {
output
.match_indices('\n')
.nth(max_lines - 1)
.map(|(index, _)| index)
.unwrap_or(output.len())
};
let cut = byte_cut.min(line_cut);
if cut >= output.len() {
return;
}
let mut remaining = cut;
let mut marked = false;
for part in content {
if let ContentPart::Text { text } = part {
if marked {
text.clear();
} else if text.len() > remaining {
text.truncate(remaining);
text.push_str("\n[truncated]");
marked = true;
} else {
remaining -= text.len();
}
}
}
}
}
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,
output_dir: Arc<OutputDir>,
gate: Option<Arc<dyn ToolGate>>,
}
#[derive(Default)]
struct OutputDir(OnceCell<TempDir>);
impl OutputDir {
async fn path(&self) -> std::io::Result<&Path> {
let dir = self
.0
.get_or_try_init(|| async {
tempfile::Builder::new()
.prefix("flashmind-output-")
.tempdir()
})
.await?;
Ok(dir.path())
}
}
impl ToolRegistry {
pub fn new() -> Self {
Self {
tools: HashMap::new(),
aliases: HashMap::new(),
truncate_on_overflow: false,
output_dir: Arc::default(),
gate: None,
}
}
pub fn set_gate(&mut self, gate: Arc<dyn ToolGate>) {
self.gate = Some(gate);
}
pub fn gate(&self) -> Option<&Arc<dyn ToolGate>> {
self.gate.as_ref()
}
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));
};
if let Some(gate) = &self.gate
&& let Some(refused) = pass_gate(gate.as_ref(), call, ctx.cancel_token).await
{
return refused;
}
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:#}")),
};
self.enforce_output_limits(
&call.id,
result,
tool.max_output_bytes(),
tool.max_output_lines(),
)
.await
}
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()
}
}
async fn pass_gate(
gate: &dyn ToolGate,
call: &crate::ToolCall,
cancel: &CancellationToken,
) -> Option<ToolResult> {
let ask = match gate.check(call).await {
GateDecision::Allow => return None,
GateDecision::Deny(reason) => return Some(ToolResult::failure(&call.id, reason)),
GateDecision::Ask(ask) => ask,
};
let answer = tokio::select! {
_ = cancel.cancelled() => {
return Some(ToolResult::failure(&call.id, format!("Tool '{}' cancelled", call.name)));
}
answer = tokio::time::timeout(ask.timeout, ask.answer) => answer,
};
match answer {
Ok(GateAnswer::Allow) => None,
Ok(GateAnswer::Deny(reason)) => Some(ToolResult::failure(&call.id, reason)),
Err(_) => Some(ToolResult::failure(
&call.id,
format!(
"Not done: nobody answered in time to allow '{}'.",
call.name
),
)),
}
}
impl ToolRegistry {
async fn enforce_output_limits(
&self,
tool_call_id: &str,
mut result: ToolResult,
max_bytes: usize,
max_lines: usize,
) -> ToolResult {
if result.is_interrupt() {
return result;
}
let output = result.output();
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 self.truncate_on_overflow {
result.truncate_output(max_bytes, max_lines);
return result;
}
match self.save_full_output(&output).await {
Ok(path) => {
result.truncate_output(max_bytes, max_lines);
result.append_text(&format!(
"\n[full output ({bytes} bytes, {lines} lines) saved to {}]",
path.display()
));
result
}
Err(e) => {
tracing::warn!(error = %e, "could not save oversized tool output");
if !result.is_success() {
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."
),
)
}
}
}
}
async fn save_full_output(&self, output: &str) -> std::io::Result<PathBuf> {
let path = self
.output_dir
.path()
.await?
.join(format!("{}.txt", Uuid::new_v4()));
let mut options = OpenOptions::new();
options.write(true).create_new(true);
#[cfg(unix)]
options.mode(0o600);
let mut file = options.open(&path).await?;
file.write_all(output.as_bytes()).await?;
file.flush().await?;
Ok(path)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn saved_path(output: &str) -> PathBuf {
let start = output
.rfind("saved to ")
.expect("output names a saved file")
+ 9;
PathBuf::from(output[start..].trim_end_matches(']'))
}
#[tokio::test]
async fn failed_output_is_truncated_not_passed_through() {
let output = "x\n".repeat(100);
let mut registry = ToolRegistry::new();
registry.set_truncate_on_overflow(true);
let result = registry
.enforce_output_limits(
"call",
ToolResult::failure("call", output.clone()),
50,
usize::MAX,
)
.await;
assert!(!result.is_success());
let output = result.output();
assert!(output.len() < 70, "{output}");
assert!(output.ends_with("[truncated]"), "{output}");
}
#[tokio::test]
async fn oversized_output_is_saved_to_temp_file() {
let registry = ToolRegistry::new();
let full = "line\n".repeat(100);
let mut paths = Vec::new();
for result in [
ToolResult::success("call", full.clone()),
ToolResult::failure("call", full.clone()),
] {
let was_success = result.is_success();
let result = registry
.enforce_output_limits("call", result, usize::MAX, 10)
.await;
assert_eq!(result.is_success(), was_success);
let output = result.output();
assert!(output.starts_with(&"line\n".repeat(9)), "{output}");
assert!(output.contains("[truncated]"), "{output}");
assert!(output.contains("(500 bytes, 100 lines)"), "{output}");
let path = saved_path(&output);
let saved = std::fs::read_to_string(&path).expect("saved file is readable");
assert_eq!(saved, full);
paths.push(path);
}
let clone = registry.clone();
drop(registry);
assert!(paths.iter().all(|p| p.exists()));
drop(clone);
assert!(paths.iter().all(|p| !p.exists()));
}
#[test]
fn mixed_result_projects_text_and_keeps_ordered_content() {
let mut result = ToolResult::success_with_content(
"call",
vec![
ContentPart::Text {
text: "before ".into(),
},
ContentPart::Image {
media_type: "image/png".into(),
data: "encoded".into(),
},
ContentPart::Text {
text: "after".into(),
},
],
);
assert_eq!(result.output(), "before after");
assert!(
matches!(&result.content()[1], ContentPart::Image { data, .. } if data == "encoded")
);
result.append_text("!");
assert_eq!(result.output(), "before after!");
assert!(!format!("{result:?}").contains("encoded"));
}
#[tokio::test]
async fn limits_apply_to_combined_text_and_leave_images_untouched() {
let registry = ToolRegistry::new();
let result = ToolResult::success_with_content(
"call",
vec![
ContentPart::Text { text: "ab".into() },
ContentPart::Image {
media_type: "image/png".into(),
data: "encoded".into(),
},
ContentPart::Text {
text: "éxyz".into(),
},
],
);
let result = registry
.enforce_output_limits("call", result, 3, usize::MAX)
.await;
assert!(result.is_success());
assert!(result.output().starts_with("ab\n[truncated]"));
assert!(
matches!(&result.content()[1], ContentPart::Image { data, .. } if data == "encoded")
);
assert_eq!(
std::fs::read_to_string(saved_path(&result.output())).unwrap(),
"abéxyz"
);
}
#[test]
fn failure_content_preserves_media_and_text_projection() {
let mut result = ToolResult::failure_with_content(
"call",
vec![
ContentPart::Text {
text: "failed".into(),
},
ContentPart::Image {
media_type: "image/png".into(),
data: "encoded".into(),
},
],
);
result.append_text(" details");
assert!(!result.is_success());
assert_eq!(result.output(), "failed details");
assert_eq!(result.content().len(), 3);
}
#[tokio::test]
async fn output_within_limits_is_unchanged() {
let result = ToolRegistry::new()
.enforce_output_limits("call", ToolResult::success("call", "ok"), 10, 10)
.await;
assert_eq!(result.output(), "ok");
}
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()
);
}
struct CountedTool(Arc<std::sync::atomic::AtomicUsize>);
#[async_trait]
impl Tool for CountedTool {
fn name(&self) -> &str {
"act"
}
fn description(&self) -> &str {
"acts"
}
fn parameters(&self) -> Value {
serde_json::json!({})
}
async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Ok(ToolResult::success(ctx.tool_call_id, "acted"))
}
}
struct ScriptedGate(std::sync::Mutex<Vec<GateDecision>>);
#[async_trait]
impl ToolGate for ScriptedGate {
async fn check(&self, _call: &crate::ToolCall) -> GateDecision {
self.0
.lock()
.expect("the test does not panic while holding it")
.pop()
.unwrap_or(GateDecision::Allow)
}
}
fn gated(decision: GateDecision) -> (ToolRegistry, Arc<std::sync::atomic::AtomicUsize>) {
let runs = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let mut registry = ToolRegistry::new();
registry.register(Arc::new(CountedTool(runs.clone())));
registry.set_gate(Arc::new(ScriptedGate(std::sync::Mutex::new(vec![
decision,
]))));
(registry, runs)
}
fn act_call() -> crate::ToolCall {
crate::ToolCall {
id: "c1".into(),
name: "act".into(),
arguments: serde_json::json!({}),
}
}
fn ask(
timeout: std::time::Duration,
) -> (GateDecision, tokio::sync::oneshot::Sender<GateAnswer>) {
let (answer, asked) = tokio::sync::oneshot::channel();
let decision = GateDecision::Ask(GateAsk {
answer: Box::pin(async move {
asked
.await
.unwrap_or(GateAnswer::Deny("nobody answered".into()))
}),
timeout,
});
(decision, answer)
}
#[tokio::test]
async fn a_denied_call_does_not_run_and_returns_the_reason() {
let (registry, runs) = gated(GateDecision::Deny("not allowed here".into()));
let result = registry
.execute(&act_call(), None, &CancellationToken::new())
.await;
assert!(!result.is_success());
assert_eq!(result.output(), "not allowed here");
assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
}
#[tokio::test]
async fn an_asked_call_waits_for_the_answer_then_runs() {
let (decision, answer) = ask(std::time::Duration::from_secs(30));
let (registry, runs) = gated(decision);
let registry = Arc::new(registry);
let running = registry.clone();
let call = tokio::spawn(async move {
running
.execute(&act_call(), None, &CancellationToken::new())
.await
});
tokio::task::yield_now().await;
assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
answer
.send(GateAnswer::Allow)
.expect("the call is still waiting");
let result = call.await.expect("the call does not panic");
assert!(result.is_success());
assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 1);
}
#[tokio::test]
async fn an_asked_call_refused_by_the_answer_does_not_run() {
let (decision, answer) = ask(std::time::Duration::from_secs(30));
let (registry, runs) = gated(decision);
answer
.send(GateAnswer::Deny("the user said no".into()))
.expect("the receiver is held by the gate's question");
let result = registry
.execute(&act_call(), None, &CancellationToken::new())
.await;
assert_eq!(result.output(), "the user said no");
assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
}
#[tokio::test]
async fn an_unanswered_call_is_refused_when_its_time_runs_out() {
let (decision, _answer) = ask(std::time::Duration::from_millis(20));
let (registry, runs) = gated(decision);
let result = registry
.execute(&act_call(), None, &CancellationToken::new())
.await;
assert!(!result.is_success());
assert!(result.output().contains("nobody answered in time"));
assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
}
#[tokio::test]
async fn clones_keep_the_gate() {
let (registry, runs) = gated(GateDecision::Deny("no".into()));
let copy = registry.clone();
assert!(copy.gate().is_some());
let result = copy
.execute(&act_call(), None, &CancellationToken::new())
.await;
assert!(!result.is_success());
assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
}
}