defect_agent/session/
tool_registry.rs

1//! [`ToolRegistry`] implementations: static registry + composite query.
2//!
3//! - [`StaticToolRegistry`]: process-level (builtin tools) or session-level (MCP tools),
4//!   immutable after construction
5//! - [`CompositeRegistry`]: chains two registries so the main loop sees a unified
6//!   interface (`get` checks session-level first, then process-level; schemas concatenate
7//!   both)
8
9use std::collections::{HashMap, HashSet};
10use std::sync::Arc;
11
12use crate::session::ToolRegistry;
13use crate::tool::{Tool, ToolSchema};
14
15/// An immutable mapping from names to tools.
16///
17/// Construct via [`StaticToolRegistry::builder`]; once built, no tools can be added or
18/// removed, ensuring that the schema order and `get` behavior remain stable.
19pub struct StaticToolRegistry {
20    schemas: Vec<ToolSchema>,
21    by_name: HashMap<String, Arc<dyn Tool>>,
22}
23
24impl StaticToolRegistry {
25    pub fn builder() -> StaticToolRegistryBuilder {
26        StaticToolRegistryBuilder::default()
27    }
28
29    /// An empty registry, useful for testing or as a placeholder.
30    pub fn empty() -> Self {
31        Self {
32            schemas: Vec::new(),
33            by_name: HashMap::new(),
34        }
35    }
36}
37
38impl ToolRegistry for StaticToolRegistry {
39    fn schemas(&self) -> Vec<ToolSchema> {
40        self.schemas.clone()
41    }
42
43    fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
44        self.by_name.get(name).cloned()
45    }
46}
47
48/// A builder for [`StaticToolRegistry`].
49#[derive(Default)]
50pub struct StaticToolRegistryBuilder {
51    schemas: Vec<ToolSchema>,
52    by_name: HashMap<String, Arc<dyn Tool>>,
53}
54
55impl StaticToolRegistryBuilder {
56    /// Registers a tool. If a tool with the same name already exists, it is overwritten
57    /// and the old entry in `schemas` is replaced with the new one (keeping `schemas`
58    /// order stable for diagnostics).
59    pub fn insert(mut self, tool: Arc<dyn Tool>) -> Self {
60        let schema = tool.schema().clone();
61        if let Some(pos) = self.schemas.iter().position(|s| s.name == schema.name) {
62            if let Some(slot) = self.schemas.get_mut(pos) {
63                *slot = schema.clone();
64            }
65        } else {
66            self.schemas.push(schema.clone());
67        }
68        self.by_name.insert(schema.name, tool);
69        self
70    }
71
72    pub fn build(self) -> StaticToolRegistry {
73        StaticToolRegistry {
74            schemas: self.schemas,
75            by_name: self.by_name,
76        }
77    }
78}
79
80/// A "composite" view of the process-level and session-level registries.
81///
82/// Lookup semantics: session-level (per-session MCP) first, then process-level
83/// (built-in). This allows session-level MCP tools to "shadow" built-in tools with the
84/// same name — a common convention in the MCP model.
85pub struct CompositeRegistry {
86    session: Arc<dyn ToolRegistry>,
87    process: Arc<dyn ToolRegistry>,
88}
89
90impl CompositeRegistry {
91    pub fn new(session: Arc<dyn ToolRegistry>, process: Arc<dyn ToolRegistry>) -> Self {
92        Self { session, process }
93    }
94}
95
96impl ToolRegistry for CompositeRegistry {
97    fn schemas(&self) -> Vec<ToolSchema> {
98        let mut session_schemas = self.session.schemas();
99        let mut process_schemas = self.process.schemas();
100        // Session-level override: remove schemas from `process` that are already declared
101        // in `session` with the same name.
102        let session_names: HashSet<&str> =
103            session_schemas.iter().map(|s| s.name.as_str()).collect();
104        process_schemas.retain(|s| !session_names.contains(s.name.as_str()));
105        session_schemas.append(&mut process_schemas);
106        session_schemas
107    }
108
109    fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
110        self.session.get(name).or_else(|| self.process.get(name))
111    }
112}
113
114#[cfg(test)]
115mod tests;
defect_agent/session/tool_registry.rs

defect_agent/session/
tool_registry.rs
