1use async_trait::async_trait;
16use futures_util::future::BoxFuture;
17use serde::de::DeserializeOwned;
18use serde::{Deserialize, Serialize};
19use serde_json::Value;
20use std::any::Any;
21use std::collections::HashMap;
22use std::path::{Path, PathBuf};
23use std::sync::Arc;
24use std::time::Duration;
25use tempfile::TempDir;
26use tokio::fs::OpenOptions;
27use tokio::io::AsyncWriteExt;
28use tokio::sync::OnceCell;
29use tokio_util::sync::CancellationToken;
30use uuid::Uuid;
31
32use crate::event::Source;
33use crate::llm::ToolDefinition;
34use crate::{AgentLlmConfig, ContentPart, LlmProvider};
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct ForbiddenCmd {
43 pub command: String,
45 pub reason: String,
47 #[serde(default)]
50 pub reconsider: bool,
51}
52
53impl ForbiddenCmd {
54 pub fn validate(&self) -> std::result::Result<(), Vec<String>> {
56 let mut errors = Vec::new();
57 if self.command.is_empty() {
58 errors.push("forbidden command name cannot be empty".to_string());
59 }
60 if regex::Regex::new(&self.command).is_err() {
61 errors.push(format!("invalid regex for name: '{}'", self.command));
62 }
63 if errors.is_empty() {
64 Ok(())
65 } else {
66 Err(errors)
67 }
68 }
69}
70
71pub trait CommandAllowList: Send + Sync {
81 fn is_allowed(&self, command: &str) -> bool;
83
84 fn add_session_pattern(&self, pattern: &str);
89}
90
91#[async_trait]
101pub trait ToolGate: Send + Sync {
102 async fn check(&self, call: &crate::ToolCall) -> GateDecision;
104}
105
106pub enum GateDecision {
108 Allow,
110 Deny(String),
112 Ask(GateAsk),
114}
115
116pub struct GateAsk {
118 pub answer: BoxFuture<'static, GateAnswer>,
120 pub timeout: Duration,
122}
123
124#[derive(Debug, Clone, PartialEq, Eq)]
126pub enum GateAnswer {
127 Allow,
128 Deny(String),
130}
131
132impl std::fmt::Debug for GateDecision {
133 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134 match self {
135 Self::Allow => f.write_str("Allow"),
136 Self::Deny(reason) => f.debug_tuple("Deny").field(reason).finish(),
137 Self::Ask(ask) => f
138 .debug_struct("Ask")
139 .field("timeout", &ask.timeout)
140 .finish(),
141 }
142 }
143}
144
145pub trait InterruptPayload: Send + Sync + std::fmt::Debug {
154 fn as_any(&self) -> &dyn Any;
156
157 fn display_output(&self) -> String;
160}
161
162#[derive(Debug, Clone, Serialize, Deserialize)]
168pub enum DiffLine {
169 Added { line: u64, content: String },
171 Removed { line: u64, content: String },
173 Context { line: u64, content: String },
176}
177
178#[derive(Debug, Clone, Serialize, Deserialize)]
180pub struct FileDiff {
181 pub path: String,
183 pub diff: Vec<DiffLine>,
185}
186
187#[derive(Clone)]
238pub enum ToolResult {
239 Success {
241 tool_call_id: String,
242 content: Vec<ContentPart>,
243 sources: Vec<Source>,
244 diffs: Vec<FileDiff>,
245 working_dir: Option<PathBuf>,
247 },
248 Failure {
250 tool_call_id: String,
251 content: Vec<ContentPart>,
252 },
253 Interrupt {
257 tool_call_id: String,
258 payload: Arc<dyn InterruptPayload>,
259 },
260}
261
262impl std::fmt::Debug for ToolResult {
263 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
264 match self {
265 Self::Success { tool_call_id, .. } => f
266 .debug_struct("Success")
267 .field("tool_call_id", tool_call_id)
268 .field("output", &self.output())
269 .finish(),
270 Self::Failure { tool_call_id, .. } => f
271 .debug_struct("Failure")
272 .field("tool_call_id", tool_call_id)
273 .field("output", &self.output())
274 .finish(),
275 Self::Interrupt {
276 tool_call_id,
277 payload,
278 } => f
279 .debug_struct("Interrupt")
280 .field("tool_call_id", tool_call_id)
281 .field("output", &payload.display_output())
282 .finish(),
283 }
284 }
285}
286
287impl ToolResult {
288 pub fn success(tool_call_id: &str, output: impl Into<String>) -> Self {
290 Self::success_with_content(
291 tool_call_id,
292 vec![ContentPart::Text {
293 text: output.into(),
294 }],
295 )
296 }
297
298 pub fn success_with_content(tool_call_id: &str, content: Vec<ContentPart>) -> Self {
300 Self::Success {
301 tool_call_id: tool_call_id.into(),
302 content,
303 sources: Vec::new(),
304 diffs: Vec::new(),
305 working_dir: None,
306 }
307 }
308
309 pub fn with_parts(mut self, parts: Vec<ContentPart>) -> Self {
311 if let Self::Success { content, .. } = &mut self {
312 content.extend(parts);
313 }
314 self
315 }
316
317 pub fn content(&self) -> &[ContentPart] {
319 match self {
320 Self::Success { content, .. } | Self::Failure { content, .. } => content,
321 Self::Interrupt { .. } => &[],
322 }
323 }
324
325 pub fn append_text(&mut self, text: &str) {
327 match self {
328 Self::Success { content, .. } | Self::Failure { content, .. } => {
329 match content.last_mut() {
330 Some(ContentPart::Text { text: last }) => last.push_str(text),
331 _ => content.push(ContentPart::Text {
332 text: text.to_owned(),
333 }),
334 }
335 }
336 Self::Interrupt { .. } => {}
337 }
338 }
339
340 pub fn failure(tool_call_id: &str, output: impl Into<String>) -> Self {
342 Self::failure_with_content(
343 tool_call_id,
344 vec![ContentPart::Text {
345 text: output.into(),
346 }],
347 )
348 }
349
350 pub fn failure_with_content(tool_call_id: &str, content: Vec<ContentPart>) -> Self {
352 Self::Failure {
353 tool_call_id: tool_call_id.into(),
354 content,
355 }
356 }
357
358 pub fn success_with_diffs(
361 tool_call_id: &str,
362 output: impl Into<String>,
363 diffs: Vec<FileDiff>,
364 ) -> Self {
365 Self::Success {
366 tool_call_id: tool_call_id.into(),
367 content: vec![ContentPart::Text {
368 text: output.into(),
369 }],
370 sources: Vec::new(),
371 diffs,
372 working_dir: None,
373 }
374 }
375
376 pub fn interrupt(tool_call_id: &str, payload: Arc<dyn InterruptPayload>) -> Self {
383 Self::Interrupt {
384 tool_call_id: tool_call_id.into(),
385 payload,
386 }
387 }
388
389 pub fn with_sources(mut self, sources: Vec<Source>) -> Self {
397 if let Self::Success {
398 sources: ref mut s, ..
399 } = self
400 {
401 *s = sources;
402 }
403 self
404 }
405
406 pub fn with_working_dir(mut self, path: PathBuf) -> Self {
409 if let Self::Success { working_dir, .. } = &mut self {
410 *working_dir = Some(path);
411 }
412 self
413 }
414
415 pub fn working_dir(&self) -> Option<&Path> {
417 match self {
418 Self::Success { working_dir, .. } => working_dir.as_deref(),
419 _ => None,
420 }
421 }
422
423 pub fn tool_call_id(&self) -> &str {
425 match self {
426 Self::Success { tool_call_id, .. }
427 | Self::Failure { tool_call_id, .. }
428 | Self::Interrupt { tool_call_id, .. } => tool_call_id,
429 }
430 }
431
432 pub fn output(&self) -> String {
435 match self {
436 Self::Success { content, .. } | Self::Failure { content, .. } => content
437 .iter()
438 .filter_map(|part| match part {
439 ContentPart::Text { text } => Some(text.as_str()),
440 _ => None,
441 })
442 .collect(),
443 Self::Interrupt { payload, .. } => payload.display_output(),
444 }
445 }
446
447 pub fn payload(&self) -> Option<&Arc<dyn InterruptPayload>> {
449 match self {
450 Self::Interrupt { payload, .. } => Some(payload),
451 _ => None,
452 }
453 }
454
455 pub fn is_success(&self) -> bool {
457 matches!(self, Self::Success { .. })
458 }
459
460 pub fn is_interrupt(&self) -> bool {
462 matches!(self, Self::Interrupt { .. })
463 }
464
465 pub fn sources(&self) -> &[Source] {
468 match self {
469 Self::Success { sources, .. } => sources,
470 _ => &[],
471 }
472 }
473
474 pub fn diffs(&self) -> &[FileDiff] {
477 match self {
478 Self::Success { diffs, .. } => diffs,
479 _ => &[],
480 }
481 }
482
483 pub fn truncate_output(&mut self, max_bytes: usize, max_lines: usize) {
486 let output = self.output();
487 let content = match self {
488 Self::Success { content, .. } | Self::Failure { content, .. } => content,
489 Self::Interrupt { .. } => return,
490 };
491 if output.len() <= max_bytes && max_lines == usize::MAX {
492 return;
493 }
494 let byte_cut = output.floor_char_boundary(max_bytes);
495 let line_cut = if max_lines == 0 {
496 0
497 } else {
498 output
499 .match_indices('\n')
500 .nth(max_lines - 1)
501 .map(|(index, _)| index)
502 .unwrap_or(output.len())
503 };
504 let cut = byte_cut.min(line_cut);
505 if cut >= output.len() {
506 return;
507 }
508 let mut remaining = cut;
509 let mut marked = false;
510 for part in content {
511 if let ContentPart::Text { text } = part {
512 if marked {
513 text.clear();
514 } else if text.len() > remaining {
515 text.truncate(remaining);
516 text.push_str("\n[truncated]");
517 marked = true;
518 } else {
519 remaining -= text.len();
520 }
521 }
522 }
523 }
524}
525
526pub fn parse_args<T: DeserializeOwned>(tool_name: &str, args: Value) -> anyhow::Result<T> {
533 serde_json::from_value(args)
534 .map_err(|e| anyhow::anyhow!("Tool '{}': Invalid arguments: {}", tool_name, e))
535}
536
537#[async_trait]
561pub trait Tool: Send + Sync {
562 fn name(&self) -> &str;
564 fn description(&self) -> &str;
566 fn parameters(&self) -> Value;
568 async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult>;
570 fn timeout_secs(&self) -> Option<u64> {
572 None
573 }
574
575 fn max_output_bytes(&self) -> usize {
579 256 * 1024
580 }
581
582 fn max_output_lines(&self) -> usize {
585 10_000
586 }
587
588 fn concurrent(&self) -> bool {
593 true
594 }
595
596 fn supports_background(&self) -> bool {
598 false
599 }
600
601 fn humanize(&self, args: &Value) -> String {
604 let Some(obj) = args.as_object() else {
605 return self.name().to_string();
606 };
607 let keys: Vec<&str> = obj.keys().map(|k| k.as_str()).collect();
608 if keys.is_empty() {
609 return self.name().to_string();
610 }
611 let summary = keys.join(", ");
612 if summary.len() > 80 {
613 let truncated: String = summary.chars().take(77).collect();
614 format!("{truncated}…")
615 } else {
616 summary
617 }
618 }
619}
620
621#[derive(Clone, Debug, Default)]
627pub struct BackgroundRequest {
628 requested: CancellationToken,
629}
630
631impl BackgroundRequest {
632 pub fn request(&self) {
633 self.requested.cancel();
634 }
635
636 pub async fn requested(&self) {
637 self.requested.cancelled().await;
638 }
639}
640
641pub struct ToolContext<'a> {
654 pub tool_call_id: &'a str,
656 pub args: Value,
658 pub llm: Option<&'a AgentLlmConfig>,
660 pub provider: Option<Arc<dyn LlmProvider>>,
662 pub working_dir: Option<&'a PathBuf>,
664 cancel_token: &'a CancellationToken,
666 progress_tx: Option<tokio::sync::mpsc::UnboundedSender<String>>,
668 background_request: Option<&'a BackgroundRequest>,
669}
670
671impl<'a> ToolContext<'a> {
672 pub fn new(
674 tool_call_id: &'a str,
675 args: Value,
676 working_dir: Option<&'a PathBuf>,
677 cancel_token: &'a CancellationToken,
678 ) -> Self {
679 Self {
680 tool_call_id,
681 args,
682 llm: None,
683 provider: None,
684 working_dir,
685 cancel_token,
686 progress_tx: None,
687 background_request: None,
688 }
689 }
690
691 pub fn with_llm(mut self, llm: &'a AgentLlmConfig) -> Self {
693 self.llm = Some(llm);
694 self
695 }
696
697 pub fn with_provider(mut self, provider: Arc<dyn LlmProvider>) -> Self {
699 self.provider = Some(provider);
700 self
701 }
702
703 pub fn with_progress(mut self, tx: tokio::sync::mpsc::UnboundedSender<String>) -> Self {
705 self.progress_tx = Some(tx);
706 self
707 }
708
709 pub fn with_background_request(mut self, request: &'a BackgroundRequest) -> Self {
711 self.background_request = Some(request);
712 self
713 }
714
715 pub fn background_request(&self) -> Option<&BackgroundRequest> {
716 self.background_request
717 }
718
719 pub fn progress_sender(&self) -> Option<tokio::sync::mpsc::UnboundedSender<String>> {
721 self.progress_tx.clone()
722 }
723
724 pub fn progress(&self, line: &str) {
726 if let Some(ref tx) = self.progress_tx {
727 let _ = tx.send(line.to_string());
728 }
729 }
730
731 pub fn parse_args<T: DeserializeOwned>(&self, tool_name: &str) -> anyhow::Result<T> {
757 T::deserialize(&self.args)
758 .map_err(|e| anyhow::anyhow!("Tool '{}': Invalid arguments: {}", tool_name, e))
759 }
760
761 pub fn cancel_token(&self) -> &CancellationToken {
766 self.cancel_token
767 }
768
769 pub fn child_token(&self) -> CancellationToken {
772 self.cancel_token.child_token()
773 }
774}
775
776#[derive(Clone)]
786pub struct ToolRegistry {
787 tools: HashMap<String, Arc<dyn Tool>>,
788 aliases: HashMap<String, String>,
789 truncate_on_overflow: bool,
790 output_dir: Arc<OutputDir>,
791 gate: Option<Arc<dyn ToolGate>>,
792}
793
794#[derive(Default)]
797struct OutputDir(OnceCell<TempDir>);
798
799impl OutputDir {
800 async fn path(&self) -> std::io::Result<&Path> {
801 let dir = self
802 .0
803 .get_or_try_init(|| async {
804 tempfile::Builder::new()
805 .prefix("flashmind-output-")
806 .tempdir()
807 })
808 .await?;
809 Ok(dir.path())
810 }
811}
812
813impl ToolRegistry {
814 pub fn new() -> Self {
816 Self {
817 tools: HashMap::new(),
818 aliases: HashMap::new(),
819 truncate_on_overflow: false,
820 output_dir: Arc::default(),
821 gate: None,
822 }
823 }
824
825 pub fn set_gate(&mut self, gate: Arc<dyn ToolGate>) {
828 self.gate = Some(gate);
829 }
830
831 pub fn gate(&self) -> Option<&Arc<dyn ToolGate>> {
833 self.gate.as_ref()
834 }
835
836 pub fn set_truncate_on_overflow(&mut self, enabled: bool) {
841 self.truncate_on_overflow = enabled;
842 }
843
844 pub fn register(&mut self, tool: Arc<dyn Tool>) {
846 tracing::debug!(tool = tool.name(), "registered tool");
847 self.tools.insert(tool.name().to_string(), tool);
848 }
849
850 pub fn remove(&mut self, name: &str) -> bool {
852 let removed = self.tools.remove(name).is_some();
853 if removed {
854 let tools = &self.tools;
855 self.aliases.retain(|_, target| tools.contains_key(target));
856 }
857 removed
858 }
859
860 pub fn strip_prefixes(&mut self, prefixes: &[&str]) {
863 let to_remove: Vec<String> = self
864 .tools
865 .keys()
866 .filter(|name| prefixes.iter().any(|p| name.starts_with(p)))
867 .cloned()
868 .collect();
869 for name in to_remove {
870 self.tools.remove(&name);
871 }
872 let tools = &self.tools;
873 self.aliases.retain(|_, target| tools.contains_key(target));
874 }
875
876 pub fn alias(&mut self, alias: &str, target: &str) {
879 tracing::debug!(alias, target, "registered tool alias");
880 self.aliases.insert(alias.to_string(), target.to_string());
881 }
882
883 pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
886 let result = self.tools.get(name).or_else(|| {
887 if let Some(target) = self.aliases.get(name) {
888 tracing::debug!(
889 alias = name,
890 target = target.as_str(),
891 "resolved tool alias"
892 );
893 self.tools.get(target)
894 } else {
895 None
896 }
897 });
898
899 if result.is_none() {
900 tracing::warn!(tool = name, "tool not found in registry");
901 }
902
903 result
904 }
905
906 pub fn contains(&self, name: &str) -> bool {
908 self.resolve(name).is_some()
909 }
910
911 fn resolve(&self, name: &str) -> Option<&Arc<dyn Tool>> {
913 self.tools.get(name).or_else(|| {
914 self.aliases
915 .get(name)
916 .and_then(|target| self.tools.get(target))
917 })
918 }
919
920 pub fn retain(&mut self, mut f: impl FnMut(&str) -> bool) {
922 self.tools.retain(|name, _| f(name));
923 let tools = &self.tools;
924 self.aliases.retain(|_, target| tools.contains_key(target));
925 }
926
927 pub fn list(&self) -> Vec<&str> {
929 let mut names: Vec<&str> = self.tools.keys().map(|s| s.as_str()).collect();
930 names.sort();
931 names
932 }
933
934 pub fn list_with_descriptions(&self) -> Vec<(&str, &str)> {
936 let mut items: Vec<(&str, &str)> = self
937 .tools
938 .values()
939 .map(|t| (t.name(), t.description()))
940 .collect();
941 items.sort_by_key(|(name, _)| *name);
942 items
943 }
944
945 pub async fn execute(
953 &self,
954 call: &crate::ToolCall,
955 working_dir: Option<&PathBuf>,
956 cancel_token: &CancellationToken,
957 ) -> ToolResult {
958 self.execute_inner(call, working_dir, cancel_token, None)
959 .await
960 }
961
962 pub async fn execute_with_progress(
965 &self,
966 call: &crate::ToolCall,
967 working_dir: Option<&PathBuf>,
968 cancel_token: &CancellationToken,
969 progress_tx: tokio::sync::mpsc::UnboundedSender<String>,
970 ) -> ToolResult {
971 self.execute_inner(call, working_dir, cancel_token, Some(progress_tx))
972 .await
973 }
974
975 async fn execute_inner(
976 &self,
977 call: &crate::ToolCall,
978 working_dir: Option<&PathBuf>,
979 cancel_token: &CancellationToken,
980 progress_tx: Option<tokio::sync::mpsc::UnboundedSender<String>>,
981 ) -> ToolResult {
982 let mut ctx = ToolContext::new(&call.id, call.arguments.clone(), working_dir, cancel_token);
983 if let Some(tx) = progress_tx {
984 ctx = ctx.with_progress(tx);
985 }
986 self.execute_in_context(call, ctx).await
987 }
988
989 pub async fn execute_in_context(
991 &self,
992 call: &crate::ToolCall,
993 ctx: ToolContext<'_>,
994 ) -> ToolResult {
995 let Some(tool) = self.get(&call.name) else {
996 return ToolResult::failure(&call.id, format!("Unknown tool: {}", call.name));
997 };
998 if let Some(gate) = &self.gate
999 && let Some(refused) = pass_gate(gate.as_ref(), call, ctx.cancel_token).await
1000 {
1001 return refused;
1002 }
1003 let cancel_token = ctx.cancel_token;
1004 let child_token = cancel_token.child_token();
1005 let ctx = ToolContext {
1006 cancel_token: &child_token,
1007 ..ctx
1008 };
1009 let fut = tool.execute(ctx);
1010
1011 let timeout_secs = tool.timeout_secs();
1014 let run = async {
1015 match timeout_secs {
1016 Some(secs) => tokio::time::timeout(std::time::Duration::from_secs(secs), fut)
1017 .await
1018 .map_err(|_| secs),
1019 None => Ok(fut.await),
1020 }
1021 };
1022 let outcome = tokio::select! {
1023 biased;
1024 _ = cancel_token.cancelled() => {
1025 child_token.cancel();
1026 return ToolResult::failure(
1027 &call.id,
1028 format!("Tool '{}' cancelled", call.name),
1029 );
1030 }
1031 r = run => match r {
1032 Ok(result) => result,
1033 Err(secs) => {
1034 child_token.cancel();
1035 return ToolResult::failure(
1036 &call.id,
1037 format!("Tool '{}' timed out after {secs}s", call.name),
1038 );
1039 }
1040 },
1041 };
1042
1043 let result = match outcome {
1044 Ok(result) => result,
1045 Err(e) => ToolResult::failure(&call.id, format!("Tool error: {e:#}")),
1046 };
1047
1048 self.enforce_output_limits(
1049 &call.id,
1050 result,
1051 tool.max_output_bytes(),
1052 tool.max_output_lines(),
1053 )
1054 .await
1055 }
1056
1057 pub fn is_concurrent(&self, call: &crate::ToolCall) -> bool {
1060 self.resolve(&call.name)
1061 .is_none_or(|tool| tool.concurrent())
1062 }
1063
1064 pub fn supports_background(&self, call: &crate::ToolCall) -> bool {
1065 self.resolve(&call.name)
1066 .is_some_and(|tool| tool.supports_background())
1067 }
1068
1069 pub fn humanize(&self, call: &crate::ToolCall) -> String {
1071 self.get(&call.name)
1072 .map(|t| t.humanize(&call.arguments))
1073 .unwrap_or_else(|| call.name.clone())
1074 }
1075
1076 pub fn definitions(&self) -> Vec<ToolDefinition> {
1078 let mut defs: Vec<ToolDefinition> = self
1079 .tools
1080 .values()
1081 .map(|t| ToolDefinition {
1082 name: t.name().to_string(),
1083 description: t.description().to_string(),
1084 parameters: t.parameters(),
1085 })
1086 .collect();
1087
1088 defs.sort_by(|a, b| a.name.cmp(&b.name));
1089
1090 tracing::debug!(count = defs.len(), "generated tool definitions");
1091
1092 defs
1093 }
1094}
1095
1096impl Default for ToolRegistry {
1097 fn default() -> Self {
1098 Self::new()
1099 }
1100}
1101
1102async fn pass_gate(
1105 gate: &dyn ToolGate,
1106 call: &crate::ToolCall,
1107 cancel: &CancellationToken,
1108) -> Option<ToolResult> {
1109 let ask = match gate.check(call).await {
1110 GateDecision::Allow => return None,
1111 GateDecision::Deny(reason) => return Some(ToolResult::failure(&call.id, reason)),
1112 GateDecision::Ask(ask) => ask,
1113 };
1114 let answer = tokio::select! {
1115 _ = cancel.cancelled() => {
1116 return Some(ToolResult::failure(&call.id, format!("Tool '{}' cancelled", call.name)));
1117 }
1118 answer = tokio::time::timeout(ask.timeout, ask.answer) => answer,
1119 };
1120 match answer {
1121 Ok(GateAnswer::Allow) => None,
1122 Ok(GateAnswer::Deny(reason)) => Some(ToolResult::failure(&call.id, reason)),
1123 Err(_) => Some(ToolResult::failure(
1124 &call.id,
1125 format!(
1126 "Not done: nobody answered in time to allow '{}'.",
1127 call.name
1128 ),
1129 )),
1130 }
1131}
1132
1133impl ToolRegistry {
1134 async fn enforce_output_limits(
1144 &self,
1145 tool_call_id: &str,
1146 mut result: ToolResult,
1147 max_bytes: usize,
1148 max_lines: usize,
1149 ) -> ToolResult {
1150 if result.is_interrupt() {
1151 return result;
1152 }
1153 let output = result.output();
1154 let bytes = output.len();
1155 let lines = output.lines().count();
1156 let over_bytes = bytes > max_bytes;
1157 let over_lines = lines > max_lines;
1158
1159 if !over_bytes && !over_lines {
1160 return result;
1161 }
1162
1163 if self.truncate_on_overflow {
1164 result.truncate_output(max_bytes, max_lines);
1165 return result;
1166 }
1167
1168 match self.save_full_output(&output).await {
1169 Ok(path) => {
1170 result.truncate_output(max_bytes, max_lines);
1171 result.append_text(&format!(
1172 "\n[full output ({bytes} bytes, {lines} lines) saved to {}]",
1173 path.display()
1174 ));
1175 result
1176 }
1177 Err(e) => {
1178 tracing::warn!(error = %e, "could not save oversized tool output");
1179 if !result.is_success() {
1180 result.truncate_output(max_bytes, max_lines);
1181 return result;
1182 }
1183 if over_bytes {
1184 ToolResult::failure(
1185 tool_call_id,
1186 format!(
1187 "Output too large ({bytes} bytes, limit {max_bytes}). \
1188 Use more specific arguments to reduce output size."
1189 ),
1190 )
1191 } else {
1192 ToolResult::failure(
1193 tool_call_id,
1194 format!(
1195 "Output too many lines ({lines}, limit {max_lines}). \
1196 Use more specific arguments to reduce output size."
1197 ),
1198 )
1199 }
1200 }
1201 }
1202 }
1203
1204 async fn save_full_output(&self, output: &str) -> std::io::Result<PathBuf> {
1207 let path = self
1208 .output_dir
1209 .path()
1210 .await?
1211 .join(format!("{}.txt", Uuid::new_v4()));
1212 let mut options = OpenOptions::new();
1213 options.write(true).create_new(true);
1214 #[cfg(unix)]
1215 options.mode(0o600);
1216 let mut file = options.open(&path).await?;
1217 file.write_all(output.as_bytes()).await?;
1218 file.flush().await?;
1219 Ok(path)
1220 }
1221}
1222
1223#[cfg(test)]
1224mod tests {
1225 use super::*;
1226
1227 fn saved_path(output: &str) -> PathBuf {
1229 let start = output
1230 .rfind("saved to ")
1231 .expect("output names a saved file")
1232 + 9;
1233 PathBuf::from(output[start..].trim_end_matches(']'))
1234 }
1235
1236 #[tokio::test]
1237 async fn failed_output_is_truncated_not_passed_through() {
1238 let output = "x\n".repeat(100);
1239 let mut registry = ToolRegistry::new();
1240 registry.set_truncate_on_overflow(true);
1241 let result = registry
1242 .enforce_output_limits(
1243 "call",
1244 ToolResult::failure("call", output.clone()),
1245 50,
1246 usize::MAX,
1247 )
1248 .await;
1249 assert!(!result.is_success());
1250 let output = result.output();
1251 assert!(output.len() < 70, "{output}");
1252 assert!(output.ends_with("[truncated]"), "{output}");
1253 }
1254
1255 #[tokio::test]
1256 async fn oversized_output_is_saved_to_temp_file() {
1257 let registry = ToolRegistry::new();
1258 let full = "line\n".repeat(100);
1259 let mut paths = Vec::new();
1260 for result in [
1261 ToolResult::success("call", full.clone()),
1262 ToolResult::failure("call", full.clone()),
1263 ] {
1264 let was_success = result.is_success();
1265 let result = registry
1266 .enforce_output_limits("call", result, usize::MAX, 10)
1267 .await;
1268 assert_eq!(result.is_success(), was_success);
1269
1270 let output = result.output();
1271 assert!(output.starts_with(&"line\n".repeat(9)), "{output}");
1272 assert!(output.contains("[truncated]"), "{output}");
1273 assert!(output.contains("(500 bytes, 100 lines)"), "{output}");
1274
1275 let path = saved_path(&output);
1276 let saved = std::fs::read_to_string(&path).expect("saved file is readable");
1277 assert_eq!(saved, full);
1278 paths.push(path);
1279 }
1280
1281 let clone = registry.clone();
1283 drop(registry);
1284 assert!(paths.iter().all(|p| p.exists()));
1285 drop(clone);
1286 assert!(paths.iter().all(|p| !p.exists()));
1287 }
1288
1289 #[test]
1290 fn mixed_result_projects_text_and_keeps_ordered_content() {
1291 let mut result = ToolResult::success_with_content(
1292 "call",
1293 vec![
1294 ContentPart::Text {
1295 text: "before ".into(),
1296 },
1297 ContentPart::Image {
1298 media_type: "image/png".into(),
1299 data: "encoded".into(),
1300 },
1301 ContentPart::Text {
1302 text: "after".into(),
1303 },
1304 ],
1305 );
1306 assert_eq!(result.output(), "before after");
1307 assert!(
1308 matches!(&result.content()[1], ContentPart::Image { data, .. } if data == "encoded")
1309 );
1310 result.append_text("!");
1311 assert_eq!(result.output(), "before after!");
1312 assert!(!format!("{result:?}").contains("encoded"));
1313 }
1314
1315 #[tokio::test]
1316 async fn limits_apply_to_combined_text_and_leave_images_untouched() {
1317 let registry = ToolRegistry::new();
1318 let result = ToolResult::success_with_content(
1319 "call",
1320 vec![
1321 ContentPart::Text { text: "ab".into() },
1322 ContentPart::Image {
1323 media_type: "image/png".into(),
1324 data: "encoded".into(),
1325 },
1326 ContentPart::Text {
1327 text: "éxyz".into(),
1328 },
1329 ],
1330 );
1331 let result = registry
1332 .enforce_output_limits("call", result, 3, usize::MAX)
1333 .await;
1334 assert!(result.is_success());
1335 assert!(result.output().starts_with("ab\n[truncated]"));
1336 assert!(
1337 matches!(&result.content()[1], ContentPart::Image { data, .. } if data == "encoded")
1338 );
1339 assert_eq!(
1340 std::fs::read_to_string(saved_path(&result.output())).unwrap(),
1341 "abéxyz"
1342 );
1343 }
1344
1345 #[test]
1346 fn failure_content_preserves_media_and_text_projection() {
1347 let mut result = ToolResult::failure_with_content(
1348 "call",
1349 vec![
1350 ContentPart::Text {
1351 text: "failed".into(),
1352 },
1353 ContentPart::Image {
1354 media_type: "image/png".into(),
1355 data: "encoded".into(),
1356 },
1357 ],
1358 );
1359 result.append_text(" details");
1360 assert!(!result.is_success());
1361 assert_eq!(result.output(), "failed details");
1362 assert_eq!(result.content().len(), 3);
1363 }
1364
1365 #[tokio::test]
1366 async fn output_within_limits_is_unchanged() {
1367 let result = ToolRegistry::new()
1368 .enforce_output_limits("call", ToolResult::success("call", "ok"), 10, 10)
1369 .await;
1370 assert_eq!(result.output(), "ok");
1371 }
1372
1373 struct DummyTool(&'static str);
1374
1375 #[async_trait]
1376 impl Tool for DummyTool {
1377 fn name(&self) -> &str {
1378 self.0
1379 }
1380 fn description(&self) -> &str {
1381 "dummy"
1382 }
1383 fn parameters(&self) -> Value {
1384 serde_json::json!({})
1385 }
1386 async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
1387 Ok(ToolResult::success(ctx.tool_call_id, "ok"))
1388 }
1389 }
1390
1391 #[test]
1392 fn test_registry_empty() {
1393 let registry = ToolRegistry::new();
1394 assert_eq!(registry.list().len(), 0);
1395 }
1396
1397 #[test]
1398 fn test_registry_definitions_empty() {
1399 let registry = ToolRegistry::new();
1400 assert!(registry.definitions().is_empty());
1401 }
1402
1403 #[test]
1404 fn test_registry_lookup() {
1405 let mut registry = ToolRegistry::new();
1406 registry.register(Arc::new(DummyTool("exec")));
1407 registry.register(Arc::new(DummyTool("slack_react")));
1408 registry.register(Arc::new(DummyTool("browser")));
1409
1410 assert_eq!(registry.definitions().len(), 3);
1411
1412 assert!(registry.get("slack_react").is_some());
1414 }
1415
1416 #[test]
1417 fn test_registry_alias_contains_and_remove_cleanup() {
1418 let mut registry = ToolRegistry::new();
1419 registry.register(Arc::new(DummyTool("exec")));
1420 registry.alias("bash_exec", "exec");
1421
1422 assert!(registry.contains("exec"));
1423 assert!(registry.contains("bash_exec"));
1424 assert!(registry.get("bash_exec").is_some());
1425
1426 assert!(registry.remove("exec"));
1427 assert!(!registry.contains("exec"));
1428 assert!(!registry.contains("bash_exec"));
1429 assert!(registry.get("bash_exec").is_none());
1430 }
1431
1432 struct SlowTimeoutTool;
1435
1436 #[async_trait]
1437 impl Tool for SlowTimeoutTool {
1438 fn name(&self) -> &str {
1439 "slow"
1440 }
1441 fn description(&self) -> &str {
1442 "sleeps forever"
1443 }
1444 fn parameters(&self) -> Value {
1445 serde_json::json!({})
1446 }
1447 fn timeout_secs(&self) -> Option<u64> {
1448 Some(600)
1449 }
1450 async fn execute(&self, _ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
1451 tokio::time::sleep(std::time::Duration::from_secs(60)).await;
1453 Ok(ToolResult::success("unused", "should not return"))
1454 }
1455 }
1456
1457 #[tokio::test]
1458 async fn registry_cancel_short_circuits_tool_with_timeout_secs() {
1459 let mut registry = ToolRegistry::new();
1462 registry.register(Arc::new(SlowTimeoutTool));
1463
1464 let cancel = CancellationToken::new();
1465 let cancel_clone = cancel.clone();
1466
1467 let call = crate::ToolCall {
1468 id: "c1".into(),
1469 name: "slow".into(),
1470 arguments: serde_json::json!({}),
1471 };
1472
1473 let registry = Arc::new(registry);
1474 let registry2 = registry.clone();
1475 let exec = tokio::spawn(async move { registry2.execute(&call, None, &cancel).await });
1476
1477 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1479 cancel_clone.cancel();
1480
1481 let result = tokio::time::timeout(std::time::Duration::from_secs(2), exec)
1483 .await
1484 .expect("registry did not return within 2s of cancellation")
1485 .unwrap();
1486
1487 assert!(!result.is_success(), "expected failure on cancel");
1488 assert!(
1489 result.output().contains("cancelled"),
1490 "expected 'cancelled' in output, got: {}",
1491 result.output()
1492 );
1493 }
1494
1495 struct CountedTool(Arc<std::sync::atomic::AtomicUsize>);
1497
1498 #[async_trait]
1499 impl Tool for CountedTool {
1500 fn name(&self) -> &str {
1501 "act"
1502 }
1503 fn description(&self) -> &str {
1504 "acts"
1505 }
1506 fn parameters(&self) -> Value {
1507 serde_json::json!({})
1508 }
1509 async fn execute(&self, ctx: ToolContext<'_>) -> anyhow::Result<ToolResult> {
1510 self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1511 Ok(ToolResult::success(ctx.tool_call_id, "acted"))
1512 }
1513 }
1514
1515 struct ScriptedGate(std::sync::Mutex<Vec<GateDecision>>);
1517
1518 #[async_trait]
1519 impl ToolGate for ScriptedGate {
1520 async fn check(&self, _call: &crate::ToolCall) -> GateDecision {
1521 self.0
1522 .lock()
1523 .expect("the test does not panic while holding it")
1524 .pop()
1525 .unwrap_or(GateDecision::Allow)
1526 }
1527 }
1528
1529 fn gated(decision: GateDecision) -> (ToolRegistry, Arc<std::sync::atomic::AtomicUsize>) {
1530 let runs = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1531 let mut registry = ToolRegistry::new();
1532 registry.register(Arc::new(CountedTool(runs.clone())));
1533 registry.set_gate(Arc::new(ScriptedGate(std::sync::Mutex::new(vec![
1534 decision,
1535 ]))));
1536 (registry, runs)
1537 }
1538
1539 fn act_call() -> crate::ToolCall {
1540 crate::ToolCall {
1541 id: "c1".into(),
1542 name: "act".into(),
1543 arguments: serde_json::json!({}),
1544 }
1545 }
1546
1547 fn ask(
1548 timeout: std::time::Duration,
1549 ) -> (GateDecision, tokio::sync::oneshot::Sender<GateAnswer>) {
1550 let (answer, asked) = tokio::sync::oneshot::channel();
1551 let decision = GateDecision::Ask(GateAsk {
1552 answer: Box::pin(async move {
1553 asked
1554 .await
1555 .unwrap_or(GateAnswer::Deny("nobody answered".into()))
1556 }),
1557 timeout,
1558 });
1559 (decision, answer)
1560 }
1561
1562 #[tokio::test]
1563 async fn a_denied_call_does_not_run_and_returns_the_reason() {
1564 let (registry, runs) = gated(GateDecision::Deny("not allowed here".into()));
1565 let result = registry
1566 .execute(&act_call(), None, &CancellationToken::new())
1567 .await;
1568 assert!(!result.is_success());
1569 assert_eq!(result.output(), "not allowed here");
1570 assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
1571 }
1572
1573 #[tokio::test]
1574 async fn an_asked_call_waits_for_the_answer_then_runs() {
1575 let (decision, answer) = ask(std::time::Duration::from_secs(30));
1576 let (registry, runs) = gated(decision);
1577 let registry = Arc::new(registry);
1578 let running = registry.clone();
1579 let call = tokio::spawn(async move {
1580 running
1581 .execute(&act_call(), None, &CancellationToken::new())
1582 .await
1583 });
1584 tokio::task::yield_now().await;
1585 assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
1586 answer
1587 .send(GateAnswer::Allow)
1588 .expect("the call is still waiting");
1589 let result = call.await.expect("the call does not panic");
1590 assert!(result.is_success());
1591 assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 1);
1592 }
1593
1594 #[tokio::test]
1595 async fn an_asked_call_refused_by_the_answer_does_not_run() {
1596 let (decision, answer) = ask(std::time::Duration::from_secs(30));
1597 let (registry, runs) = gated(decision);
1598 answer
1599 .send(GateAnswer::Deny("the user said no".into()))
1600 .expect("the receiver is held by the gate's question");
1601 let result = registry
1602 .execute(&act_call(), None, &CancellationToken::new())
1603 .await;
1604 assert_eq!(result.output(), "the user said no");
1605 assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
1606 }
1607
1608 #[tokio::test]
1609 async fn an_unanswered_call_is_refused_when_its_time_runs_out() {
1610 let (decision, _answer) = ask(std::time::Duration::from_millis(20));
1611 let (registry, runs) = gated(decision);
1612 let result = registry
1613 .execute(&act_call(), None, &CancellationToken::new())
1614 .await;
1615 assert!(!result.is_success());
1616 assert!(result.output().contains("nobody answered in time"));
1617 assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
1618 }
1619
1620 #[tokio::test]
1621 async fn clones_keep_the_gate() {
1622 let (registry, runs) = gated(GateDecision::Deny("no".into()));
1623 let copy = registry.clone();
1624 assert!(copy.gate().is_some());
1625 let result = copy
1626 .execute(&act_call(), None, &CancellationToken::new())
1627 .await;
1628 assert!(!result.is_success());
1629 assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), 0);
1630 }
1631}