zeph_tools/
registry.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::borrow::Cow;
5use std::fmt::Write;
6
7#[non_exhaustive]
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum InvocationHint {
10    /// Tool invoked via ```{tag}\n...\n``` fenced block in LLM response
11    FencedBlock(&'static str),
12    /// Tool invoked via structured `ToolCall` JSON
13    ToolCall,
14}
15
16#[derive(Debug, Clone)]
17pub struct ToolDef {
18    pub id: Cow<'static, str>,
19    pub description: Cow<'static, str>,
20    pub schema: schemars::Schema,
21    pub invocation: InvocationHint,
22    /// Raw output schema from an MCP server, if present.
23    ///
24    /// DO NOT convert to `schemars::Schema` — lossy; see #2931 critique P0-1.
25    pub output_schema: Option<serde_json::Value>,
26}
27
28#[derive(Debug, Default)]
29pub struct ToolRegistry {
30    tools: Vec<ToolDef>,
31}
32
33impl ToolRegistry {
34    #[must_use]
35    pub fn from_definitions(tools: Vec<ToolDef>) -> Self {
36        Self { tools }
37    }
38
39    #[must_use]
40    pub fn tools(&self) -> &[ToolDef] {
41        &self.tools
42    }
43
44    #[must_use]
45    pub fn find(&self, id: &str) -> Option<&ToolDef> {
46        self.tools.iter().find(|t| t.id.as_ref() == id)
47    }
48
49    /// Format tools for prompt, excluding tools fully denied by policy.
50    #[must_use]
51    pub fn format_for_prompt_filtered(
52        &self,
53        policy: &crate::permissions::PermissionPolicy,
54    ) -> String {
55        let mut out = String::from("<tools>\n");
56        for tool in &self.tools {
57            if policy.is_fully_denied(&tool.id) {
58                continue;
59            }
60            format_tool(&mut out, tool);
61        }
62        out.push_str("</tools>");
63        out
64    }
65}
66
67fn format_tool(out: &mut String, tool: &ToolDef) {
68    let _ = writeln!(out, "## {}", tool.id);
69    let _ = writeln!(out, "{}", tool.description);
70    match tool.invocation {
71        InvocationHint::FencedBlock(tag) => {
72            let _ = writeln!(out, "Invocation: use ```{tag} fenced block");
73        }
74        InvocationHint::ToolCall => {
75            let _ = writeln!(
76                out,
77                "Invocation: use tool_call with {{\"tool_id\": \"{}\", \"params\": {{...}}}}",
78                tool.id
79            );
80        }
81    }
82    format_schema_params(out, &tool.schema);
83    out.push('\n');
84}
85
86/// Extract the primary type when schemars renders `Option<T>` as `"type": ["T", "null"]`
87/// or `"anyOf": [{"type": "T"}, {"type": "null"}]`.
88fn extract_non_null_type(obj: &serde_json::Map<String, serde_json::Value>) -> Option<&str> {
89    if let Some(arr) = obj.get("type").and_then(|v| v.as_array()) {
90        return arr.iter().filter_map(|v| v.as_str()).find(|t| *t != "null");
91    }
92    obj.get("anyOf")?
93        .as_array()?
94        .iter()
95        .filter_map(|v| v.as_object())
96        .filter_map(|o| o.get("type")?.as_str())
97        .find(|t| *t != "null")
98}
99
100fn format_schema_params(out: &mut String, schema: &schemars::Schema) {
101    let Some(obj) = schema.as_object() else {
102        return;
103    };
104    let Some(serde_json::Value::Object(props)) = obj.get("properties") else {
105        return;
106    };
107    if props.is_empty() {
108        return;
109    }
110
111    let required: Vec<&str> = obj
112        .get("required")
113        .and_then(|v| v.as_array())
114        .map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
115        .unwrap_or_default();
116
117    let _ = writeln!(out, "Parameters:");
118    for (name, prop) in props {
119        let prop_obj = prop.as_object();
120        let ty = prop_obj
121            .and_then(|o| {
122                o.get("type")
123                    .and_then(|v| v.as_str())
124                    .or_else(|| extract_non_null_type(o))
125            })
126            .unwrap_or("string");
127        let desc = prop_obj
128            .and_then(|o| o.get("description"))
129            .and_then(|v| v.as_str())
130            .unwrap_or("");
131        let req = if required.contains(&name.as_str()) {
132            "required"
133        } else {
134            "optional"
135        };
136        let _ = writeln!(out, "  - {name}: {desc} ({ty}, {req})");
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use crate::file::ReadParams;
144    use crate::shell::BashParams;
145
146    fn sample_tools() -> Vec<ToolDef> {
147        vec![
148            ToolDef {
149                id: "bash".into(),
150                description: "Execute a shell command".into(),
151                schema: schemars::schema_for!(BashParams),
152                invocation: InvocationHint::FencedBlock("bash"),
153                output_schema: None,
154            },
155            ToolDef {
156                id: "read".into(),
157                description: "Read file contents".into(),
158                schema: schemars::schema_for!(ReadParams),
159                invocation: InvocationHint::ToolCall,
160                output_schema: None,
161            },
162        ]
163    }
164
165    #[test]
166    fn from_definitions_stores_tools() {
167        let reg = ToolRegistry::from_definitions(sample_tools());
168        assert_eq!(reg.tools().len(), 2);
169    }
170
171    #[test]
172    fn default_registry_is_empty() {
173        let reg = ToolRegistry::default();
174        assert!(reg.tools().is_empty());
175    }
176
177    #[test]
178    fn find_existing_tool() {
179        let reg = ToolRegistry::from_definitions(sample_tools());
180        assert!(reg.find("bash").is_some());
181        assert!(reg.find("read").is_some());
182    }
183
184    #[test]
185    fn find_nonexistent_returns_none() {
186        let reg = ToolRegistry::from_definitions(sample_tools());
187        assert!(reg.find("nonexistent").is_none());
188    }
189
190    #[test]
191    fn format_for_prompt_contains_tools() {
192        let reg = ToolRegistry::from_definitions(sample_tools());
193        let prompt =
194            reg.format_for_prompt_filtered(&crate::permissions::PermissionPolicy::default());
195        assert!(prompt.contains("<tools>"));
196        assert!(prompt.contains("</tools>"));
197        assert!(prompt.contains("## bash"));
198        assert!(prompt.contains("## read"));
199    }
200
201    #[test]
202    fn format_for_prompt_shows_invocation_fenced() {
203        let reg = ToolRegistry::from_definitions(sample_tools());
204        let prompt =
205            reg.format_for_prompt_filtered(&crate::permissions::PermissionPolicy::default());
206        assert!(prompt.contains("Invocation: use ```bash fenced block"));
207    }
208
209    #[test]
210    fn format_for_prompt_shows_invocation_tool_call() {
211        let reg = ToolRegistry::from_definitions(sample_tools());
212        let prompt =
213            reg.format_for_prompt_filtered(&crate::permissions::PermissionPolicy::default());
214        assert!(prompt.contains("Invocation: use tool_call"));
215        assert!(prompt.contains("\"tool_id\": \"read\""));
216    }
217
218    #[test]
219    fn format_for_prompt_shows_param_info() {
220        let reg = ToolRegistry::from_definitions(sample_tools());
221        let prompt =
222            reg.format_for_prompt_filtered(&crate::permissions::PermissionPolicy::default());
223        assert!(prompt.contains("command:"));
224        assert!(prompt.contains("required"));
225        assert!(prompt.contains("string"));
226    }
227
228    #[test]
229    fn format_for_prompt_shows_optional_params() {
230        let reg = ToolRegistry::from_definitions(sample_tools());
231        let prompt =
232            reg.format_for_prompt_filtered(&crate::permissions::PermissionPolicy::default());
233        assert!(prompt.contains("offset:"));
234        assert!(prompt.contains("optional"));
235        assert!(
236            prompt.contains("(integer, optional)"),
237            "Option<u32> should render as integer, not string: {prompt}"
238        );
239    }
240
241    #[test]
242    fn format_filtered_excludes_fully_denied() {
243        use crate::permissions::{PermissionAction, PermissionPolicy, PermissionRule};
244        use std::collections::HashMap;
245        let mut rules = HashMap::new();
246        rules.insert(
247            "bash".to_owned(),
248            vec![PermissionRule {
249                pattern: "*".to_owned(),
250                action: PermissionAction::Deny,
251            }],
252        );
253        let policy = PermissionPolicy::new(rules);
254        let reg = ToolRegistry::from_definitions(sample_tools());
255        let prompt = reg.format_for_prompt_filtered(&policy);
256        assert!(!prompt.contains("## bash"));
257        assert!(prompt.contains("## read"));
258    }
259
260    #[test]
261    fn format_filtered_includes_mixed_rules() {
262        use crate::permissions::{PermissionAction, PermissionPolicy, PermissionRule};
263        use std::collections::HashMap;
264        let mut rules = HashMap::new();
265        rules.insert(
266            "bash".to_owned(),
267            vec![
268                PermissionRule {
269                    pattern: "echo *".to_owned(),
270                    action: PermissionAction::Allow,
271                },
272                PermissionRule {
273                    pattern: "*".to_owned(),
274                    action: PermissionAction::Deny,
275                },
276            ],
277        );
278        let policy = PermissionPolicy::new(rules);
279        let reg = ToolRegistry::from_definitions(sample_tools());
280        let prompt = reg.format_for_prompt_filtered(&policy);
281        assert!(prompt.contains("## bash"));
282    }
283
284    #[test]
285    fn format_filtered_no_rules_includes_all() {
286        let policy = crate::permissions::PermissionPolicy::default();
287        let reg = ToolRegistry::from_definitions(sample_tools());
288        let prompt = reg.format_for_prompt_filtered(&policy);
289        assert!(prompt.contains("## bash"));
290        assert!(prompt.contains("## read"));
291    }
292}
zeph_tools/registry.rs

zeph_tools/
registry.rs
