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
46impl<Input, Handler> FunctionTool<Input, ToolOutput, Handler>
47where
48    Input: JsonSchema,
49{
50    /// Creates a typed Tool whose handler returns canonical rich content.
51    ///
52    /// Unlike [`Self::new`], this constructor preserves the returned
53    /// [`ToolOutput`] instead of serializing it into JSON text. The default
54    /// output schema is permissive because rich presentation content and
55    /// optional structured content are validated by the canonical Tool
56    /// boundary rather than one generated Rust output type.
57    pub fn new_rich(
58        name: impl Into<String>,
59        description: impl Into<String>,
60        handler: Handler,
61    ) -> Self {
62        Self {
63            descriptor: ToolDescriptor {
64                id: CapabilityId::new(),
65                name: name.into(),
66                version: "1".into(),
67                description: description.into(),
68                input_schema: schema_for!(Input).to_value(),
69                output_schema: serde_json::json!({}),
70                effect: EffectClass::Pure,
71                risk: RiskLevel::Low,
72                metadata: std::collections::BTreeMap::new(),
73            },
74            handler,
75            types: PhantomData,
76        }
77    }
78}
79
80impl<Input, Output, Handler> FunctionTool<Input, Output, Handler> {
81    /// Replaces the stable capability identity.
82    #[must_use]
83    pub const fn capability_id(mut self, id: CapabilityId) -> Self {
84        self.descriptor.id = id;
85        self
86    }
87
88    /// Sets the semantic Tool contract version.
89    #[must_use]
90    pub fn version(mut self, version: impl Into<String>) -> Self {
91        self.descriptor.version = version.into();
92        self
93    }
94
95    /// Declares external-effect behavior.
96    #[must_use]
97    pub const fn effect(mut self, effect: EffectClass) -> Self {
98        self.descriptor.effect = effect;
99        self
100    }
101
102    /// Declares policy risk.
103    #[must_use]
104    pub const fn risk(mut self, risk: RiskLevel) -> Self {
105        self.descriptor.risk = risk;
106        self
107    }
108
109    /// Adds host-only namespaced metadata.
110    #[must_use]
111    pub fn metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
112        self.descriptor.metadata.insert(key.into(), value);
113        self
114    }
115
116    /// Replaces the successful output schema.
117    ///
118    /// This is primarily useful for rich Tools that attach typed
119    /// `structured_content` alongside media. The registry compiles and
120    /// enforces the schema before exposing a successful result to an Agent.
121    #[must_use]
122    pub fn output_schema(mut self, schema: serde_json::Value) -> Self {
123        self.descriptor.output_schema = schema;
124        self
125    }
126}
127
128impl<Input, Output, Handler> std::fmt::Debug for FunctionTool<Input, Output, Handler> {
129    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
130        formatter
131            .debug_struct("FunctionTool")
132            .field("descriptor", &self.descriptor)
133            .finish_non_exhaustive()
134    }
135}
136
137impl<Input, Output, Handler, HandlerFuture> Tool for FunctionTool<Input, Output, Handler>
138where
139    Input: DeserializeOwned + JsonSchema + Send + 'static,
140    Output: JsonSchema + Serialize + Send + 'static,
141    Handler: Fn(Input, ToolContext) -> HandlerFuture + Send + Sync,
142    HandlerFuture: Future<Output = Result<Output, ToolError>> + Send + 'static,
143{
144    fn descriptor(&self) -> &ToolDescriptor {
145        &self.descriptor
146    }
147
148    fn invoke(
149        &self,
150        input: serde_json::Value,
151        context: ToolContext,
152    ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
153        let input = match decode_input(input) {
154            Ok(input) => input,
155            Err(error) => {
156                return Box::pin(async move { Err(error) });
157            }
158        };
159        let future = (self.handler)(input, context);
160        Box::pin(async move {
161            let output = future.await?;
162            let value = serde_json::to_value(output).map_err(|error| {
163                ToolError::local(
164                    ToolErrorKind::InvalidOutput,
165                    format!("typed Tool output cannot be serialized: {error}"),
166                )
167            })?;
168            Ok(ToolOutput::model_visible(value))
169        })
170    }
171}
172
173impl<Input, Handler, HandlerFuture> Tool for FunctionTool<Input, ToolOutput, Handler>
174where
175    Input: DeserializeOwned + JsonSchema + Send + 'static,
176    Handler: Fn(Input, ToolContext) -> HandlerFuture + Send + Sync,
177    HandlerFuture: Future<Output = Result<ToolOutput, ToolError>> + Send + 'static,
178{
179    fn descriptor(&self) -> &ToolDescriptor {
180        &self.descriptor
181    }
182
183    fn invoke(
184        &self,
185        input: serde_json::Value,
186        context: ToolContext,
187    ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
188        let input = match decode_input(input) {
189            Ok(input) => input,
190            Err(error) => {
191                return Box::pin(async move { Err(error) });
192            }
193        };
194        let future = (self.handler)(input, context);
195        Box::pin(future)
196    }
197}
198
199fn decode_input<Input: DeserializeOwned>(input: serde_json::Value) -> Result<Input, ToolError> {
200    serde_json::from_value(input).map_err(|error| {
201        ToolError::local(
202            ToolErrorKind::InvalidInput,
203            format!("typed Tool input is invalid: {error}"),
204        )
205    })
206}
207
208#[cfg(test)]
209mod tests {
210    use std::sync::{
211        Arc,
212        atomic::{AtomicUsize, Ordering},
213    };
214
215    use runifold_core::{Budget, BudgetTracker, CapabilitySet, RunContext};
216    use runifold_model::{ContentPart, MediaSource};
217    use schemars::JsonSchema;
218    use serde::{Deserialize, Serialize};
219    use serde_json::json;
220
221    use super::FunctionTool;
222    use crate::{Tool, ToolError, ToolErrorKind, ToolOutput, ToolRegistry};
223
224    #[derive(Deserialize, JsonSchema)]
225    struct AddInput {
226        left: i64,
227        right: i64,
228    }
229
230    #[derive(JsonSchema, Serialize)]
231    struct AddOutput {
232        sum: i64,
233    }
234
235    #[test]
236    fn typed_function_generates_schemas_and_runs_through_registry() {
237        let calls = Arc::new(AtomicUsize::new(0));
238        let observed = calls.clone();
239        let tool = Arc::new(FunctionTool::new(
240            "add",
241            "adds two integers",
242            move |input: AddInput, _context| {
243                let observed = observed.clone();
244                async move {
245                    observed.fetch_add(1, Ordering::SeqCst);
246                    Ok(AddOutput {
247                        sum: input.left + input.right,
248                    })
249                }
250            },
251        ));
252        let descriptor = tool.descriptor();
253        assert_eq!(
254            descriptor.input_schema["required"],
255            json!(["left", "right"])
256        );
257        assert_eq!(descriptor.output_schema["required"], json!(["sum"]));
258
259        let mut capabilities = CapabilitySet::new();
260        capabilities.grant(descriptor.capability());
261        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
262        let mut registry = ToolRegistry::new();
263        registry.register(tool).unwrap();
264
265        let output = futures_executor::block_on(registry.invoke(
266            "add",
267            json!({"left": 2, "right": 3}),
268            &run,
269        ))
270        .unwrap();
271
272        assert_eq!(output.structured_content, Some(json!({"sum": 5})));
273        assert_eq!(calls.load(Ordering::SeqCst), 1);
274    }
275
276    #[test]
277    fn invalid_typed_input_never_calls_handler() {
278        let calls = Arc::new(AtomicUsize::new(0));
279        let observed = calls.clone();
280        let tool = Arc::new(FunctionTool::new(
281            "add",
282            "adds two integers",
283            move |_input: AddInput, _context| {
284                let observed = observed.clone();
285                async move {
286                    observed.fetch_add(1, Ordering::SeqCst);
287                    Ok(AddOutput { sum: 0 })
288                }
289            },
290        ));
291        let mut capabilities = CapabilitySet::new();
292        capabilities.grant(tool.descriptor().capability());
293        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
294        let mut registry = ToolRegistry::new();
295        registry.register(tool).unwrap();
296
297        let error = futures_executor::block_on(registry.invoke("add", json!({"left": 2}), &run))
298            .unwrap_err();
299
300        assert_eq!(error.kind, ToolErrorKind::InvalidInput);
301        assert_eq!(calls.load(Ordering::SeqCst), 0);
302    }
303
304    #[test]
305    fn rich_function_preserves_image_and_structured_content() {
306        let tool = Arc::new(
307            FunctionTool::new_rich(
308                "kline",
309                "returns a K-line chart",
310                |input: AddInput, _context| async move {
311                    Ok::<_, ToolError>(
312                        ToolOutput::rich(vec![
313                            ContentPart::text("K-line chart"),
314                            ContentPart::Image {
315                                source: MediaSource::Url {
316                                    url: "https://example.com/kline.png".into(),
317                                    media_type: Some("image/png".into()),
318                                },
319                            },
320                        ])
321                        .with_structured_content(json!({
322                            "left": input.left,
323                            "right": input.right,
324                        })),
325                    )
326                },
327            )
328            .output_schema(json!({
329                "type": "object",
330                "required": ["left", "right"],
331                "properties": {
332                    "left": { "type": "integer" },
333                    "right": { "type": "integer" }
334                }
335            })),
336        );
337        let mut capabilities = CapabilitySet::new();
338        capabilities.grant(tool.descriptor().capability());
339        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
340        let mut registry = ToolRegistry::new();
341        registry.register(tool).unwrap();
342
343        let output = futures_executor::block_on(registry.invoke(
344            "kline",
345            json!({"left": 20, "right": 22}),
346            &run,
347        ))
348        .unwrap();
349
350        assert_eq!(
351            output.structured_content,
352            Some(json!({"left": 20, "right": 22}))
353        );
354        assert!(matches!(
355            &output.content[1],
356            ContentPart::Image {
357                source: MediaSource::Url { media_type, .. }
358            } if media_type.as_deref() == Some("image/png")
359        ));
360    }
361}
runifold_tool/function.rs

runifold_tool/
function.rs
