1use async_trait::async_trait;
8use serde_json::Value;
9use std::collections::{BTreeMap, HashMap};
10use std::sync::{Arc, Mutex};
11use tokio::sync::RwLock;
12
13use crate::acp::ToolKind;
14use crate::agent::PermissionDecision;
15use crate::error::Result;
16use crate::llm::ToolSpec;
17use crate::permissions::auto_classifier::AutoClassifier;
18use crate::permissions::SharedPermissions;
19use crate::permissions::{PermissionMode, PermissionsConfig};
20use crate::tools::fs::ReadFileState;
21
22use super::audit::TouchedFiles;
23use super::policy_sandbox;
24
25pub type SpecWithHint = (ToolSpec, Option<String>);
28
29#[async_trait]
30pub trait Tool: Send + Sync {
31 fn spec(&self) -> ToolSpec;
32 async fn execute(&self, arguments: Value) -> Result<String>;
33
34 fn side_effect_class(&self) -> super::audit::ToolSideEffect {
39 super::audit::ToolSideEffect::External
40 }
41
42 fn is_deferred(&self) -> bool {
46 false
47 }
48
49 fn kind(&self) -> ToolKind {
53 ToolKind::Other
54 }
55
56 fn is_readonly(&self) -> bool {
61 matches!(
62 self.side_effect_class(),
63 super::audit::ToolSideEffect::ReadOnly
64 )
65 }
66
67 fn is_readonly_for_args(&self, _arguments: &Value) -> bool {
73 self.is_readonly()
74 }
75}
76
77#[async_trait]
86pub trait PermissionHook: Send + Sync {
87 async fn check(&self, tool_name: &str, args: &serde_json::Value) -> PermissionDecision;
89}
90
91pub struct PermissionHookDisabled;
96
97#[async_trait]
98impl PermissionHook for PermissionHookDisabled {
99 async fn check(&self, _tool_name: &str, _args: &serde_json::Value) -> PermissionDecision {
100 PermissionDecision::Allow
101 }
102}
103
104#[derive(Clone)]
106pub struct ToolRegistry {
107 tools: BTreeMap<String, Arc<dyn Tool>>,
108 aliases: BTreeMap<String, String>,
111 transport: Arc<dyn super::transport::ToolTransport>,
112 pub(crate) permissions: Option<SharedPermissions>,
117 pub(crate) permission_mode: PermissionMode,
120 pub(crate) touched: Option<Arc<Mutex<TouchedFiles>>>,
121 read_file_state: Option<Arc<Mutex<ReadFileState>>>,
124 session_roots: Option<super::dispatch::SharedSandboxRoots>,
128 #[cfg(feature = "mcp")]
130 elicitation: Option<crate::mcp::SharedElicitationHandler>,
131 pub(crate) permission_hook: Option<Arc<dyn PermissionHook>>,
135 pub(crate) policy: Option<policy_sandbox::PolicyConfig>,
139 pub(crate) headless: bool,
142 pub(crate) hook_runner: crate::hooks::ExternalHookRunner,
144
145 pub(crate) auto_classifier: Option<Arc<tokio::sync::Mutex<AutoClassifier>>>,
150}
151
152impl Default for ToolRegistry {
153 fn default() -> Self {
154 Self::local()
155 }
156}
157
158impl ToolRegistry {
159 pub fn new(transport: Arc<dyn super::transport::ToolTransport>) -> Self {
160 Self {
161 tools: BTreeMap::new(),
162 aliases: BTreeMap::new(),
163 transport,
164 permissions: None,
165 auto_classifier: None,
166 permission_mode: PermissionMode::Default,
167 touched: None,
168 read_file_state: None,
169 session_roots: None,
170 #[cfg(feature = "mcp")]
171 elicitation: None,
172 permission_hook: None,
173 policy: None,
174 headless: false,
175 hook_runner: crate::hooks::ExternalHookRunner::discover(&[]),
176 }
177 }
178
179 pub fn local() -> Self {
181 Self::new(Arc::new(super::transport::LocalTransport))
182 }
183
184 pub fn transport(&self) -> &Arc<dyn super::transport::ToolTransport> {
186 &self.transport
187 }
188
189 pub fn with_same_transport(&self) -> Self {
191 Self {
192 tools: BTreeMap::new(),
193 aliases: BTreeMap::new(),
194 transport: self.transport.clone(),
195 permissions: self.permissions.clone(),
196 auto_classifier: self.auto_classifier.clone(),
197 permission_mode: self.permission_mode.clone(),
198 touched: self.touched.clone(),
199 read_file_state: self.read_file_state.clone(),
200 session_roots: self.session_roots.clone(),
201 #[cfg(feature = "mcp")]
202 elicitation: self.elicitation.clone(),
203 permission_hook: self.permission_hook.clone(),
204 policy: self.policy.clone(),
205 headless: self.headless,
206 hook_runner: self.hook_runner.clone(),
207 }
208 }
209
210 pub fn fork(&self) -> Self {
223 self.clone()
224 }
225
226 pub fn with_permission_hook(mut self, hook: Arc<dyn PermissionHook>) -> Self {
230 self.permission_hook = Some(hook);
231 self
232 }
233
234 pub fn set_permission_hook(&mut self, hook: Arc<dyn PermissionHook>) {
237 self.permission_hook = Some(hook);
238 }
239
240 pub fn clear_permission_hook(&mut self) {
242 self.permission_hook = None;
243 }
244
245 pub fn with_policy(mut self, policy: policy_sandbox::PolicyConfig) -> Self {
249 self.policy = Some(policy);
250 self
251 }
252
253 pub fn set_policy(&mut self, policy: policy_sandbox::PolicyConfig) {
255 self.policy = Some(policy);
256 }
257
258 pub fn policy(&self) -> Option<&policy_sandbox::PolicyConfig> {
260 self.policy.as_ref()
261 }
262
263 pub fn with_headless(mut self, headless: bool) -> Self {
266 self.headless = headless;
267 self
268 }
269
270 pub fn set_headless(&mut self, headless: bool) {
272 self.headless = headless;
273 }
274
275 pub fn with_hook_runner(mut self, hook_runner: crate::hooks::ExternalHookRunner) -> Self {
277 self.hook_runner = hook_runner;
278 self
279 }
280
281 pub fn set_hook_runner(&mut self, hook_runner: crate::hooks::ExternalHookRunner) {
283 self.hook_runner = hook_runner;
284 }
285
286 pub fn with_permissions(mut self, permissions: PermissionsConfig) -> Self {
288 self.permission_mode = permissions.mode.clone();
289 self.permissions = Some(Arc::new(RwLock::new(permissions)));
290 self
291 }
292
293 pub fn with_shared_permissions(mut self, sp: SharedPermissions) -> Self {
301 if let Ok(guard) = sp.try_read() {
303 self.permission_mode = guard.mode.clone();
304 }
305 self.permissions = Some(sp);
306 self
307 }
308
309 pub fn with_auto_classifier(mut self, classifier: AutoClassifier) -> Self {
316 self.auto_classifier = Some(Arc::new(tokio::sync::Mutex::new(classifier)));
317 self
318 }
319
320 pub fn permission_mode(&self) -> PermissionMode {
322 self.permission_mode.clone()
323 }
324
325 pub fn set_permission_mode(&mut self, mode: PermissionMode) {
331 self.permission_mode = mode.clone();
332 match &self.permissions {
333 Some(sp) => {
334 if let Ok(mut guard) = sp.try_write() {
335 guard.mode = mode;
336 }
337 }
338 None => {
339 self.permissions = Some(Arc::new(RwLock::new(
340 crate::permissions::LayeredPermissionsConfig {
341 mode,
342 layers: Vec::new(),
343 },
344 )));
345 }
346 }
347 }
348
349 pub fn shared_permissions(&self) -> Option<SharedPermissions> {
351 self.permissions.clone()
352 }
353
354 pub fn permissions_config(&self) -> Option<PermissionsConfig> {
361 self.permissions
362 .as_ref()
363 .and_then(|sp| sp.try_read().ok())
364 .map(|guard| guard.clone())
365 }
366
367 pub fn is_plan_mode(&self, tool_name: &str) -> bool {
370 self.permissions
371 .as_ref()
372 .and_then(|sp| sp.try_read().ok())
373 .map(|guard| guard.is_plan_mode(tool_name))
374 .unwrap_or(false)
375 }
376
377 pub fn with_touched_files(mut self, slot: Arc<Mutex<TouchedFiles>>) -> Self {
382 self.touched = Some(slot);
383 self
384 }
385
386 pub fn clear_touched_files(&mut self) {
388 self.touched = None;
389 }
390
391 pub fn touched_files(&self) -> Option<Arc<Mutex<TouchedFiles>>> {
393 self.touched.clone()
394 }
395
396 pub fn with_read_file_state(mut self, slot: Arc<Mutex<ReadFileState>>) -> Self {
399 self.read_file_state = Some(slot);
400 self
401 }
402
403 pub fn read_file_state(&self) -> Option<Arc<Mutex<ReadFileState>>> {
405 self.read_file_state.clone()
406 }
407
408 pub fn with_session_roots(mut self, roots: super::dispatch::SharedSandboxRoots) -> Self {
410 self.session_roots = Some(roots);
411 self
412 }
413
414 pub fn session_roots(&self) -> Option<super::dispatch::SharedSandboxRoots> {
416 self.session_roots.clone()
417 }
418
419 #[cfg(feature = "mcp")]
421 pub fn with_elicitation_slot(mut self, slot: crate::mcp::SharedElicitationHandler) -> Self {
422 self.elicitation = Some(slot);
423 self
424 }
425
426 #[cfg(feature = "mcp")]
428 pub fn elicitation_slot(&self) -> Option<crate::mcp::SharedElicitationHandler> {
429 self.elicitation.clone()
430 }
431
432 pub fn register(mut self, tool: Arc<dyn Tool>) -> Self {
433 let name = tool.spec().name;
434 self.tools.insert(name, tool);
435 self
436 }
437
438 pub fn register_with_aliases(mut self, tool: Arc<dyn Tool>, aliases: &[&str]) -> Self {
444 let name = tool.spec().name.clone();
445 for &alias in aliases {
446 self.aliases.insert(alias.to_string(), name.clone());
447 }
448 self.tools.insert(name, tool);
449 self
450 }
451
452 pub fn register_mut(&mut self, tool: Arc<dyn Tool>) {
454 let name = tool.spec().name;
455 self.tools.insert(name, tool);
456 }
457
458 pub fn register_mut_with_aliases(&mut self, tool: Arc<dyn Tool>, aliases: &[&str]) {
460 let name = tool.spec().name.clone();
461 for &alias in aliases {
462 self.aliases.insert(alias.to_string(), name.clone());
463 }
464 self.tools.insert(name, tool);
465 }
466
467 pub fn find_by_name(&self, name: &str) -> Option<Arc<dyn Tool>> {
472 if let Some(tool) = self.tools.get(name) {
474 return Some(tool.clone());
475 }
476 if let Some(primary) = self.aliases.get(name) {
478 return self.tools.get(primary).cloned();
479 }
480 None
481 }
482
483 pub fn specs(&self) -> Vec<ToolSpec> {
484 self.tools.values().map(|t| t.spec()).collect()
485 }
486
487 pub fn specs_partitioned(&self) -> (Vec<ToolSpec>, Vec<ToolSpec>) {
492 let mut eager = Vec::new();
493 let mut deferred = Vec::new();
494 for tool in self.tools.values() {
495 if tool.is_deferred() {
496 deferred.push(tool.spec());
497 } else {
498 eager.push(tool.spec());
499 }
500 }
501 (eager, deferred)
502 }
503
504 pub fn retain_tools(&mut self, allow: &[String]) {
509 let allowed: std::collections::HashSet<String> =
510 allow.iter().map(|n| n.to_lowercase()).collect();
511 self.tools
512 .retain(|name, _| allowed.contains(&name.to_lowercase()));
513 self.aliases
514 .retain(|_, primary| self.tools.contains_key(primary));
515 }
516
517 pub fn split_eager_deferred(&self) -> (Vec<SpecWithHint>, Vec<SpecWithHint>) {
527 let mut eager = Vec::new();
528 let mut deferred = Vec::new();
529 for tool in self.tools.values() {
530 let spec = tool.spec();
531 let hint = spec
532 .description
533 .split('.')
534 .next()
535 .map(|s| s.trim().to_string())
536 .filter(|s| !s.is_empty());
537 if tool.is_deferred() {
538 deferred.push((spec, hint));
539 } else {
540 eager.push((spec, hint));
541 }
542 }
543 (eager, deferred)
544 }
545
546 pub fn is_deferred_spec(&self, spec: &ToolSpec) -> bool {
548 self.tools
549 .get(&spec.name)
550 .map(|t| t.is_deferred())
551 .unwrap_or(false)
552 }
553
554 pub fn freeze_deferred_specs(&mut self) {
560 let deferred_specs: Vec<ToolSpec> = self
561 .tools
562 .values()
563 .filter(|t| t.is_deferred())
564 .map(|t| t.spec())
565 .collect();
566
567 if deferred_specs.is_empty() {
568 return;
569 }
570
571 let catalog: super::tool_search::DeferredCatalog =
572 Arc::new(std::sync::RwLock::new(deferred_specs));
573 let tool = Arc::new(super::tool_search::ToolSearchTool::new(catalog));
574 self.tools
575 .insert(super::tool_search::TOOL_SEARCH_TOOL_NAME.to_string(), tool);
576 }
577
578 pub fn names(&self) -> Vec<String> {
579 self.tools.keys().cloned().collect()
580 }
581
582 pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
583 self.tools.get(name).cloned()
584 }
585
586 pub fn is_readonly(&self, name: &str) -> bool {
588 self.tools
589 .get(name)
590 .map(|t| t.is_readonly())
591 .unwrap_or(false)
592 }
593
594 pub fn build_kind_map(&self) -> HashMap<String, ToolKind> {
598 self.tools
599 .iter()
600 .map(|(name, tool)| (name.clone(), tool.kind()))
601 .collect()
602 }
603
604 pub fn is_readonly_for_call(&self, name: &str, args: &Value) -> bool {
608 self.tools
609 .get(name)
610 .map(|t| t.is_readonly_for_args(args))
611 .unwrap_or(false)
612 }
613}
614
615pub fn build_standard_tools(
624 workspace: &std::path::Path,
625 skills: &[crate::skills::Skill],
626 shell_timeout_secs: u64,
627) -> ToolRegistry {
628 build_standard_tools_with_roots(
629 workspace,
630 &[],
631 None,
632 skills,
633 shell_timeout_secs,
634 None,
635 None,
636 None,
637 None,
638 )
639}
640
641#[allow(clippy::too_many_arguments)]
672pub fn build_standard_tools_with_roots(
673 workspace: &std::path::Path,
674 extra_roots: &[(std::path::PathBuf, super::dispatch::AccessTier)],
675 session_roots: Option<super::dispatch::SharedSandboxRoots>,
676 skills: &[crate::skills::Skill],
677 shell_timeout_secs: u64,
678 web_search_provider: Option<String>,
679 web_search_api_key: Option<String>,
680 web_search_jina_key: Option<String>,
681 bg_manager: Option<Arc<tokio::sync::Mutex<super::run_background::BackgroundJobManager>>>,
682) -> ToolRegistry {
683 let bg_manager = bg_manager.unwrap_or_else(|| {
684 Arc::new(tokio::sync::Mutex::new(
685 super::run_background::BackgroundJobManager::new(),
686 ))
687 });
688 let todo_list = Arc::new(std::sync::RwLock::new(Vec::<super::todo::TodoItem>::new()));
689 let read_state = Arc::new(Mutex::new(ReadFileState::new()));
690 let mut registry = ToolRegistry::local()
691 .with_read_file_state(read_state.clone())
692 .register(Arc::new(
693 super::fs::ReadFile::new(workspace)
694 .with_extra_roots(extra_roots.iter().cloned())
695 .with_session_roots_opt(session_roots.clone())
696 .with_read_state(read_state.clone()),
697 ))
698 .register(Arc::new(
699 super::fs::WriteFile::new(workspace)
700 .with_extra_roots(extra_roots.iter().cloned())
701 .with_session_roots_opt(session_roots.clone())
702 .with_read_state(read_state.clone()),
703 ))
704 .register(Arc::new(
705 super::edit::EditTool::new(workspace)
706 .with_extra_roots(extra_roots.iter().cloned())
707 .with_session_roots_opt(session_roots.clone())
708 .with_read_state(read_state.clone()),
709 ))
710 .register(Arc::new(
711 super::shell::RunShell::new(workspace)
712 .with_timeout(std::time::Duration::from_secs(shell_timeout_secs)),
713 ))
714 .register(Arc::new(
715 super::search::SearchFiles::new(workspace)
716 .with_extra_roots(extra_roots.iter().cloned())
717 .with_session_roots_opt(session_roots.clone()),
718 ))
719 .register(Arc::new(
720 super::glob::GlobTool::new(workspace)
721 .with_extra_roots(extra_roots.iter().cloned())
722 .with_session_roots_opt(session_roots.clone()),
723 ))
724 .register(Arc::new(super::run_background::RunBackground::new(
725 workspace,
726 bg_manager.clone(),
727 )))
728 .register(Arc::new(super::run_background::CheckBackground::new(
729 bg_manager.clone(),
730 )))
731 .register(Arc::new(super::watch_file::WatchFile::new(
732 workspace,
733 bg_manager.clone(),
734 )))
735 .register(Arc::new(super::stop_loop::StopLoop::new(
736 workspace, bg_manager,
737 )))
738 .register(Arc::new(
739 super::estimate_tokens::EstimateTokens::new(workspace)
740 .with_extra_roots(extra_roots.iter().cloned())
741 .with_session_roots_opt(session_roots.clone()),
742 ))
743 .register(Arc::new(super::memory::Remember::new(workspace)))
744 .register(Arc::new(super::memory::Recall::new(workspace)))
745 .register(Arc::new(super::memory::Forget::new(workspace)))
746 .register(Arc::new(super::facts::RememberFact::new(workspace)))
747 .register(Arc::new(super::facts::RecallFact::new(workspace)))
748 .register(Arc::new(super::facts::ForgetFact::new(workspace)))
749 .register(Arc::new(super::facts::UpdateFact::new(workspace)))
750 .register(Arc::new(super::episodic_recall::EpisodicRecall::new(
751 workspace,
752 )))
753 .register(Arc::new(super::memory::WorkingMemoryTool::new(workspace)))
754 .register(Arc::new(super::memory::ScratchpadGet::new(workspace)))
755 .register(Arc::new(super::memory::ScratchpadDelete::new(workspace)))
756 .register(Arc::new(super::memory::ScratchpadList::new(workspace)))
757 .register(Arc::new(super::todo::TodoWriteTool::new(
758 todo_list,
759 Arc::new(crate::event::NullSink),
760 )))
761 .register(Arc::new(super::a2a::A2aCallTool::new()))
762 .register(Arc::new(super::a2a::A2aCardTool::new()))
763 .register(Arc::new(super::a2a::A2aTaskCheckTool::new()));
764
765 #[cfg(feature = "web_fetch")]
773 {
774 registry = registry.register(Arc::new(super::web_fetch::WebFetch::new()));
775 }
776
777 #[cfg(feature = "web_search")]
778 {
779 let search = super::web_search::WebSearch::new().with_search_config(
780 web_search_provider,
781 web_search_api_key,
782 web_search_jina_key,
783 );
784 registry = registry.register(Arc::new(search));
785 }
786
787 if !skills.is_empty() {
788 registry = registry.register(Arc::new(super::load_skill::LoadSkill::new(skills.to_vec())));
789 }
790
791 registry = registry.register(Arc::new(
797 super::client_fs::ClientReadFile::new(workspace)
798 .with_extra_roots(extra_roots.iter().cloned())
799 .with_session_roots_opt(session_roots.clone()),
800 ));
801 registry = registry.register(Arc::new(
802 super::client_fs::ClientWriteFile::new(workspace)
803 .with_extra_roots(extra_roots.iter().cloned())
804 .with_session_roots_opt(session_roots.clone()),
805 ));
806
807 registry
808}
809
810#[cfg(test)]
811mod tests {
812 use super::*;
813 use async_trait::async_trait;
814
815 struct ReadOnlyTool {
817 name: &'static str,
818 }
819
820 #[async_trait]
821 impl Tool for ReadOnlyTool {
822 fn spec(&self) -> ToolSpec {
823 ToolSpec {
824 name: self.name.to_string(),
825 description: "read-only test tool".into(),
826 parameters: serde_json::json!({"type":"object","properties":{}}),
827 }
828 }
829 fn side_effect_class(&self) -> super::super::audit::ToolSideEffect {
830 super::super::audit::ToolSideEffect::ReadOnly
831 }
832 fn kind(&self) -> ToolKind {
833 ToolKind::Read
834 }
835 async fn execute(&self, _args: Value) -> crate::error::Result<String> {
836 Ok("read-result".into())
837 }
838 }
839
840 struct MutatingTool {
842 name: &'static str,
843 }
844
845 #[async_trait]
846 impl Tool for MutatingTool {
847 fn spec(&self) -> ToolSpec {
848 ToolSpec {
849 name: self.name.to_string(),
850 description: "mutating test tool".into(),
851 parameters: serde_json::json!({"type":"object","properties":{}}),
852 }
853 }
854 fn kind(&self) -> ToolKind {
855 ToolKind::Other
856 }
857 async fn execute(&self, _args: Value) -> crate::error::Result<String> {
858 Ok("mutated".into())
859 }
860 }
861
862 fn make_registry() -> ToolRegistry {
863 ToolRegistry::local()
864 }
865
866 #[test]
869 fn register_and_find_by_name() {
870 let reg = make_registry().register(Arc::new(ReadOnlyTool { name: "Alpha" }));
871 let found = reg.find_by_name("Alpha");
872 assert!(found.is_some(), "registered tool must be findable by name");
873 assert!(reg.find_by_name("NoSuch").is_none());
874 }
875
876 #[test]
877 fn names_returns_registered_tool_names() {
878 let reg = make_registry()
879 .register(Arc::new(ReadOnlyTool { name: "T1" }))
880 .register(Arc::new(MutatingTool { name: "T2" }));
881 let mut names = reg.names();
882 names.sort();
883 assert_eq!(names, vec!["T1", "T2"]);
884 }
885
886 #[test]
887 fn names_empty_on_fresh_registry() {
888 assert!(make_registry().names().is_empty());
889 }
890
891 #[test]
892 fn specs_returns_tool_specs() {
893 let reg = make_registry().register(Arc::new(ReadOnlyTool { name: "SpecTool" }));
894 let specs = reg.specs();
895 assert_eq!(specs.len(), 1);
896 assert_eq!(specs[0].name, "SpecTool");
897 }
898
899 #[test]
900 fn get_returns_tool_by_name() {
901 let reg = make_registry().register(Arc::new(MutatingTool { name: "Getter" }));
902 assert!(reg.get("Getter").is_some());
903 assert!(reg.get("Missing").is_none());
904 }
905
906 #[test]
909 fn register_mut_adds_tool() {
910 let mut reg = make_registry();
911 reg.register_mut(Arc::new(ReadOnlyTool { name: "MutReg" }));
912 assert!(reg.find_by_name("MutReg").is_some());
913 }
914
915 #[test]
918 fn is_readonly_true_for_readonly_tool() {
919 let reg = make_registry().register(Arc::new(ReadOnlyTool { name: "RO" }));
920 assert!(
921 reg.is_readonly("RO"),
922 "ReadOnly tool must be is_readonly=true"
923 );
924 assert!(
925 !reg.is_readonly("Missing"),
926 "unknown tool must be is_readonly=false"
927 );
928 }
929
930 #[test]
931 fn is_readonly_false_for_mutating_tool() {
932 let reg = make_registry().register(Arc::new(MutatingTool { name: "Mut" }));
933 assert!(
934 !reg.is_readonly("Mut"),
935 "External tool must be is_readonly=false"
936 );
937 }
938
939 #[test]
940 fn is_readonly_for_call_delegates_to_tool() {
941 let reg = make_registry().register(Arc::new(ReadOnlyTool { name: "Ro2" }));
942 let args = serde_json::json!({});
943 assert!(reg.is_readonly_for_call("Ro2", &args));
944 assert!(!reg.is_readonly_for_call("Missing", &args));
945 }
946
947 #[test]
950 fn retain_tools_keeps_only_allowed() {
951 let mut reg = make_registry()
952 .register(Arc::new(ReadOnlyTool { name: "Keep" }))
953 .register(Arc::new(MutatingTool { name: "Drop" }));
954 reg.retain_tools(&["Keep".to_string()]);
955 assert!(reg.find_by_name("Keep").is_some());
956 assert!(reg.find_by_name("Drop").is_none());
957 }
958
959 #[test]
962 fn with_same_transport_returns_empty_registry() {
963 let reg = make_registry().register(Arc::new(ReadOnlyTool { name: "Src" }));
964 let empty = reg.with_same_transport();
965 assert!(
966 empty.names().is_empty(),
967 "with_same_transport must return empty registry"
968 );
969 }
970
971 #[test]
974 fn fork_clones_tools() {
975 let reg = make_registry().register(Arc::new(ReadOnlyTool { name: "ForkMe" }));
976 let forked = reg.fork();
977 assert!(forked.find_by_name("ForkMe").is_some());
978 }
979
980 #[test]
983 fn build_standard_tools_is_nonempty() {
984 let tmp = tempfile::tempdir().expect("tempdir");
985 let registry = build_standard_tools(tmp.path(), &[], 30);
986 assert!(
987 !registry.names().is_empty(),
988 "build_standard_tools must register at least one tool"
989 );
990 }
991
992 struct DeferredTool {
995 name: &'static str,
996 description: &'static str,
997 }
998
999 #[async_trait]
1000 impl Tool for DeferredTool {
1001 fn spec(&self) -> ToolSpec {
1002 ToolSpec {
1003 name: self.name.to_string(),
1004 description: self.description.to_string(),
1005 parameters: serde_json::json!({"type":"object","properties":{}}),
1006 }
1007 }
1008 fn is_deferred(&self) -> bool {
1009 true
1010 }
1011 fn kind(&self) -> ToolKind {
1012 ToolKind::Other
1013 }
1014 async fn execute(&self, _args: Value) -> crate::error::Result<String> {
1015 Ok("deferred-result".into())
1016 }
1017 }
1018
1019 #[test]
1022 fn register_with_aliases_finds_by_primary_and_alias() {
1023 let reg = make_registry().register_with_aliases(
1024 Arc::new(ReadOnlyTool { name: "Primary" }),
1025 &["alias1", "alias2"],
1026 );
1027 assert!(
1028 reg.find_by_name("Primary").is_some(),
1029 "primary name must work"
1030 );
1031 assert!(reg.find_by_name("alias1").is_some(), "alias1 must resolve");
1032 assert!(reg.find_by_name("alias2").is_some(), "alias2 must resolve");
1033 assert!(
1034 reg.find_by_name("unknown").is_none(),
1035 "unknown must return None"
1036 );
1037 }
1038
1039 #[test]
1042 fn register_mut_with_aliases_adds_alias() {
1043 let mut reg = make_registry();
1044 reg.register_mut_with_aliases(Arc::new(ReadOnlyTool { name: "MutAliased" }), &["malias"]);
1045 assert!(reg.find_by_name("MutAliased").is_some());
1046 assert!(
1047 reg.find_by_name("malias").is_some(),
1048 "mutably-registered alias must resolve"
1049 );
1050 }
1051
1052 #[test]
1055 fn specs_partitioned_separates_eager_from_deferred() {
1056 let reg = make_registry()
1057 .register(Arc::new(ReadOnlyTool { name: "Eager1" }))
1058 .register(Arc::new(ReadOnlyTool { name: "Eager2" }))
1059 .register(Arc::new(DeferredTool {
1060 name: "Defer1",
1061 description: "deferred.",
1062 }));
1063 let (eager, deferred) = reg.specs_partitioned();
1064 assert_eq!(eager.len(), 2, "two eager tools expected");
1065 assert_eq!(deferred.len(), 1, "one deferred tool expected");
1066 assert_eq!(deferred[0].name, "Defer1");
1067 }
1068
1069 #[test]
1070 fn specs_partitioned_all_eager_when_no_deferred() {
1071 let reg = make_registry()
1072 .register(Arc::new(ReadOnlyTool { name: "E1" }))
1073 .register(Arc::new(MutatingTool { name: "E2" }));
1074 let (eager, deferred) = reg.specs_partitioned();
1075 assert_eq!(eager.len(), 2);
1076 assert!(deferred.is_empty(), "no deferred tools expected");
1077 }
1078
1079 #[test]
1082 fn split_eager_deferred_puts_deferred_in_second_slot() {
1083 let reg = make_registry()
1084 .register(Arc::new(ReadOnlyTool { name: "EagerX" }))
1085 .register(Arc::new(DeferredTool {
1086 name: "DeferX",
1087 description: "First sentence. Second sentence.",
1088 }));
1089 let (eager, deferred) = reg.split_eager_deferred();
1090 assert_eq!(eager.len(), 1);
1091 assert_eq!(deferred.len(), 1);
1092 let (_spec, hint) = &deferred[0];
1094 assert_eq!(
1095 hint.as_deref(),
1096 Some("First sentence"),
1097 "hint must be first sentence"
1098 );
1099 }
1100
1101 #[test]
1102 fn split_eager_deferred_hint_is_none_for_empty_description() {
1103 let reg = make_registry().register(Arc::new(DeferredTool {
1104 name: "NoDesc",
1105 description: ".",
1106 }));
1107 let (_eager, deferred) = reg.split_eager_deferred();
1108 let (_spec, hint) = &deferred[0];
1110 assert!(
1111 hint.is_none(),
1112 "empty sentence before '.' → hint must be None"
1113 );
1114 }
1115
1116 #[test]
1119 fn is_deferred_spec_true_for_deferred_tool() {
1120 let reg = make_registry().register(Arc::new(DeferredTool {
1121 name: "DS",
1122 description: "d.",
1123 }));
1124 let specs = reg.specs();
1125 let ds_spec = specs.iter().find(|s| s.name == "DS").unwrap();
1126 assert!(
1127 reg.is_deferred_spec(ds_spec),
1128 "deferred tool spec must be deferred"
1129 );
1130 }
1131
1132 #[test]
1133 fn is_deferred_spec_false_for_eager_tool() {
1134 let reg = make_registry().register(Arc::new(ReadOnlyTool { name: "EagerSpec" }));
1135 let specs = reg.specs();
1136 let spec = &specs[0];
1137 assert!(
1138 !reg.is_deferred_spec(spec),
1139 "eager tool must not be deferred"
1140 );
1141 }
1142
1143 #[test]
1144 fn is_deferred_spec_false_for_unknown_spec() {
1145 let reg = make_registry();
1146 let fake_spec = ToolSpec {
1147 name: "Phantom".to_string(),
1148 description: "not registered".to_string(),
1149 parameters: serde_json::json!({}),
1150 };
1151 assert!(
1152 !reg.is_deferred_spec(&fake_spec),
1153 "unknown spec must return false"
1154 );
1155 }
1156
1157 #[test]
1160 fn retain_tools_removes_aliases_of_dropped_tools() {
1161 let mut reg = make_registry()
1162 .register_with_aliases(Arc::new(ReadOnlyTool { name: "Keep2" }), &["k_alias"])
1163 .register_with_aliases(Arc::new(MutatingTool { name: "Drop2" }), &["d_alias"]);
1164 reg.retain_tools(&["Keep2".to_string()]);
1165 assert!(reg.find_by_name("Keep2").is_some());
1166 assert!(
1167 reg.find_by_name("k_alias").is_some(),
1168 "kept tool's alias must survive"
1169 );
1170 assert!(reg.find_by_name("Drop2").is_none());
1171 assert!(
1172 reg.find_by_name("d_alias").is_none(),
1173 "dropped tool's alias must be removed"
1174 );
1175 }
1176
1177 #[test]
1180 fn with_touched_files_and_clear() {
1181 use std::sync::{Arc, Mutex};
1182 let slot: Arc<Mutex<TouchedFiles>> = Arc::new(Mutex::new(TouchedFiles::default()));
1183 let mut reg = make_registry().with_touched_files(slot);
1184 assert!(
1185 reg.touched_files().is_some(),
1186 "slot must be present after with_touched_files"
1187 );
1188 reg.clear_touched_files();
1189 assert!(
1190 reg.touched_files().is_none(),
1191 "slot must be gone after clear_touched_files"
1192 );
1193 }
1194
1195 #[test]
1198 fn with_read_file_state_roundtrip() {
1199 let slot: Arc<Mutex<ReadFileState>> = Arc::new(Mutex::new(ReadFileState::default()));
1200 let reg = make_registry().with_read_file_state(slot);
1201 assert!(
1202 reg.read_file_state().is_some(),
1203 "read_file_state must be Some after setter"
1204 );
1205 let reg2 = make_registry();
1206 assert!(
1207 reg2.read_file_state().is_none(),
1208 "fresh registry has no read_file_state"
1209 );
1210 }
1211
1212 #[test]
1215 fn permission_mode_default_on_fresh_registry() {
1216 let reg = make_registry();
1217 assert!(
1218 matches!(reg.permission_mode(), PermissionMode::Default),
1219 "fresh registry must have Default permission mode"
1220 );
1221 }
1222
1223 #[test]
1224 fn with_permissions_sets_permission_mode() {
1225 let config = PermissionsConfig {
1226 mode: PermissionMode::AcceptEdits,
1227 ..Default::default()
1228 };
1229 let reg = make_registry().with_permissions(config);
1230 assert!(
1231 matches!(reg.permission_mode(), PermissionMode::AcceptEdits),
1232 "permission_mode must reflect config.mode"
1233 );
1234 }
1235
1236 #[test]
1239 fn permissions_config_none_on_fresh_registry() {
1240 let reg = make_registry();
1241 assert!(
1242 reg.permissions_config().is_none(),
1243 "fresh registry has no permissions_config"
1244 );
1245 }
1246
1247 #[test]
1248 fn permissions_config_some_after_with_permissions() {
1249 let config = PermissionsConfig {
1250 mode: PermissionMode::Default,
1251 ..Default::default()
1252 };
1253 let reg = make_registry().with_permissions(config);
1254 assert!(
1255 reg.permissions_config().is_some(),
1256 "permissions_config must be Some"
1257 );
1258 }
1259
1260 #[test]
1263 fn is_plan_mode_false_for_default_tool_and_no_config() {
1264 let reg = make_registry().register(Arc::new(ReadOnlyTool { name: "PlanTool" }));
1265 assert!(
1267 !reg.is_plan_mode("PlanTool"),
1268 "no plan-mode config → must be false"
1269 );
1270 assert!(!reg.is_plan_mode("Unknown"), "unknown tool → must be false");
1271 }
1272
1273 struct DenyAllHook;
1276 #[async_trait]
1277 impl PermissionHook for DenyAllHook {
1278 async fn check(&self, _tool_name: &str, _args: &serde_json::Value) -> PermissionDecision {
1279 PermissionDecision::Deny("no".into())
1280 }
1281 }
1282
1283 #[test]
1284 fn with_permission_hook_installs_hook() {
1285 let reg = make_registry().with_permission_hook(Arc::new(DenyAllHook));
1286 assert!(reg.permission_hook.is_some(), "hook must be installed");
1287 }
1288
1289 #[test]
1290 fn clear_permission_hook_removes_hook() {
1291 let mut reg = make_registry().with_permission_hook(Arc::new(DenyAllHook));
1292 reg.clear_permission_hook();
1293 assert!(reg.permission_hook.is_none(), "hook must be cleared");
1294 }
1295
1296 #[test]
1299 fn with_policy_and_policy_roundtrip() {
1300 use super::policy_sandbox::PolicyConfig;
1301 let reg = make_registry();
1302 assert!(reg.policy().is_none(), "fresh registry has no policy");
1303 let cfg = PolicyConfig::default();
1304 let reg2 = reg.with_policy(cfg);
1305 assert!(
1306 reg2.policy().is_some(),
1307 "policy must be Some after with_policy"
1308 );
1309 }
1310
1311 #[test]
1314 fn with_headless_sets_flag() {
1315 let reg = make_registry();
1316 assert!(!reg.headless, "fresh registry must not be headless");
1317 let reg2 = reg.with_headless(true);
1318 assert!(
1319 reg2.headless,
1320 "headless must be true after with_headless(true)"
1321 );
1322 }
1323
1324 #[test]
1327 fn transport_accessor_returns_same_arc() {
1328 let reg = make_registry();
1329 let _transport = reg.transport();
1331 }
1332
1333 #[test]
1336 fn fork_is_independent_from_original() {
1337 let mut reg = make_registry().register(Arc::new(ReadOnlyTool { name: "ForkOrig" }));
1338 let mut forked = reg.fork();
1339 forked.register_mut(Arc::new(MutatingTool { name: "ForkedOnly" }));
1341 assert!(
1342 reg.find_by_name("ForkedOnly").is_none(),
1343 "forked-only tool must not appear in original"
1344 );
1345 reg.retain_tools(&["ForkOrig".to_string()]);
1347 assert!(
1348 forked.find_by_name("ForkOrig").is_some(),
1349 "original retain must not affect fork"
1350 );
1351 }
1352
1353 #[test]
1356 fn names_does_not_include_aliases() {
1357 let reg = make_registry()
1358 .register_with_aliases(Arc::new(ReadOnlyTool { name: "PrimeName" }), &["anAlias"]);
1359 let names = reg.names();
1360 assert!(
1361 names.contains(&"PrimeName".to_string()),
1362 "primary name must be in names()"
1363 );
1364 assert!(
1365 !names.contains(&"anAlias".to_string()),
1366 "alias must NOT be in names()"
1367 );
1368 }
1369
1370 struct DefaultKindTool;
1373
1374 #[async_trait]
1375 impl Tool for DefaultKindTool {
1376 fn spec(&self) -> ToolSpec {
1377 ToolSpec {
1378 name: "DefaultKind".into(),
1379 description: "defaults to Other".into(),
1380 parameters: serde_json::json!({"type":"object","properties":{}}),
1381 }
1382 }
1383 async fn execute(&self, _args: Value) -> crate::error::Result<String> {
1384 Ok("ok".into())
1385 }
1386 }
1387
1388 #[test]
1389 fn default_kind_is_other() {
1390 let tool = DefaultKindTool;
1391 assert_eq!(tool.kind(), ToolKind::Other);
1392 }
1393
1394 #[test]
1395 fn build_kind_map_works() {
1396 let reg = make_registry()
1397 .register(Arc::new(ReadOnlyTool { name: "ReadTool" }))
1398 .register(Arc::new(MutatingTool { name: "MutTool" }));
1399 let map = reg.build_kind_map();
1400 assert_eq!(map.len(), 2);
1401 assert_eq!(map.get("ReadTool"), Some(&ToolKind::Read));
1402 assert_eq!(map.get("MutTool"), Some(&ToolKind::Other));
1403 }
1404
1405 #[test]
1406 fn build_kind_map_empty_registry() {
1407 let map = make_registry().build_kind_map();
1408 assert!(map.is_empty());
1409 }
1410}