Skip to main content

runifold_tool/
function.rs

1use std::{future::Future, marker::PhantomData};
2
3use runifold_core::{CapabilityId, EffectClass, RiskLevel};
4use schemars::{JsonSchema, schema_for};
5use serde::{Serialize, de::DeserializeOwned};
6
7use crate::{Tool, ToolContext, ToolDescriptor, ToolError, ToolErrorKind, ToolFuture, ToolOutput};
8
9/// A typed asynchronous Rust function exposed through the canonical Tool
10/// boundary.
11pub struct FunctionTool<Input, Output, Handler> {
12    descriptor: ToolDescriptor,
13    handler: Handler,
14    types: PhantomData<fn(Input) -> Output>,
15}
16
17impl<Input, Output, Handler> FunctionTool<Input, Output, Handler>
18where
19    Input: JsonSchema,
20    Output: JsonSchema,
21{
22    /// Creates a typed Tool with generated input and output JSON Schemas.
23    ///
24    /// The default effect is [`EffectClass::Pure`] and the default risk is
25    /// [`RiskLevel::Low`]. Callers must explicitly override these values for
26    /// functions that read or modify external state.
27    pub fn new(name: impl Into<String>, description: impl Into<String>, handler: Handler) -> Self {
28        Self {
29            descriptor: ToolDescriptor {
30                id: CapabilityId::new(),
31                name: name.into(),
32                version: "1".into(),
33                description: description.into(),
34                input_schema: schema_for!(Input).to_value(),
35                output_schema: schema_for!(Output).to_value(),
36                effect: EffectClass::Pure,
37                risk: RiskLevel::Low,
38                metadata: std::collections::BTreeMap::new(),
39            },
40            handler,
41            types: PhantomData,
42        }
43    }
44
45    /// Replaces the stable capability identity.
46    #[must_use]
47    pub const fn capability_id(mut self, id: CapabilityId) -> Self {
48        self.descriptor.id = id;
49        self
50    }
51
52    /// Sets the semantic Tool contract version.
53    #[must_use]
54    pub fn version(mut self, version: impl Into<String>) -> Self {
55        self.descriptor.version = version.into();
56        self
57    }
58
59    /// Declares external-effect behavior.
60    #[must_use]
61    pub const fn effect(mut self, effect: EffectClass) -> Self {
62        self.descriptor.effect = effect;
63        self
64    }
65
66    /// Declares policy risk.
67    #[must_use]
68    pub const fn risk(mut self, risk: RiskLevel) -> Self {
69        self.descriptor.risk = risk;
70        self
71    }
72
73    /// Adds host-only namespaced metadata.
74    #[must_use]
75    pub fn metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
76        self.descriptor.metadata.insert(key.into(), value);
77        self
78    }
79}
80
81impl<Input, Output, Handler> std::fmt::Debug for FunctionTool<Input, Output, Handler> {
82    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        formatter
84            .debug_struct("FunctionTool")
85            .field("descriptor", &self.descriptor)
86            .finish_non_exhaustive()
87    }
88}
89
90impl<Input, Output, Handler, HandlerFuture> Tool for FunctionTool<Input, Output, Handler>
91where
92    Input: DeserializeOwned + JsonSchema + Send + 'static,
93    Output: JsonSchema + Serialize + Send + 'static,
94    Handler: Fn(Input, ToolContext) -> HandlerFuture + Send + Sync,
95    HandlerFuture: Future<Output = Result<Output, ToolError>> + Send + 'static,
96{
97    fn descriptor(&self) -> &ToolDescriptor {
98        &self.descriptor
99    }
100
101    fn invoke(
102        &self,
103        input: serde_json::Value,
104        context: ToolContext,
105    ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
106        let input = match serde_json::from_value(input) {
107            Ok(input) => input,
108            Err(error) => {
109                return Box::pin(async move {
110                    Err(ToolError::local(
111                        ToolErrorKind::InvalidInput,
112                        format!("typed Tool input is invalid: {error}"),
113                    ))
114                });
115            }
116        };
117        let future = (self.handler)(input, context);
118        Box::pin(async move {
119            let output = future.await?;
120            let value = serde_json::to_value(output).map_err(|error| {
121                ToolError::local(
122                    ToolErrorKind::InvalidOutput,
123                    format!("typed Tool output cannot be serialized: {error}"),
124                )
125            })?;
126            Ok(ToolOutput::model_visible(value))
127        })
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use std::sync::{
134        Arc,
135        atomic::{AtomicUsize, Ordering},
136    };
137
138    use runifold_core::{Budget, BudgetTracker, CapabilitySet, RunContext};
139    use schemars::JsonSchema;
140    use serde::{Deserialize, Serialize};
141    use serde_json::json;
142
143    use super::FunctionTool;
144    use crate::{Tool, ToolErrorKind, ToolRegistry};
145
146    #[derive(Deserialize, JsonSchema)]
147    struct AddInput {
148        left: i64,
149        right: i64,
150    }
151
152    #[derive(JsonSchema, Serialize)]
153    struct AddOutput {
154        sum: i64,
155    }
156
157    #[test]
158    fn typed_function_generates_schemas_and_runs_through_registry() {
159        let calls = Arc::new(AtomicUsize::new(0));
160        let observed = calls.clone();
161        let tool = Arc::new(FunctionTool::new(
162            "add",
163            "adds two integers",
164            move |input: AddInput, _context| {
165                let observed = observed.clone();
166                async move {
167                    observed.fetch_add(1, Ordering::SeqCst);
168                    Ok(AddOutput {
169                        sum: input.left + input.right,
170                    })
171                }
172            },
173        ));
174        let descriptor = tool.descriptor();
175        assert_eq!(
176            descriptor.input_schema["required"],
177            json!(["left", "right"])
178        );
179        assert_eq!(descriptor.output_schema["required"], json!(["sum"]));
180
181        let mut capabilities = CapabilitySet::new();
182        capabilities.grant(descriptor.capability());
183        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
184        let mut registry = ToolRegistry::new();
185        registry.register(tool).unwrap();
186
187        let output = futures_executor::block_on(registry.invoke(
188            "add",
189            json!({"left": 2, "right": 3}),
190            &run,
191        ))
192        .unwrap();
193
194        assert_eq!(output.structured_content, Some(json!({"sum": 5})));
195        assert_eq!(calls.load(Ordering::SeqCst), 1);
196    }
197
198    #[test]
199    fn invalid_typed_input_never_calls_handler() {
200        let calls = Arc::new(AtomicUsize::new(0));
201        let observed = calls.clone();
202        let tool = Arc::new(FunctionTool::new(
203            "add",
204            "adds two integers",
205            move |_input: AddInput, _context| {
206                let observed = observed.clone();
207                async move {
208                    observed.fetch_add(1, Ordering::SeqCst);
209                    Ok(AddOutput { sum: 0 })
210                }
211            },
212        ));
213        let mut capabilities = CapabilitySet::new();
214        capabilities.grant(tool.descriptor().capability());
215        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
216        let mut registry = ToolRegistry::new();
217        registry.register(tool).unwrap();
218
219        let error = futures_executor::block_on(registry.invoke("add", json!({"left": 2}), &run))
220            .unwrap_err();
221
222        assert_eq!(error.kind, ToolErrorKind::InvalidInput);
223        assert_eq!(calls.load(Ordering::SeqCst), 0);
224    }
225}