agent_engine/tools/
extension.rs

1use std::sync::Arc;
2
3use serde_json::Value;
4
5use crate::extensions::runtime::process::RegisteredExtensionToolSpec;
6use crate::extensions::runtime::ExtensionHandler;
7use crate::{Result, RuntimeError, ToolContext};
8
9pub struct ExtensionTool {
10    runtime_name: String,
11    description: String,
12    input_schema: Value,
13    handler: Arc<dyn ExtensionHandler>,
14    tool_name: String,
15    plugin_id: String,
16}
17
18impl ExtensionTool {
19    pub fn new(
20        plugin_id: &str,
21        spec: RegisteredExtensionToolSpec,
22        handler: Arc<dyn ExtensionHandler>,
23    ) -> Self {
24        Self {
25            runtime_name: format!("{}:{}", plugin_id, spec.name),
26            description: spec.description,
27            input_schema: spec.input_schema,
28            handler,
29            tool_name: spec.name,
30            plugin_id: plugin_id.to_string(),
31        }
32    }
33}
34
35#[async_trait::async_trait]
36impl crate::tools::Tool for ExtensionTool {
37    fn name(&self) -> &str {
38        &self.runtime_name
39    }
40
41    fn description(&self) -> &str {
42        &self.description
43    }
44
45    fn parameters(&self) -> Value {
46        self.input_schema.clone()
47    }
48
49    async fn execute(&self, params: Value, ctx: ToolContext) -> Result<String> {
50        let binding = ctx
51            .capabilities
52            .memory_backend
53            .unwrap_or_else(crate::memory_backend::MemoryBinding::configured_current);
54        if binding.exclusive() && crate::memory_backend::competing_memory_tool(self) {
55            return Err(RuntimeError::Tool("independent extension memory is disabled by the selected host backend; use short memory tools".into()));
56        }
57        let value = self
58            .handler
59            .call_tool_for(
60                &self.tool_name,
61                params,
62                ctx.capabilities.session_id.as_deref(),
63            )
64            .await
65            .map_err(RuntimeError::Tool)?;
66        if let Some(text) = value.as_str() {
67            Ok(text.to_string())
68        } else if let Some(text) = value.get("content").and_then(Value::as_str) {
69            Ok(text.to_string())
70        } else {
71            Ok(value.to_string())
72        }
73    }
74
75    fn extension_id(&self) -> Option<&str> {
76        Some(&self.plugin_id)
77    }
78
79    fn source_operation(&self) -> Option<&str> {
80        Some(&self.tool_name)
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use std::sync::atomic::{AtomicUsize, Ordering};
87
88    use agent_core::config::MemoryBackendKind;
89    use serde_json::json;
90
91    use super::*;
92    use crate::extensions::hooks::events::{HookEvent, HookResult};
93    use crate::tools::test_helpers::{context_with_memory, memory_binding};
94    use crate::tools::{Tool, ToolRegistry};
95
96    const REFUSAL: &str = "independent extension memory is disabled by the selected host backend; use short memory tools";
97
98    /// An extension process stand-in that counts the tool calls it receives.
99    #[derive(Default)]
100    struct Counting(AtomicUsize);
101
102    #[async_trait::async_trait]
103    impl ExtensionHandler for Counting {
104        fn id(&self) -> &str {
105            "fixture"
106        }
107        async fn handle(&self, _event: &HookEvent) -> HookResult {
108            HookResult::Continue
109        }
110        async fn shutdown(&self) {}
111        async fn call_tool(&self, _name: &str, _input: Value) -> std::result::Result<Value, String> {
112            self.0.fetch_add(1, Ordering::SeqCst);
113            Ok(json!("extension result"))
114        }
115    }
116
117    fn tool(plugin: &str, operation: &str, handler: &Arc<Counting>) -> ExtensionTool {
118        ExtensionTool::new(
119            plugin,
120            RegisteredExtensionToolSpec {
121                name: operation.into(),
122                description: format!("{plugin} {operation}"),
123                input_schema: json!({"type": "object"}),
124            },
125            handler.clone(),
126        )
127    }
128
129    async fn call(tool: &dyn Tool, kind: MemoryBackendKind) -> Result<String> {
130        let dir = tempfile::tempdir().unwrap();
131        let binding = memory_binding(kind, dir.path());
132        tool.execute(json!({}), context_with_memory(binding)).await
133    }
134
135    fn refusal(result: Result<String>) -> String {
136        match result {
137            Err(RuntimeError::Tool(message)) => message,
138            other => panic!("expected a tool refusal, got {other:?}"),
139        }
140    }
141
142    /// Regression: an extension `memory` registering `search` is offered as
143    /// `memory_search` (`memory_search_2` beside the host's own tool), yet
144    /// only its operation `search` was checked.
145    #[tokio::test]
146    async fn an_exclusive_binding_refuses_extension_memory_aliases() {
147        let handler = Arc::new(Counting::default());
148        for (plugin, operation) in [
149            ("memory", "search"),
150            ("memory", "fetch"),
151            ("memory", "store"),
152            ("memory", "forget"),
153            ("plug", "memory.search"),
154            ("Memory", "Search"),
155        ] {
156            let tool = tool(plugin, operation, &handler);
157            assert_eq!(
158                refusal(call(&tool, MemoryBackendKind::Axel).await),
159                REFUSAL,
160                "{plugin}:{operation}"
161            );
162        }
163        assert_eq!(handler.0.load(Ordering::SeqCst), 0, "no call reached the extension");
164    }
165
166    #[tokio::test]
167    async fn an_exclusive_binding_still_refuses_exact_memory_tool_names() {
168        let handler = Arc::new(Counting::default());
169        for operation in ["memory_search", "memory_fetch", "memory_store", "memory_forget"] {
170            let tool = tool("axel-memory-manager", operation, &handler);
171            assert_eq!(refusal(call(&tool, MemoryBackendKind::Axel).await), REFUSAL);
172        }
173        assert_eq!(handler.0.load(Ordering::SeqCst), 0);
174    }
175
176    #[tokio::test]
177    async fn an_exclusive_binding_runs_other_extension_tools() {
178        let handler = Arc::new(Counting::default());
179        for (plugin, operation) in [("weather", "search"), ("memory", "summarize")] {
180            let tool = tool(plugin, operation, &handler);
181            assert_eq!(
182                call(&tool, MemoryBackendKind::Axel).await.unwrap(),
183                "extension result"
184            );
185        }
186        assert_eq!(handler.0.load(Ordering::SeqCst), 2);
187    }
188
189    /// The suffixed wire name the model sees beside the host's own tool
190    /// reaches the extension tool, which refuses.
191    #[tokio::test]
192    async fn the_suffixed_wire_name_of_an_alias_is_refused() {
193        let handler = Arc::new(Counting::default());
194        let registry = ToolRegistry::from_tools_for_tests(vec![
195            Arc::new(crate::tools::memory::MemorySearchTool),
196            Arc::new(tool("memory", "search", &handler)),
197        ]);
198        assert_eq!(registry.runtime_name_for_api("memory_search"), "memory_search");
199        assert_eq!(registry.runtime_name_for_api("memory_search_2"), "memory:search");
200        let alias = registry.get("memory:search").expect("alias registered");
201        assert_eq!(
202            refusal(call(alias.as_ref(), MemoryBackendKind::Axel).await),
203            REFUSAL
204        );
205        assert_eq!(handler.0.load(Ordering::SeqCst), 0);
206    }
207
208    /// No behaviour change without an exclusive binding.
209    #[tokio::test]
210    async fn a_legacy_binding_runs_extension_memory_tools_unchanged() {
211        let handler = Arc::new(Counting::default());
212        for (plugin, operation) in [("memory", "search"), ("axel-memory-manager", "memory_search")] {
213            let tool = tool(plugin, operation, &handler);
214            assert_eq!(
215                call(&tool, MemoryBackendKind::Legacy).await.unwrap(),
216                "extension result"
217            );
218        }
219        assert_eq!(handler.0.load(Ordering::SeqCst), 2);
220    }
221}
agent_engine/tools/extension.rs

agent_engine/tools/
extension.rs
