runifold_tool/
registry.rs

1use std::{collections::BTreeMap, fmt, sync::Arc};
2
3use futures_util::future::{Either, select};
4use jsonschema::Validator;
5use runifold_core::RunContext;
6use runifold_model::{ArtifactScope, ArtifactStore, ToolSpec};
7use serde_json::Value;
8
9use crate::{
10    Tool, ToolContext, ToolDescriptor, ToolError, ToolErrorKind, ToolFuture, ToolOutput,
11    ToolRegistrationError,
12};
13
14/// Immutable-name registry and capability gate for tools.
15#[derive(Clone, Copy, Debug, Eq, PartialEq)]
16pub struct ToolLimits {
17    /// Maximum serialized invocation input size.
18    pub max_input_bytes: usize,
19    /// Maximum serialized canonical output size.
20    pub max_output_bytes: usize,
21}
22
23impl Default for ToolLimits {
24    fn default() -> Self {
25        Self {
26            max_input_bytes: 1024 * 1024,
27            max_output_bytes: 8 * 1024 * 1024,
28        }
29    }
30}
31
32#[derive(Clone)]
33struct RegisteredTool {
34    tool: Arc<dyn Tool>,
35    input_validator: Validator,
36    output_validator: Validator,
37}
38
39/// Immutable-name Tool registry with compiled contracts and bounded I/O.
40#[derive(Clone, Default)]
41pub struct ToolRegistry {
42    tools: BTreeMap<String, RegisteredTool>,
43    limits: ToolLimits,
44    artifact_store: Option<Arc<dyn ArtifactStore>>,
45    artifact_scope: Option<ArtifactScope>,
46}
47
48impl ToolRegistry {
49    /// Creates an empty registry.
50    pub fn new() -> Self {
51        Self::default()
52    }
53
54    /// Replaces the serialized Tool I/O limits.
55    #[must_use]
56    pub const fn with_limits(mut self, limits: ToolLimits) -> Self {
57        self.limits = limits;
58        self
59    }
60
61    /// Makes an artifact store available to every Tool invocation context.
62    #[must_use]
63    pub fn with_artifact_store(
64        mut self,
65        scope: ArtifactScope,
66        store: Arc<dyn ArtifactStore>,
67    ) -> Self {
68        self.artifact_scope = Some(scope);
69        self.artifact_store = Some(store);
70        self
71    }
72
73    /// Registers a tool without replacing an existing name.
74    ///
75    /// # Errors
76    ///
77    /// Returns [`ToolRegistrationError`] for blank or duplicate names.
78    pub fn register(&mut self, tool: Arc<dyn Tool>) -> Result<(), ToolRegistrationError> {
79        let name = tool.descriptor().name.trim();
80        if name.is_empty() {
81            return Err(ToolRegistrationError::EmptyName);
82        }
83        if self.tools.contains_key(name) {
84            return Err(ToolRegistrationError::DuplicateName(name.into()));
85        }
86        let input_validator = compile_schema(name, "input", &tool.descriptor().input_schema)?;
87        let output_validator = compile_schema(name, "output", &tool.descriptor().output_schema)?;
88        self.tools.insert(
89            name.into(),
90            RegisteredTool {
91                tool,
92                input_validator,
93                output_validator,
94            },
95        );
96        Ok(())
97    }
98
99    /// Returns model-facing specifications in deterministic name order.
100    pub fn model_specs(&self) -> Vec<ToolSpec> {
101        self.tools
102            .values()
103            .map(|registered| registered.tool.descriptor().model_spec())
104            .collect()
105    }
106
107    /// Returns the number of registered tools.
108    pub fn len(&self) -> usize {
109        self.tools.len()
110    }
111
112    /// Returns whether no tools are registered.
113    pub fn is_empty(&self) -> bool {
114        self.tools.is_empty()
115    }
116
117    /// Returns whether a tool is registered under `name`.
118    pub fn contains(&self, name: &str) -> bool {
119        self.tools.contains_key(name)
120    }
121
122    /// Returns the immutable descriptor registered under `name`.
123    pub fn descriptor(&self, name: &str) -> Option<&ToolDescriptor> {
124        self.tools
125            .get(name)
126            .map(|registered| registered.tool.descriptor())
127    }
128
129    /// Invokes a registered tool after checking the owning run's explicit
130    /// capability grant.
131    pub fn invoke<'a>(
132        &'a self,
133        name: &'a str,
134        input: Value,
135        run: &'a RunContext,
136    ) -> ToolFuture<'a, Result<ToolOutput, ToolError>> {
137        Box::pin(async move {
138            let registered = self.tools.get(name).ok_or_else(|| {
139                ToolError::local(
140                    ToolErrorKind::NotFound,
141                    format!("tool `{name}` is not registered"),
142                )
143            })?;
144            let descriptor = registered.tool.descriptor();
145            if !run.capabilities().contains(descriptor.id) {
146                return Err(ToolError::local(
147                    ToolErrorKind::CapabilityDenied,
148                    format!("run is not granted tool capability `{name}`"),
149                ));
150            }
151            let context = ToolContext::for_run(run)
152                .with_artifact_store(self.artifact_scope.clone(), self.artifact_store.clone());
153            validate_size(
154                "Tool input",
155                &input,
156                self.limits.max_input_bytes,
157                ToolErrorKind::InvalidInput,
158            )?;
159            registered
160                .input_validator
161                .validate(&input)
162                .map_err(|error| {
163                    ToolError::local(
164                        ToolErrorKind::InvalidInput,
165                        format!(
166                            "Tool input violates its declared schema at `{}`",
167                            error.schema_path()
168                        ),
169                    )
170                })?;
171            if context
172                .remaining()
173                .is_some_and(|remaining| remaining.is_zero())
174            {
175                return Err(ToolError::local(
176                    ToolErrorKind::DeadlineExceeded,
177                    "tool invocation deadline already elapsed",
178                ));
179            }
180            let cancellation = context.cancellation().clone();
181            match select(
182                Box::pin(cancellation.cancelled()),
183                Box::pin(registered.tool.invoke(input, context)),
184            )
185            .await
186            {
187                Either::Left(_) => Err(ToolError::local(
188                    ToolErrorKind::Cancelled,
189                    "tool invocation was cancelled",
190                )),
191                Either::Right((result, _)) => {
192                    let output = result?;
193                    validate_output(&output, &registered.output_validator, self.limits)?;
194                    Ok(output)
195                }
196            }
197        })
198    }
199}
200
201fn compile_schema(
202    tool: &str,
203    direction: &'static str,
204    schema: &Value,
205) -> Result<Validator, ToolRegistrationError> {
206    jsonschema::validator_for(schema).map_err(|error| ToolRegistrationError::InvalidSchema {
207        tool: tool.into(),
208        direction,
209        message: error.to_string(),
210    })
211}
212
213fn validate_output(
214    output: &ToolOutput,
215    validator: &Validator,
216    limits: ToolLimits,
217) -> Result<(), ToolError> {
218    if output.content.is_empty() {
219        return Err(ToolError::local(
220            ToolErrorKind::InvalidOutput,
221            "Tool output content cannot be empty",
222        ));
223    }
224    validate_size(
225        "Tool output",
226        output,
227        limits.max_output_bytes,
228        ToolErrorKind::InvalidOutput,
229    )?;
230    if output.is_error {
231        return Ok(());
232    }
233    let instance = output
234        .structured_content
235        .clone()
236        .unwrap_or_else(|| serde_json::to_value(&output.content).unwrap_or(Value::Null));
237    validator.validate(&instance).map_err(|error| {
238        ToolError::local(
239            ToolErrorKind::InvalidOutput,
240            format!(
241                "Tool output violates its declared schema at `{}`",
242                error.schema_path()
243            ),
244        )
245    })
246}
247
248fn validate_size<T: serde::Serialize>(
249    label: &str,
250    value: &T,
251    limit: usize,
252    kind: ToolErrorKind,
253) -> Result<(), ToolError> {
254    let size = serde_json::to_vec(value)
255        .map_err(|error| {
256            ToolError::local(kind.clone(), format!("{label} cannot be encoded: {error}"))
257        })?
258        .len();
259    if size > limit {
260        return Err(ToolError::local(
261            kind,
262            format!("{label} is {size} bytes and exceeds the {limit}-byte limit"),
263        ));
264    }
265    Ok(())
266}
267
268impl fmt::Debug for ToolRegistry {
269    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
270        formatter
271            .debug_struct("ToolRegistry")
272            .field("tools", &self.tools.keys().collect::<Vec<_>>())
273            .field("limits", &self.limits)
274            .field("artifact_store", &self.artifact_store.is_some())
275            .field("artifact_scope", &self.artifact_scope)
276            .finish()
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use std::{collections::BTreeMap, sync::Arc};
283
284    use runifold_core::{
285        Budget, BudgetTracker, CapabilityId, CapabilitySet, EffectClass, RiskLevel, RunContext,
286    };
287    use serde_json::{Value, json};
288
289    use crate::{
290        Tool, ToolContext, ToolDescriptor, ToolError, ToolErrorKind, ToolFuture, ToolLimits,
291        ToolOutput, ToolRegistrationError,
292    };
293
294    use super::ToolRegistry;
295
296    #[derive(Debug)]
297    struct EchoTool {
298        descriptor: ToolDescriptor,
299    }
300
301    impl EchoTool {
302        fn new(name: &str) -> Self {
303            Self {
304                descriptor: ToolDescriptor {
305                    id: CapabilityId::new(),
306                    name: name.into(),
307                    version: "1".into(),
308                    description: "Echo structured input".into(),
309                    input_schema: json!({"type": "object"}),
310                    output_schema: json!({"type": "object"}),
311                    effect: EffectClass::Pure,
312                    risk: RiskLevel::Low,
313                    metadata: BTreeMap::new(),
314                },
315            }
316        }
317    }
318
319    impl Tool for EchoTool {
320        fn descriptor(&self) -> &ToolDescriptor {
321            &self.descriptor
322        }
323
324        fn invoke(
325            &self,
326            input: Value,
327            _context: ToolContext,
328        ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
329            Box::pin(async move { Ok(ToolOutput::model_visible(input)) })
330        }
331    }
332
333    #[test]
334    fn registry_requires_explicit_capability_grants() {
335        let tool = Arc::new(EchoTool::new("echo"));
336        let mut registry = ToolRegistry::new();
337        registry.register(tool).unwrap();
338        let run = RunContext::root(BudgetTracker::new(Budget::default()), CapabilitySet::new());
339
340        let error =
341            futures_executor::block_on(registry.invoke("echo", json!({"x": 1}), &run)).unwrap_err();
342
343        assert_eq!(error.kind, ToolErrorKind::CapabilityDenied);
344    }
345
346    #[test]
347    fn granted_tools_execute_through_the_object_safe_boundary() {
348        let tool = Arc::new(EchoTool::new("echo"));
349        let mut capabilities = CapabilitySet::new();
350        capabilities.grant(tool.descriptor().capability());
351        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
352        let mut registry = ToolRegistry::new();
353        registry.register(tool).unwrap();
354
355        let output =
356            futures_executor::block_on(registry.invoke("echo", json!({"x": 1}), &run)).unwrap();
357
358        assert_eq!(output.structured_content, Some(json!({"x": 1})));
359    }
360
361    #[test]
362    fn duplicate_names_are_rejected_instead_of_replaced() {
363        let mut registry = ToolRegistry::new();
364        registry.register(Arc::new(EchoTool::new("echo"))).unwrap();
365
366        let error = registry
367            .register(Arc::new(EchoTool::new("echo")))
368            .unwrap_err();
369
370        assert_eq!(error, ToolRegistrationError::DuplicateName("echo".into()));
371    }
372
373    #[test]
374    fn invalid_schemas_fail_during_registration() {
375        let mut tool = EchoTool::new("invalid");
376        tool.descriptor.input_schema = json!({"type":"not-a-json-schema-type"});
377        let error = ToolRegistry::new().register(Arc::new(tool)).unwrap_err();
378        assert!(matches!(
379            error,
380            ToolRegistrationError::InvalidSchema {
381                direction: "input",
382                ..
383            }
384        ));
385    }
386
387    #[test]
388    fn output_contract_and_size_are_enforced_after_execution() {
389        let tool = Arc::new(EchoTool::new("bounded"));
390        let mut capabilities = CapabilitySet::new();
391        capabilities.grant(tool.descriptor().capability());
392        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
393        let mut registry = ToolRegistry::new().with_limits(ToolLimits {
394            max_input_bytes: 1024,
395            max_output_bytes: 32,
396        });
397        registry.register(tool).unwrap();
398
399        let error = futures_executor::block_on(registry.invoke(
400            "bounded",
401            json!({"payload":"this output is intentionally larger than the limit"}),
402            &run,
403        ))
404        .unwrap_err();
405        assert_eq!(error.kind, ToolErrorKind::InvalidOutput);
406    }
407}
runifold_tool/registry.rs

runifold_tool/
registry.rs
