Skip to main content

rig_compose/
tool.rs

1//! [`Tool`] — the only side-effectful interface available to skills and agents.
2//!
3//! A [`Tool`] is a typed, named, async function with a JSON-Schema-compatible
4//! signature. Two transports satisfy the trait today: [`LocalTool`] (a closure
5//! over a Rust async fn) and — under the `mcp` feature in a later phase — a
6//! remote MCP server. Skills never know the difference.
7
8use std::future::Future;
9use std::pin::Pin;
10use std::sync::Arc;
11
12use async_trait::async_trait;
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15
16use crate::registry::KernelError;
17
18/// Stable, registry-unique identifier for a tool (e.g. `"grammar.query"`,
19/// `"memory.lookup"`, `"sampler.expand"`).
20pub type ToolName = String;
21
22/// Lightweight description of a tool's I/O contract. The `args_schema` and
23/// `result_schema` are JSON-Schema fragments; the LLM-facing rendering layer
24/// uses them to generate `rig` / MCP tool definitions automatically. We do
25/// **not** validate against them at the kernel — validation is the tool's
26/// responsibility — but downstream MCP exporters need them.
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct ToolSchema {
29    pub name: ToolName,
30    pub description: String,
31    pub args_schema: Value,
32    pub result_schema: Value,
33}
34
35/// A composable, side-effectful capability.
36///
37/// Implementations MUST be cheap to clone (typically `Arc`-wrapped state) so
38/// the same tool instance can be referenced from multiple agents'
39/// [`super::registry::ToolRegistry`] slices.
40#[async_trait]
41pub trait Tool: Send + Sync {
42    /// Return this tool's JSON-Schema-compatible contract.
43    fn schema(&self) -> ToolSchema;
44
45    /// Return this tool's registry name.
46    fn name(&self) -> ToolName {
47        self.schema().name
48    }
49
50    /// Invoke the tool with JSON arguments.
51    async fn invoke(&self, args: Value) -> Result<Value, KernelError>;
52}
53
54/// Adapter that turns any `async Fn(Value) -> Result<Value, KernelError>`
55/// into a [`Tool`]. Hosts can use this to surface existing async functions
56/// to the kernel without writing a dedicated tool type.
57pub struct LocalTool {
58    schema: ToolSchema,
59    #[allow(clippy::type_complexity)]
60    f: Arc<
61        dyn Fn(Value) -> Pin<Box<dyn Future<Output = Result<Value, KernelError>> + Send>>
62            + Send
63            + Sync,
64    >,
65}
66
67impl LocalTool {
68    pub fn new<F, Fut>(schema: ToolSchema, f: F) -> Self
69    where
70        F: Fn(Value) -> Fut + Send + Sync + 'static,
71        Fut: Future<Output = Result<Value, KernelError>> + Send + 'static,
72    {
73        Self {
74            schema,
75            f: Arc::new(move |v| Box::pin(f(v))),
76        }
77    }
78}
79
80#[async_trait]
81impl Tool for LocalTool {
82    fn schema(&self) -> ToolSchema {
83        self.schema.clone()
84    }
85
86    fn name(&self) -> ToolName {
87        self.schema.name.clone()
88    }
89
90    async fn invoke(&self, args: Value) -> Result<Value, KernelError> {
91        (self.f)(args).await
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use crate::*;
98    use serde_json::json;
99
100    #[tokio::test]
101    async fn local_tool_roundtrip() {
102        let schema = ToolSchema {
103            name: "test.echo".into(),
104            description: "echoes the input".into(),
105            args_schema: json!({"type": "object"}),
106            result_schema: json!({"type": "object"}),
107        };
108        let tool = LocalTool::new(schema, |v| async move { Ok(v) });
109        let out = tool.invoke(json!({"hello": "world"})).await.unwrap();
110        assert_eq!(out, json!({"hello": "world"}));
111        assert_eq!(tool.name(), "test.echo");
112    }
113}