aegis_tools/
registry.rs

1use anyhow::Result;
2use async_trait::async_trait;
3use serde_json::Value;
4use std::collections::HashMap;
5use std::path::PathBuf;
6use std::sync::Arc;
7
8/// Context passed to every tool execution.
9pub struct ToolContext<'a> {
10    pub cwd: PathBuf,
11    pub session_id: String,
12    /// Callback: returns true if user approves the command.
13    pub approve_fn: &'a (dyn Fn(&str) -> bool + Send + Sync),
14    /// Whether to skip all approval (YOLO mode).
15    pub yolo: bool,
16    /// Who is invoking this tool call. `None` means "assume `LocalOwner`"
17    /// — the default for legacy callers that predate the identity system.
18    /// New callers should always populate this.
19    pub identity: Option<aegis_security::Identity>,
20    /// Whether the sandbox layer is enabled at runtime (from
21    /// `[sandbox] enabled` in config.toml). When `false`, tools skip the
22    /// pre_exec hook regardless of identity — this matches the "opt-in"
23    /// design promise so upgrading users see no behavior change.
24    pub sandbox_enabled: bool,
25}
26
27impl ToolContext<'_> {
28    /// Check whether the given command is approved, either by YOLO mode or the approval callback.
29    pub fn approve(&self, command: &str) -> bool {
30        self.yolo || (self.approve_fn)(command)
31    }
32
33    /// Return the effective identity for this tool call, defaulting to
34    /// [`aegis_security::Identity::LocalOwner`] when unset.
35    pub fn effective_identity(&self) -> aegis_security::Identity {
36        self.identity
37            .clone()
38            .unwrap_or(aegis_security::Identity::LocalOwner)
39    }
40}
41
42/// A tool that the agent can invoke.
43#[async_trait]
44pub trait Tool: Send + Sync {
45    fn name(&self) -> &str;
46    fn description(&self) -> &str;
47    fn parameters(&self) -> Value;
48    async fn execute(&self, args: Value, ctx: &ToolContext<'_>) -> Result<String>;
49}
50
51/// Registry of available tools.
52pub struct ToolRegistry {
53    tools: HashMap<String, Arc<dyn Tool>>,
54}
55
56impl ToolRegistry {
57    /// Create an empty tool registry.
58    pub fn new() -> Self {
59        Self {
60            tools: HashMap::new(),
61        }
62    }
63
64    /// Register a tool, keyed by its name. Overwrites any existing tool with the same name.
65    pub fn register(&mut self, tool: Arc<dyn Tool>) {
66        self.tools.insert(tool.name().to_string(), tool);
67    }
68
69    /// Look up a registered tool by name.
70    pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
71        self.tools.get(name)
72    }
73
74    /// Return the names of all registered tools.
75    pub fn names(&self) -> Vec<String> {
76        self.tools.keys().cloned().collect()
77    }
78
79    /// Serialize all registered tools into an OpenAI-compatible function-calling schema.
80    pub fn to_openai_schema(&self) -> Value {
81        let arr: Vec<Value> = self
82            .tools
83            .values()
84            .map(|t| {
85                serde_json::json!({
86                    "type": "function",
87                    "function": {
88                        "name": t.name(),
89                        "description": t.description(),
90                        "parameters": t.parameters(),
91                    }
92                })
93            })
94            .collect();
95        Value::Array(arr)
96    }
97
98    /// Return a human-readable, sorted list of tool names and their descriptions.
99    pub fn tool_descriptions(&self) -> String {
100        let mut names: Vec<_> = self.tools.keys().collect();
101        names.sort();
102        names
103            .iter()
104            .map(|n| {
105                let t = &self.tools[n.as_str()];
106                format!("- {}: {}", t.name(), t.description())
107            })
108            .collect::<Vec<_>>()
109            .join("\n")
110    }
111}
112
113impl Default for ToolRegistry {
114    fn default() -> Self {
115        Self::new()
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    struct DummyTool;
124
125    #[async_trait]
126    impl Tool for DummyTool {
127        fn name(&self) -> &str { "dummy" }
128        fn description(&self) -> &str { "A test tool" }
129        fn parameters(&self) -> Value { serde_json::json!({"type": "object", "properties": {}}) }
130        async fn execute(&self, _args: Value, _ctx: &ToolContext<'_>) -> Result<String> {
131            Ok("ok".into())
132        }
133    }
134
135    #[test]
136    fn test_register_and_get() {
137        let mut reg = ToolRegistry::new();
138        reg.register(Arc::new(DummyTool));
139        assert!(reg.get("dummy").is_some());
140        assert!(reg.get("nonexistent").is_none());
141    }
142
143    #[test]
144    fn test_names() {
145        let mut reg = ToolRegistry::new();
146        reg.register(Arc::new(DummyTool));
147        let names = reg.names();
148        assert!(names.contains(&"dummy".to_string()));
149    }
150
151    #[test]
152    fn test_openai_schema() {
153        let mut reg = ToolRegistry::new();
154        reg.register(Arc::new(DummyTool));
155        let schema = reg.to_openai_schema();
156        let arr = schema.as_array().unwrap();
157        assert_eq!(arr.len(), 1);
158        assert_eq!(arr[0]["type"], "function");
159        assert_eq!(arr[0]["function"]["name"], "dummy");
160    }
161
162    #[test]
163    fn test_tool_descriptions() {
164        let mut reg = ToolRegistry::new();
165        reg.register(Arc::new(DummyTool));
166        let desc = reg.tool_descriptions();
167        assert!(desc.contains("dummy: A test tool"));
168    }
169}
aegis_tools/registry.rs

aegis_tools/
registry.rs
