Skip to main content

arcp_runtime/runtime/
tools.rs

1//! Tool registry and handler trait.
2//!
3//! User code registers a [`ToolHandler`] for each tool the runtime should
4//! be able to execute. The runtime dispatches `tool.invoke` envelopes by
5//! looking up the handler in the [`ToolRegistry`] and driving it inside a
6//! per-job tokio task with a [`tokio_util::sync::CancellationToken`].
7
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use async_trait::async_trait;
12
13use super::context::ToolContext;
14use arcp_core::error::ARCPError;
15
16/// Application-supplied tool handler.
17///
18/// Implementations should poll `cancel` at safe checkpoints to honour
19/// cooperative cancellation (RFC ยง10.4).
20#[async_trait]
21pub trait ToolHandler: Send + Sync {
22    /// Tool identifier (matches `tool.invoke.payload.tool`).
23    fn name(&self) -> &str;
24
25    /// Run the tool. Return either an inline JSON result or an error.
26    ///
27    /// `arguments` is the raw `arguments` block from the envelope.
28    /// `ctx` is the per-job [`ToolContext`] โ€” the handler polls
29    /// `ctx.cancel` for cooperative cancellation.
30    ///
31    /// # Errors
32    ///
33    /// Implementations return [`ARCPError`] for any failure path. The
34    /// runtime maps the error to a `job.failed` (or `job.cancelled`)
35    /// envelope on the wire.
36    async fn invoke(
37        &self,
38        arguments: serde_json::Value,
39        ctx: ToolContext,
40    ) -> Result<serde_json::Value, ARCPError>;
41}
42
43/// Runtime-owned registry of tools.
44#[derive(Clone, Default)]
45pub struct ToolRegistry {
46    tools: Arc<HashMap<String, Arc<dyn ToolHandler>>>,
47}
48
49impl std::fmt::Debug for ToolRegistry {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        f.debug_struct("ToolRegistry")
52            .field("names", &self.tools.keys().collect::<Vec<_>>())
53            .finish()
54    }
55}
56
57impl ToolRegistry {
58    /// Construct an empty registry.
59    #[must_use]
60    pub fn new() -> Self {
61        Self::default()
62    }
63
64    /// Look up a tool by name.
65    #[must_use]
66    pub fn get(&self, name: &str) -> Option<Arc<dyn ToolHandler>> {
67        self.tools.get(name).cloned()
68    }
69
70    /// Number of registered tools.
71    #[must_use]
72    pub fn len(&self) -> usize {
73        self.tools.len()
74    }
75
76    /// True if no tools are registered.
77    #[must_use]
78    pub fn is_empty(&self) -> bool {
79        self.tools.is_empty()
80    }
81}
82
83/// Builder for [`ToolRegistry`] โ€” accumulate handlers, then `build`.
84#[derive(Default)]
85pub struct ToolRegistryBuilder {
86    tools: HashMap<String, Arc<dyn ToolHandler>>,
87}
88
89impl std::fmt::Debug for ToolRegistryBuilder {
90    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        f.debug_struct("ToolRegistryBuilder")
92            .field("names", &self.tools.keys().collect::<Vec<_>>())
93            .finish()
94    }
95}
96
97impl ToolRegistryBuilder {
98    /// Construct an empty builder.
99    #[must_use]
100    pub fn new() -> Self {
101        Self::default()
102    }
103
104    /// Register `handler` under its declared `name()`.
105    #[must_use]
106    pub fn with(mut self, handler: Arc<dyn ToolHandler>) -> Self {
107        let name = handler.name().to_owned();
108        self.tools.insert(name, handler);
109        self
110    }
111
112    /// Finalise the registry.
113    #[must_use]
114    pub fn build(self) -> ToolRegistry {
115        ToolRegistry {
116            tools: Arc::new(self.tools),
117        }
118    }
119}
120
121#[cfg(test)]
122#[allow(
123    clippy::expect_used,
124    clippy::unwrap_used,
125    clippy::panic,
126    clippy::missing_panics_doc
127)]
128mod tests {
129    use tokio_util::sync::CancellationToken;
130
131    use super::*;
132
133    struct EchoTool;
134
135    #[async_trait]
136    impl ToolHandler for EchoTool {
137        fn name(&self) -> &'static str {
138            "echo"
139        }
140
141        async fn invoke(
142            &self,
143            arguments: serde_json::Value,
144            _ctx: ToolContext,
145        ) -> Result<serde_json::Value, ARCPError> {
146            Ok(arguments)
147        }
148    }
149
150    #[tokio::test]
151    async fn registry_round_trips_through_builder() {
152        let reg = ToolRegistryBuilder::new().with(Arc::new(EchoTool)).build();
153        assert!(!reg.is_empty());
154        assert_eq!(reg.len(), 1);
155        let echo = reg.get("echo").expect("registered");
156        assert_eq!(echo.name(), "echo");
157
158        // Invoking the handler through the trait obj exercises the dyn dispatch.
159        let (tx, _rx) = tokio::sync::mpsc::channel(1);
160        let ctx = ToolContext {
161            cancel: CancellationToken::new(),
162            job_id: arcp_core::ids::JobId::new(),
163            session_id: arcp_core::ids::SessionId::new(),
164            correlation_id: arcp_core::ids::MessageId::new(),
165            out: tx,
166            budget: crate::runtime::context::BudgetTracker::new(),
167            lease: None,
168        };
169        let result = echo
170            .invoke(serde_json::json!({"k": 1}), ctx)
171            .await
172            .expect("invoke");
173        assert_eq!(result, serde_json::json!({"k": 1}));
174    }
175
176    #[test]
177    fn empty_registry_reports_empty() {
178        let reg = ToolRegistry::new();
179        assert!(reg.is_empty());
180        assert_eq!(reg.len(), 0);
181        assert!(reg.get("missing").is_none());
182    }
183
184    #[test]
185    fn debug_impls_render_without_panicking() {
186        let reg = ToolRegistryBuilder::new().with(Arc::new(EchoTool)).build();
187        let s = format!("{reg:?}");
188        assert!(s.contains("echo"));
189        let builder = ToolRegistryBuilder::new().with(Arc::new(EchoTool));
190        let bs = format!("{builder:?}");
191        assert!(bs.contains("echo"));
192    }
193}