funera_core/re_act/
tool.rs

1#![cfg(feature = "tool")]
2
3use std::{collections::HashMap, fmt::Display, sync::Arc};
4
5use anyhow::Result;
6use async_trait::async_trait;
7use serde::{Deserialize, Serialize};
8use serde_json::Value as JsonValue;
9use thiserror::Error;
10
11/// The type of a tool, as communicated to the LLM.
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
13pub enum ToolType {
14    /// Standard OpenAI-compatible function tool.
15    Function,
16}
17
18impl Display for ToolType {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        write!(f, "function")
21    }
22}
23
24/// A callable tool exposed to the LLM agent.
25///
26/// Implement this trait to define custom tools. The framework will expose
27/// the tool's [`schema`](Tool::schema) to the LLM and invoke
28/// [`execute`](Tool::execute) when the LLM requests it.
29#[async_trait]
30pub trait Tool: Send + Sync {
31    /// Unique name for this tool (e.g. `"read"`, `"shell"`).
32    fn name(&self) -> &str;
33
34    /// Human-readable description sent to the LLM.
35    fn description(&self) -> &str;
36
37    /// Execute the tool with the given JSON arguments.
38    ///
39    /// Returns a string result on success, or a [`ToolCallError`] on failure.
40    async fn execute(&self, args: JsonValue) -> Result<String, ToolCallError>;
41
42    /// Returns the tool type (defaults to [`ToolType::Function`]).
43    fn get_type(&self) -> ToolType {
44        ToolType::Function
45    }
46
47    /// Returns the JSON schema describing this tool's parameters.
48    ///
49    /// This is sent to the LLM so it can generate well-formed invocations.
50    fn schema(&self) -> JsonValue;
51}
52
53/// Errors that can occur during tool execution.
54#[derive(Debug, Error)]
55pub enum ToolCallError {
56    /// The arguments did not match the expected schema.
57    #[error("parameter mismatch: {0}")]
58    ParameterMismatch(JsonValue),
59
60    /// The tool encountered a runtime error during execution.
61    #[error("tool execution error: {0}")]
62    ToolExecutionError(#[from] anyhow::Error),
63
64    /// The tool exists but is currently unavailable (e.g. disabled by policy).
65    #[error("tool unavailable: {0}")]
66    ToolUnavailable(String),
67
68    /// No tool with the given name is registered.
69    #[error("tool not found: {0}")]
70    ToolNotFound(String),
71
72    /// The tool call requires user approval before proceeding.
73    #[error("approval required for {tool_name}: {reason}")]
74    ApprovalRequired {
75        call_id: String,
76        tool_name: String,
77        reason: String,
78    },
79
80    /// The tool call was rejected after an approval request was denied.
81    #[error("tool call rejected: {reason}")]
82    Rejected { reason: String },
83}
84
85/// An entry in the tool registry, pairing a tool with its availability status.
86#[derive(Clone)]
87pub struct ToolRegistryEntry {
88    pub tool: Arc<dyn Tool>,
89    pub available: bool,
90}
91impl ToolRegistryEntry {
92    /// Create a new registry entry with explicit availability.
93    pub fn new(tool: Arc<dyn Tool>, available: bool) -> Self {
94        Self { tool, available }
95    }
96
97    /// Whether the tool is currently available for execution.
98    pub fn is_available(&self) -> bool {
99        self.available
100    }
101
102    /// Create a new registry entry with the tool available.
103    pub fn new_available(tool: Arc<dyn Tool>) -> Self {
104        Self::new(tool, true)
105    }
106
107    /// Create a new registry entry with the tool unavailable.
108    pub fn new_unavailable(tool: Arc<dyn Tool>) -> Self {
109        Self::new(tool, false)
110    }
111}
112
113/// Raw tool registry (no security checks).
114///
115/// When the `security` feature is enabled, [`ToolRegistry`] aliases to
116/// [`GuardedToolRegistry`](crate::security::registry::GuardedToolRegistry)
117/// instead, which wraps this registry with policy checks and audit logging.
118#[doc(hidden)]
119#[derive(Clone)]
120pub struct RawToolRegistry {
121    tools: HashMap<String, ToolRegistryEntry>,
122}
123
124impl Default for RawToolRegistry {
125    fn default() -> Self {
126        Self::new()
127    }
128}
129
130impl RawToolRegistry {
131    pub fn new() -> Self {
132        Self {
133            tools: HashMap::new(),
134        }
135    }
136
137    pub fn add_tool(&mut self, tool: Arc<dyn Tool>) {
138        self.tools.insert(
139            tool.name().to_string(),
140            ToolRegistryEntry::new_available(tool),
141        );
142    }
143    pub fn get_tool(&self, name: &str) -> Option<&ToolRegistryEntry> {
144        self.tools.get(name)
145    }
146
147    /// Clone the tool's `Arc` if it exists and is available.
148    ///
149    /// Used by the executor to run a tool outside the registry lock.
150    pub fn get_tool_arc(&self, name: &str) -> Option<Arc<dyn Tool>> {
151        self.get_tool(name)
152            .filter(|entry| entry.is_available())
153            .map(|entry| entry.tool.clone())
154    }
155    pub fn remove_tool(&mut self, name: &str) {
156        self.tools.remove(name);
157    }
158
159    /// Remove the tool only if the registered entry is the same `Arc` value.
160    /// This prevents a stale disposer from deleting a replacement tool that
161    /// reuses the same name (e.g. HMR replacement).
162    pub fn remove_tool_if_same(&mut self, name: &str, tool: &Arc<dyn Tool>) -> bool {
163        match self.tools.get(name) {
164            Some(entry) if Arc::ptr_eq(&entry.tool, tool) => {
165                self.tools.remove(name);
166                true
167            }
168            _ => false,
169        }
170    }
171    pub fn tool_exists(&self, name: &str) -> bool {
172        self.tools.contains_key(name)
173    }
174    pub fn tool_count(&self) -> usize {
175        self.tools.len()
176    }
177    pub fn get_all_tools(&self) -> &HashMap<String, ToolRegistryEntry> {
178        &self.tools
179    }
180    pub fn available_tools_json(&self) -> JsonValue {
181        self.tools
182            .values()
183            .filter_map(|tool| {
184                if tool.is_available() {
185                    Some(tool.tool.schema())
186                } else {
187                    None
188                }
189            })
190            .collect::<Vec<_>>()
191            .into()
192    }
193    pub async fn call_tool(&self, name: &str, args: JsonValue) -> Result<String, ToolCallError> {
194        if let Some(tool) = self.get_tool(name) {
195            if tool.is_available() {
196                tool.tool.execute(args).await
197            } else {
198                Err(ToolCallError::ToolUnavailable(name.to_string()))
199            }
200        } else {
201            Err(ToolCallError::ToolNotFound(name.to_string()))
202        }
203    }
204}
205
206/// The active tool registry type.
207///
208/// When the `security` feature is enabled, this aliases to
209/// [`GuardedToolRegistry`](crate::security::registry::GuardedToolRegistry)
210/// which enforces tool policies and logs audit events on every tool call.
211/// Without `security`, it is the raw registry with no policy checks.
212#[cfg(feature = "security")]
213pub use crate::security::registry::GuardedToolRegistry as ToolRegistry;
214
215#[cfg(not(feature = "security"))]
216pub use RawToolRegistry as ToolRegistry;
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221    use serde_json::json;
222
223    #[test]
224    fn rejected_error_display() {
225        let e = ToolCallError::Rejected {
226            reason: "access denied".into(),
227        };
228        let msg = format!("{e}");
229        assert!(msg.contains("access denied"), "msg: {msg}");
230    }
231
232    #[test]
233    fn approval_required_error_display() {
234        let e = ToolCallError::ApprovalRequired {
235            call_id: "c1".into(),
236            tool_name: "shell".into(),
237            reason: "needs approval".into(),
238        };
239        let msg = format!("{e}");
240        assert!(msg.contains("shell"), "msg: {msg}");
241        assert!(msg.contains("approval"), "msg: {msg}");
242    }
243
244    struct MockTool;
245    #[async_trait]
246    impl Tool for MockTool {
247        fn name(&self) -> &str {
248            "mock"
249        }
250        fn description(&self) -> &str {
251            "mock tool"
252        }
253        fn schema(&self) -> JsonValue {
254            json!({"type": "function", "function": {"name": "mock"}})
255        }
256        async fn execute(&self, _args: JsonValue) -> Result<String, ToolCallError> {
257            Ok("done".into())
258        }
259    }
260
261    #[test]
262    fn get_tool_arc_returns_available_tool_or_none() {
263        let mut reg = RawToolRegistry::new();
264        assert!(reg.get_tool_arc("mock").is_none());
265
266        reg.add_tool(Arc::new(MockTool));
267        let tool = reg.get_tool_arc("mock");
268        assert!(
269            tool.is_some(),
270            "registered tool must be clonable via get_tool_arc"
271        );
272        assert_eq!(tool.unwrap().name(), "mock");
273
274        assert!(reg.get_tool_arc("missing").is_none());
275    }
276
277    #[test]
278    fn remove_tool_if_same_only_removes_matching_arc() {
279        let mut reg = RawToolRegistry::new();
280        let original: Arc<dyn Tool> = Arc::new(MockTool);
281        reg.add_tool(Arc::clone(&original));
282
283        // A different Arc (e.g. a replacement with the same name) is kept.
284        let other: Arc<dyn Tool> = Arc::new(MockTool);
285        assert!(!reg.remove_tool_if_same("mock", &other));
286        assert!(reg.tool_exists("mock"));
287
288        // Removing the registered Arc succeeds.
289        assert!(reg.remove_tool_if_same("mock", &original));
290        assert!(!reg.tool_exists("mock"));
291    }
292}
funera_core/re_act/tool.rs

funera_core/re_act/
tool.rs
