#![allow(clippy::all)]
use super::context::AgentContext;
use async_trait::async_trait;
use dashmap::DashMap;
use nexo_llm::ToolDef;
use serde_json::Value;
use std::sync::Arc;
#[derive(Debug, Clone, Default)]
pub struct ToolMeta {
pub deferred: bool,
pub search_hint: Option<String>,
}
impl ToolMeta {
pub fn deferred() -> Self {
Self {
deferred: true,
search_hint: None,
}
}
pub fn deferred_with_hint(hint: impl Into<String>) -> Self {
Self {
deferred: true,
search_hint: Some(hint.into()),
}
}
pub fn with_search_hint(mut self, hint: impl Into<String>) -> Self {
self.search_hint = Some(hint.into());
self
}
}
#[async_trait]
pub trait ToolHandler: Send + Sync {
async fn call(&self, ctx: &AgentContext, args: Value) -> anyhow::Result<Value>;
}
pub type HandlerEntry = (ToolDef, Arc<dyn ToolHandler>);
#[derive(Default, Clone)]
pub struct ToolRegistry {
handlers: Arc<DashMap<String, HandlerEntry>>,
meta: Arc<DashMap<String, ToolMeta>>,
}
impl ToolRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register(&self, def: ToolDef, handler: impl ToolHandler + 'static) {
let name = def.name.clone();
if self.handlers.contains_key(&name) {
tracing::warn!(
tool = %name,
"tool registered twice — previous handler overwritten (use `register_if_absent` to preserve)"
);
}
self.handlers.insert(name, (def, Arc::new(handler)));
}
pub fn register_arc(&self, def: ToolDef, handler: Arc<dyn ToolHandler>) {
let name = def.name.clone();
if self.handlers.contains_key(&name) {
tracing::warn!(
tool = %name,
"tool registered twice — previous handler overwritten (use `register_if_absent` to preserve)"
);
}
self.handlers.insert(name, (def, handler));
}
pub fn register_with_meta(
&self,
def: ToolDef,
handler: impl ToolHandler + 'static,
meta: ToolMeta,
) {
let name = def.name.clone();
self.register(def, handler);
self.meta.insert(name, meta);
}
pub fn set_meta(&self, tool_name: &str, meta: ToolMeta) {
if self.handlers.contains_key(tool_name) {
self.meta.insert(tool_name.to_string(), meta);
}
}
pub fn meta(&self, tool_name: &str) -> Option<ToolMeta> {
self.meta.get(tool_name).map(|e| e.value().clone())
}
pub fn deferred_tools(&self) -> Vec<(String, ToolMeta)> {
self.meta
.iter()
.filter(|e| e.value().deferred)
.map(|e| (e.key().clone(), e.value().clone()))
.collect()
}
pub fn register_if_absent(&self, def: ToolDef, handler: impl ToolHandler + 'static) -> bool {
use dashmap::mapref::entry::Entry;
match self.handlers.entry(def.name.clone()) {
Entry::Occupied(_) => false,
Entry::Vacant(slot) => {
slot.insert((def, Arc::new(handler)));
true
}
}
}
pub fn register_if_absent_arc(&self, def: ToolDef, handler: Arc<dyn ToolHandler>) -> bool {
use dashmap::mapref::entry::Entry;
match self.handlers.entry(def.name.clone()) {
Entry::Occupied(_) => false,
Entry::Vacant(slot) => {
slot.insert((def, handler));
true
}
}
}
pub fn contains(&self, name: &str) -> bool {
self.handlers.contains_key(name)
}
pub fn get(&self, name: &str) -> Option<(ToolDef, Arc<dyn ToolHandler>)> {
self.handlers.get(name).map(|e| e.value().clone())
}
pub fn to_tool_defs(&self) -> Vec<ToolDef> {
self.handlers.iter().map(|e| e.value().0.clone()).collect()
}
pub fn to_tool_defs_non_deferred(&self) -> Vec<ToolDef> {
self.handlers
.iter()
.filter(|e| !self.meta.get(e.key()).map(|m| m.deferred).unwrap_or(false))
.map(|e| e.value().0.clone())
.collect()
}
pub fn deferred_tools_summary(&self) -> Option<String> {
let deferred: Vec<(String, String)> = self
.meta
.iter()
.filter(|e| e.value().deferred)
.filter_map(|e| {
let name = e.key().clone();
let desc = self
.handlers
.get(&name)
.map(|h| h.0.description.clone())
.unwrap_or_default();
Some((name, desc))
})
.collect();
if deferred.is_empty() {
return None;
}
let mut parts: Vec<String> = Vec::with_capacity(deferred.len() + 2);
parts.push(
"<deferred-tools>\n\
The following tools are available. Their full schemas are omitted to save tokens.\n\
Use ToolSearch(select:<name>) to load one when needed.\n"
.to_string(),
);
for (name, desc) in &deferred {
let desc = if desc.len() > 120 {
format!("{}…", &desc[..119])
} else {
desc.clone()
};
parts.push(format!("- {name}: {desc}"));
}
parts.push("</deferred-tools>".to_string());
Some(parts.join("\n"))
}
pub fn names(&self) -> Vec<String> {
self.handlers.iter().map(|e| e.key().clone()).collect()
}
pub fn plan_mode_unclassified(&self) -> Vec<String> {
crate::plan_mode::unclassified_tools(self.names())
}
pub fn assert_plan_mode_classified(&self) {
crate::plan_mode::assert_registry_classified(self.names());
}
pub fn retain_matching(&self, patterns: &[String]) -> usize {
if patterns.is_empty() {
return 0;
}
let victims: Vec<String> = self
.handlers
.iter()
.filter(|e| !super::effective::allowlist_matches(patterns, e.key()))
.map(|e| e.key().clone())
.collect();
let n = victims.len();
for k in victims {
self.handlers.remove(&k);
self.meta.remove(&k);
}
n
}
pub fn filtered_clone(&self, allowed_tools: &[String]) -> ToolRegistry {
let clone = ToolRegistry {
handlers: Arc::new(DashMap::new()),
meta: Arc::new(DashMap::new()),
};
for entry in self.handlers.iter() {
clone
.handlers
.insert(entry.key().clone(), entry.value().clone());
}
for entry in self.meta.iter() {
clone
.meta
.insert(entry.key().clone(), entry.value().clone());
}
clone.retain_matching(allowed_tools);
clone
}
pub fn apply_dispatch_capability(
&self,
policy: &nexo_config::DispatchPolicy,
is_admin: bool,
) -> usize {
use nexo_dispatch_tools::tool_names::{
should_register, ToolGroup, ADMIN_TOOL_NAMES, READ_TOOL_NAMES, WRITE_TOOL_NAMES,
};
let mut removed = 0usize;
let drop_set: Vec<&'static str> = [
(ToolGroup::Read, READ_TOOL_NAMES),
(ToolGroup::Write, WRITE_TOOL_NAMES),
(ToolGroup::Admin, ADMIN_TOOL_NAMES),
]
.into_iter()
.filter(|(group, _)| !should_register(policy, *group, is_admin))
.flat_map(|(_, names)| names.iter().copied())
.collect();
for name in drop_set {
if self.handlers.remove(name).is_some() {
removed += 1;
}
}
removed
}
pub fn clear_by_prefix(&self, prefix: &str) -> usize {
let keys: Vec<String> = self
.handlers
.iter()
.filter(|e| e.key().starts_with(prefix))
.map(|e| e.key().clone())
.collect();
let n = keys.len();
for k in keys {
self.handlers.remove(&k);
}
n
}
}
#[cfg(test)]
mod tests {
use super::*;
struct Noop;
#[async_trait]
impl ToolHandler for Noop {
async fn call(&self, _ctx: &AgentContext, _args: Value) -> anyhow::Result<Value> {
Ok(Value::Null)
}
}
fn mk_def(name: &str) -> ToolDef {
ToolDef {
name: name.to_string(),
description: "".into(),
parameters: serde_json::json!({"type": "object" }),
}
}
#[test]
fn clear_by_prefix_removes_matching() {
let reg = ToolRegistry::new();
reg.register(mk_def("mcp_srv_a"), Noop);
reg.register(mk_def("mcp_srv_b"), Noop);
reg.register(mk_def("memory_recall"), Noop);
let n = reg.clear_by_prefix("mcp_srv_");
assert_eq!(n, 2);
assert!(reg.get("memory_recall").is_some());
assert!(reg.get("mcp_srv_a").is_none());
}
struct TagHandler(&'static str);
#[async_trait]
impl ToolHandler for TagHandler {
async fn call(&self, _ctx: &AgentContext, _args: Value) -> anyhow::Result<Value> {
Ok(Value::String(self.0.into()))
}
}
fn tagged_def(name: &str, desc: &str) -> ToolDef {
ToolDef {
name: name.into(),
description: desc.into(),
parameters: serde_json::json!({"type": "object" }),
}
}
#[test]
fn register_if_absent_preserves_original() {
let reg = ToolRegistry::new();
assert!(reg.register_if_absent(tagged_def("tool_a", "first"), TagHandler("first")));
assert!(!reg.register_if_absent(tagged_def("tool_a", "second"), TagHandler("second")));
let (def, _) = reg.get("tool_a").unwrap();
assert_eq!(def.description, "first");
}
#[test]
fn register_if_absent_accepts_empty_slot() {
let reg = ToolRegistry::new();
assert!(reg.register_if_absent(mk_def("fresh"), Noop));
assert!(reg.contains("fresh"));
}
#[test]
fn retain_matching_keeps_globs_and_exacts_drops_rest() {
let reg = ToolRegistry::new();
reg.register(mk_def("memory_recall"), Noop);
reg.register(mk_def("memory_save"), Noop);
reg.register(mk_def("ext_weather_now"), Noop);
reg.register(mk_def("ext_github_comment"), Noop);
reg.register(mk_def("delegate"), Noop);
assert_eq!(reg.to_tool_defs().len(), 5);
let removed = reg.retain_matching(&[
"memory_*".into(),
"ext_weather_now".into(),
"delegate".into(),
]);
assert_eq!(removed, 1);
assert!(reg.contains("memory_recall"));
assert!(reg.contains("memory_save"));
assert!(reg.contains("ext_weather_now"));
assert!(reg.contains("delegate"));
assert!(!reg.contains("ext_github_comment"));
}
#[test]
fn retain_matching_empty_patterns_is_noop() {
let reg = ToolRegistry::new();
reg.register(mk_def("tool_a"), Noop);
reg.register(mk_def("tool_b"), Noop);
assert_eq!(reg.retain_matching(&[]), 0);
assert_eq!(reg.to_tool_defs().len(), 2);
}
#[test]
fn apply_dispatch_capability_none_drops_every_dispatch_tool() {
let reg = ToolRegistry::new();
for n in nexo_dispatch_tools::READ_TOOL_NAMES {
reg.register(mk_def(n), Noop);
}
for n in nexo_dispatch_tools::WRITE_TOOL_NAMES {
reg.register(mk_def(n), Noop);
}
for n in nexo_dispatch_tools::ADMIN_TOOL_NAMES {
reg.register(mk_def(n), Noop);
}
reg.register(mk_def("memory_recall"), Noop);
let policy = nexo_config::DispatchPolicy {
mode: nexo_config::DispatchCapability::None,
..Default::default()
};
let removed = reg.apply_dispatch_capability(&policy, true);
assert!(removed > 0);
assert!(reg.contains("memory_recall"));
assert!(!reg.contains("project_status"));
assert!(!reg.contains("program_phase"));
assert!(!reg.contains("set_concurrency_cap"));
}
#[test]
fn apply_dispatch_capability_read_only_keeps_reads_drops_writes() {
let reg = ToolRegistry::new();
for n in nexo_dispatch_tools::READ_TOOL_NAMES {
reg.register(mk_def(n), Noop);
}
for n in nexo_dispatch_tools::WRITE_TOOL_NAMES {
reg.register(mk_def(n), Noop);
}
let policy = nexo_config::DispatchPolicy {
mode: nexo_config::DispatchCapability::ReadOnly,
..Default::default()
};
reg.apply_dispatch_capability(&policy, false);
assert!(reg.contains("project_status"));
assert!(!reg.contains("program_phase"));
}
#[test]
fn apply_dispatch_capability_full_admin_keeps_admin_tools() {
let reg = ToolRegistry::new();
for n in nexo_dispatch_tools::ADMIN_TOOL_NAMES {
reg.register(mk_def(n), Noop);
}
let policy = nexo_config::DispatchPolicy {
mode: nexo_config::DispatchCapability::Full,
..Default::default()
};
reg.apply_dispatch_capability(&policy, true);
assert!(reg.contains("set_concurrency_cap"));
reg.apply_dispatch_capability(&policy, false);
assert!(!reg.contains("set_concurrency_cap"));
}
#[test]
fn contains_reflects_state() {
let reg = ToolRegistry::new();
assert!(!reg.contains("missing"));
reg.register(mk_def("present"), Noop);
assert!(reg.contains("present"));
assert!(!reg.contains("still-missing"));
}
#[test]
fn plan_mode_unclassified_reports_unknown_names() {
let reg = ToolRegistry::new();
reg.register(mk_def("FileRead"), Noop); reg.register(mk_def("totally_unknown_tool"), Noop); let bad = reg.plan_mode_unclassified();
assert_eq!(bad, vec!["totally_unknown_tool".to_string()]);
}
#[test]
fn plan_mode_unclassified_empty_when_all_known() {
let reg = ToolRegistry::new();
reg.register(mk_def("FileRead"), Noop);
reg.register(mk_def("FileEdit"), Noop);
reg.register(mk_def("EnterPlanMode"), Noop);
reg.register(mk_def("whatsapp.send"), Noop); assert!(reg.plan_mode_unclassified().is_empty());
}
#[test]
#[should_panic(expected = "plan_mode:")]
fn assert_plan_mode_classified_panics_on_unknown() {
let reg = ToolRegistry::new();
reg.register(mk_def("totally_unknown_tool"), Noop);
reg.assert_plan_mode_classified();
}
#[test]
fn to_tool_defs_non_deferred_skips_deferred() {
let reg = ToolRegistry::new();
reg.register(mk_def("normal_a"), Noop);
reg.register(mk_def("normal_b"), Noop);
reg.register_with_meta(mk_def("deferred_x"), Noop, ToolMeta::deferred());
let defs = reg.to_tool_defs_non_deferred();
let mut names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect();
names.sort_unstable();
assert_eq!(names, vec!["normal_a", "normal_b"]);
}
#[test]
fn to_tool_defs_non_deferred_empty_when_all_deferred() {
let reg = ToolRegistry::new();
reg.register_with_meta(mk_def("a"), Noop, ToolMeta::deferred());
reg.register_with_meta(mk_def("b"), Noop, ToolMeta::deferred());
assert!(reg.to_tool_defs_non_deferred().is_empty());
}
#[test]
fn deferred_tools_summary_returns_none_when_empty() {
let reg = ToolRegistry::new();
reg.register(mk_def("normal"), Noop);
assert!(reg.deferred_tools_summary().is_none());
}
#[test]
fn deferred_tools_summary_includes_deferred_names_and_descriptions() {
let reg = ToolRegistry::new();
reg.register_with_meta(
tagged_def("mcp_x", "Create issues on GitHub"),
Noop,
ToolMeta::deferred(),
);
reg.register_with_meta(
tagged_def("mcp_y", "Search Slack messages"),
Noop,
ToolMeta::deferred(),
);
reg.register(tagged_def("normal", "Just a normal tool"), Noop);
let summary = reg.deferred_tools_summary().unwrap();
assert!(summary.contains("<deferred-tools>"));
assert!(summary.contains("</deferred-tools>"));
assert!(summary.contains("mcp_x"));
assert!(summary.contains("Create issues on GitHub"));
assert!(summary.contains("mcp_y"));
assert!(summary.contains("Search Slack messages"));
assert!(!summary.contains("normal"));
}
#[test]
fn mark_built_in_deferred_excludes_listed_tools() {
use super::super::built_in_deferred::mark_built_in_deferred;
let reg = ToolRegistry::new();
reg.register(mk_def("TodoWrite"), Noop);
reg.register(mk_def("Lsp"), Noop);
reg.register(mk_def("Repl"), Noop);
reg.register(mk_def("FileRead"), Noop);
mark_built_in_deferred(®);
let defs = reg.to_tool_defs_non_deferred();
let mut names: Vec<&str> = defs.iter().map(|d| d.name.as_str()).collect();
names.sort_unstable();
assert_eq!(names, vec!["FileRead"]);
let deferred: Vec<String> = reg.deferred_tools().into_iter().map(|(n, _)| n).collect();
for expected in &["TodoWrite", "Lsp", "Repl"] {
assert!(
deferred.iter().any(|n| n == expected),
"expected {expected} to appear in deferred_tools(): got {deferred:?}",
);
}
}
#[test]
fn mark_built_in_deferred_skips_absent_tools() {
use super::super::built_in_deferred::mark_built_in_deferred;
let reg = ToolRegistry::new();
mark_built_in_deferred(®);
assert!(reg.to_tool_defs_non_deferred().is_empty());
}
#[test]
fn mark_built_in_deferred_propagates_search_hints() {
use super::super::built_in_deferred::mark_built_in_deferred;
let reg = ToolRegistry::new();
reg.register(mk_def("TodoWrite"), Noop);
mark_built_in_deferred(®);
let meta = reg.meta("TodoWrite").expect("TodoWrite meta present");
assert!(meta.deferred);
assert_eq!(
meta.search_hint.as_deref(),
Some("todo, tasks, in-progress checklist"),
);
}
}