runifold_tool/
registry.rs

1use std::{collections::BTreeMap, fmt, fmt::Write as _, sync::Arc};
2
3use futures_util::future::{Either, select};
4use jsonschema::{ValidationError, Validator, error::ValidationErrorKind};
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| input_validation_error(name, &error))?;
163            if context
164                .remaining()
165                .is_some_and(|remaining| remaining.is_zero())
166            {
167                return Err(ToolError::local(
168                    ToolErrorKind::DeadlineExceeded,
169                    "tool invocation deadline already elapsed",
170                ));
171            }
172            let cancellation = context.cancellation().clone();
173            match select(
174                Box::pin(cancellation.cancelled()),
175                Box::pin(registered.tool.invoke(input, context)),
176            )
177            .await
178            {
179                Either::Left(_) => Err(ToolError::local(
180                    ToolErrorKind::Cancelled,
181                    "tool invocation was cancelled",
182                )),
183                Either::Right((result, _)) => {
184                    let output = result?;
185                    validate_output(&output, &registered.output_validator, self.limits)?;
186                    Ok(output)
187                }
188            }
189        })
190    }
191}
192
193fn input_validation_error(tool: &str, error: &ValidationError<'_>) -> ToolError {
194    let input_path = error.instance_path().to_string();
195    let schema_path = error.schema_path().to_string();
196    let keyword = error.kind().keyword().to_owned();
197    let mut diagnostic = ToolError::local(
198        ToolErrorKind::InvalidInput,
199        format!(
200            "Tool `{tool}` rejected its arguments: input at `{input_path}` violates `{keyword}` at schema `{schema_path}`; call the tool again with corrected arguments and do not guess its result"
201        ),
202    );
203    diagnostic
204        .metadata
205        .insert("validation.input_path".into(), Value::String(input_path));
206    diagnostic
207        .metadata
208        .insert("validation.schema_path".into(), Value::String(schema_path));
209    diagnostic
210        .metadata
211        .insert("validation.keyword".into(), Value::String(keyword));
212    diagnostic
213        .metadata
214        .insert("tool.name".into(), Value::String(tool.into()));
215    if let ValidationErrorKind::Enum { options } = error.kind() {
216        diagnostic
217            .metadata
218            .insert("validation.allowed_values".into(), options.clone());
219        let _ = write!(
220            diagnostic.message,
221            "; allowed values: {}",
222            bounded_json(options, 512)
223        );
224    }
225    if is_safe_scalar(error.instance()) {
226        diagnostic.metadata.insert(
227            "validation.actual_value".into(),
228            error.instance().as_ref().clone(),
229        );
230        let _ = write!(
231            diagnostic.message,
232            "; received: {}",
233            bounded_json(error.instance(), 128)
234        );
235    }
236    diagnostic
237}
238
239fn is_safe_scalar(value: &Value) -> bool {
240    match value {
241        Value::Null | Value::Bool(_) | Value::Number(_) => true,
242        Value::String(value) => value.len() <= 64,
243        Value::Array(_) | Value::Object(_) => false,
244    }
245}
246
247fn bounded_json(value: &Value, maximum: usize) -> String {
248    let encoded = value.to_string();
249    if encoded.len() <= maximum {
250        encoded
251    } else {
252        "<redacted: value exceeds diagnostic limit>".into()
253    }
254}
255
256fn compile_schema(
257    tool: &str,
258    direction: &'static str,
259    schema: &Value,
260) -> Result<Validator, ToolRegistrationError> {
261    jsonschema::validator_for(schema).map_err(|error| ToolRegistrationError::InvalidSchema {
262        tool: tool.into(),
263        direction,
264        message: error.to_string(),
265    })
266}
267
268fn validate_output(
269    output: &ToolOutput,
270    validator: &Validator,
271    limits: ToolLimits,
272) -> Result<(), ToolError> {
273    if output.content.is_empty() {
274        return Err(ToolError::local(
275            ToolErrorKind::InvalidOutput,
276            "Tool output content cannot be empty",
277        ));
278    }
279    validate_size(
280        "Tool output",
281        output,
282        limits.max_output_bytes,
283        ToolErrorKind::InvalidOutput,
284    )?;
285    if output.is_error {
286        return Ok(());
287    }
288    let instance = output
289        .structured_content
290        .clone()
291        .unwrap_or_else(|| serde_json::to_value(&output.content).unwrap_or(Value::Null));
292    validator.validate(&instance).map_err(|error| {
293        ToolError::local(
294            ToolErrorKind::InvalidOutput,
295            format!(
296                "Tool output violates its declared schema at `{}`",
297                error.schema_path()
298            ),
299        )
300    })
301}
302
303fn validate_size<T: serde::Serialize>(
304    label: &str,
305    value: &T,
306    limit: usize,
307    kind: ToolErrorKind,
308) -> Result<(), ToolError> {
309    let size = serde_json::to_vec(value)
310        .map_err(|error| {
311            ToolError::local(kind.clone(), format!("{label} cannot be encoded: {error}"))
312        })?
313        .len();
314    if size > limit {
315        return Err(ToolError::local(
316            kind,
317            format!("{label} is {size} bytes and exceeds the {limit}-byte limit"),
318        ));
319    }
320    Ok(())
321}
322
323impl fmt::Debug for ToolRegistry {
324    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
325        formatter
326            .debug_struct("ToolRegistry")
327            .field("tools", &self.tools.keys().collect::<Vec<_>>())
328            .field("limits", &self.limits)
329            .field("artifact_store", &self.artifact_store.is_some())
330            .field("artifact_scope", &self.artifact_scope)
331            .finish()
332    }
333}
334
335#[cfg(test)]
336mod tests {
337    use std::{collections::BTreeMap, sync::Arc};
338
339    use runifold_core::{
340        Budget, BudgetTracker, CapabilityId, CapabilitySet, EffectClass, RiskLevel, RunContext,
341    };
342    use serde_json::{Value, json};
343
344    use crate::{
345        Tool, ToolContext, ToolDescriptor, ToolError, ToolErrorKind, ToolFuture, ToolLimits,
346        ToolOutput, ToolRegistrationError,
347    };
348
349    use super::ToolRegistry;
350
351    #[derive(Debug)]
352    struct EchoTool {
353        descriptor: ToolDescriptor,
354    }
355
356    impl EchoTool {
357        fn new(name: &str) -> Self {
358            Self {
359                descriptor: ToolDescriptor {
360                    id: CapabilityId::new(),
361                    name: name.into(),
362                    version: "1".into(),
363                    description: "Echo structured input".into(),
364                    input_schema: json!({"type": "object"}),
365                    output_schema: json!({"type": "object"}),
366                    effect: EffectClass::Pure,
367                    risk: RiskLevel::Low,
368                    metadata: BTreeMap::new(),
369                },
370            }
371        }
372    }
373
374    impl Tool for EchoTool {
375        fn descriptor(&self) -> &ToolDescriptor {
376            &self.descriptor
377        }
378
379        fn invoke(
380            &self,
381            input: Value,
382            _context: ToolContext,
383        ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
384            Box::pin(async move { Ok(ToolOutput::model_visible(input)) })
385        }
386    }
387
388    #[test]
389    fn registry_requires_explicit_capability_grants() {
390        let tool = Arc::new(EchoTool::new("echo"));
391        let mut registry = ToolRegistry::new();
392        registry.register(tool).unwrap();
393        let run = RunContext::root(BudgetTracker::new(Budget::default()), CapabilitySet::new());
394
395        let error =
396            futures_executor::block_on(registry.invoke("echo", json!({"x": 1}), &run)).unwrap_err();
397
398        assert_eq!(error.kind, ToolErrorKind::CapabilityDenied);
399    }
400
401    #[test]
402    fn granted_tools_execute_through_the_object_safe_boundary() {
403        let tool = Arc::new(EchoTool::new("echo"));
404        let mut capabilities = CapabilitySet::new();
405        capabilities.grant(tool.descriptor().capability());
406        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
407        let mut registry = ToolRegistry::new();
408        registry.register(tool).unwrap();
409
410        let output =
411            futures_executor::block_on(registry.invoke("echo", json!({"x": 1}), &run)).unwrap();
412
413        assert_eq!(output.structured_content, Some(json!({"x": 1})));
414    }
415
416    #[test]
417    fn invalid_enum_reports_input_path_allowed_values_and_safe_actual_value() {
418        let mut tool = EchoTool::new("market_history");
419        tool.descriptor.input_schema = json!({
420            "type": "object",
421            "properties": {
422                "lookback": {"type": "string", "enum": ["1d", "7d", "30d"]}
423            },
424            "required": ["lookback"]
425        });
426        let tool = Arc::new(tool);
427        let mut capabilities = CapabilitySet::new();
428        capabilities.grant(tool.descriptor().capability());
429        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
430        let mut registry = ToolRegistry::new();
431        registry.register(tool).unwrap();
432
433        let error = futures_executor::block_on(registry.invoke(
434            "market_history",
435            json!({"lookback": "forever"}),
436            &run,
437        ))
438        .unwrap_err();
439
440        assert_eq!(error.kind, ToolErrorKind::InvalidInput);
441        assert_eq!(error.metadata["tool.name"], "market_history");
442        assert_eq!(error.metadata["validation.input_path"], "/lookback");
443        assert_eq!(error.metadata["validation.keyword"], "enum");
444        assert_eq!(error.metadata["validation.actual_value"], "forever");
445        assert_eq!(
446            error.metadata["validation.allowed_values"],
447            json!(["1d", "7d", "30d"])
448        );
449        assert!(error.message.contains("call the tool again"));
450    }
451
452    #[test]
453    fn duplicate_names_are_rejected_instead_of_replaced() {
454        let mut registry = ToolRegistry::new();
455        registry.register(Arc::new(EchoTool::new("echo"))).unwrap();
456
457        let error = registry
458            .register(Arc::new(EchoTool::new("echo")))
459            .unwrap_err();
460
461        assert_eq!(error, ToolRegistrationError::DuplicateName("echo".into()));
462    }
463
464    #[test]
465    fn invalid_schemas_fail_during_registration() {
466        let mut tool = EchoTool::new("invalid");
467        tool.descriptor.input_schema = json!({"type":"not-a-json-schema-type"});
468        let error = ToolRegistry::new().register(Arc::new(tool)).unwrap_err();
469        assert!(matches!(
470            error,
471            ToolRegistrationError::InvalidSchema {
472                direction: "input",
473                ..
474            }
475        ));
476    }
477
478    #[test]
479    fn output_contract_and_size_are_enforced_after_execution() {
480        let tool = Arc::new(EchoTool::new("bounded"));
481        let mut capabilities = CapabilitySet::new();
482        capabilities.grant(tool.descriptor().capability());
483        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
484        let mut registry = ToolRegistry::new().with_limits(ToolLimits {
485            max_input_bytes: 1024,
486            max_output_bytes: 32,
487        });
488        registry.register(tool).unwrap();
489
490        let error = futures_executor::block_on(registry.invoke(
491            "bounded",
492            json!({"payload":"this output is intentionally larger than the limit"}),
493            &run,
494        ))
495        .unwrap_err();
496        assert_eq!(error.kind, ToolErrorKind::InvalidOutput);
497    }
498}
runifold_tool/registry.rs

runifold_tool/
registry.rs
