cairn_runtime/
tool.rs

1use std::future::Future;
2use std::marker::PhantomData;
3use std::time::Duration;
4
5use async_trait::async_trait;
6use schemars::JsonSchema;
7use serde::de::DeserializeOwned;
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10use sha2::{Digest, Sha256};
11
12/// What a tool does to the world. Policies and retry semantics key off this.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum Effect {
16    /// No side effects; safe to retry freely.
17    Read,
18    /// Changes state but can be undone.
19    Write,
20    /// Cannot be undone (payments, emails, deletes).
21    Irreversible,
22}
23
24#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
25pub struct ToolSpec {
26    pub name: String,
27    pub description: String,
28    pub input_schema: Value,
29    pub effect: Effect,
30}
31
32/// Passed to every tool invocation.
33#[derive(Debug, Clone)]
34pub struct ToolCtx {
35    pub run_id: String,
36    pub call_id: String,
37    /// Stable across retries and crash-recovery re-executions of the same
38    /// call. Forward it to downstream APIs to get exactly-once effects.
39    pub idempotency_key: String,
40}
41
42impl ToolCtx {
43    pub fn new(run_id: &str, call_id: &str) -> Self {
44        let digest = Sha256::digest(format!("{run_id}:{call_id}").as_bytes());
45        Self { run_id: run_id.to_owned(), call_id: call_id.to_owned(), idempotency_key: hex::encode(&digest[..16]) }
46    }
47}
48
49#[derive(Debug, Clone, thiserror::Error)]
50#[error("{0}")]
51pub struct ToolError(pub String);
52
53impl ToolError {
54    pub fn new(msg: impl Into<String>) -> Self {
55        Self(msg.into())
56    }
57}
58
59#[async_trait]
60pub trait Tool: Send + Sync {
61    fn spec(&self) -> &ToolSpec;
62    async fn call(&self, input: Value, ctx: &ToolCtx) -> Result<Value, ToolError>;
63
64    /// Overrides the agent's default tool timeout for this tool.
65    fn timeout(&self) -> Option<Duration> {
66        None
67    }
68}
69
70/// Builds a typed tool from an async function. The JSON schema shown to the
71/// model is derived from `I`, and model input is validated against it.
72///
73/// ```ignore
74/// #[derive(Deserialize, JsonSchema)]
75/// struct GetOrder { order_id: String }
76///
77/// let get_order = tool("get_order", "Look up an order", Effect::Read,
78///     |i: GetOrder, _ctx| async move { Ok(json!({"id": i.order_id})) });
79/// ```
80pub fn tool<I, O, F, Fut>(name: &str, description: &str, effect: Effect, f: F) -> FnTool<I, O, F>
81where
82    I: DeserializeOwned + JsonSchema + Send + 'static,
83    O: Serialize + Send + 'static,
84    F: Fn(I, ToolCtx) -> Fut + Send + Sync + 'static,
85    Fut: Future<Output = Result<O, ToolError>> + Send + 'static,
86{
87    let mut schema = serde_json::to_value(schemars::schema_for!(I)).unwrap_or(Value::Null);
88    if let Value::Object(map) = &mut schema {
89        map.remove("$schema");
90        map.remove("title");
91    }
92    FnTool {
93        spec: ToolSpec { name: name.to_owned(), description: description.to_owned(), input_schema: schema, effect },
94        f,
95        timeout: None,
96        _marker: PhantomData,
97    }
98}
99
100pub struct FnTool<I, O, F> {
101    spec: ToolSpec,
102    f: F,
103    timeout: Option<Duration>,
104    _marker: PhantomData<fn(I) -> O>,
105}
106
107impl<I, O, F> FnTool<I, O, F> {
108    /// Cancel this tool if it runs longer than `limit`; the model receives
109    /// an error result and the run continues.
110    pub fn with_timeout(mut self, limit: Duration) -> Self {
111        self.timeout = Some(limit);
112        self
113    }
114}
115
116#[async_trait]
117impl<I, O, F, Fut> Tool for FnTool<I, O, F>
118where
119    I: DeserializeOwned + JsonSchema + Send + 'static,
120    O: Serialize + Send + 'static,
121    F: Fn(I, ToolCtx) -> Fut + Send + Sync + 'static,
122    Fut: Future<Output = Result<O, ToolError>> + Send + 'static,
123{
124    fn spec(&self) -> &ToolSpec {
125        &self.spec
126    }
127
128    fn timeout(&self) -> Option<Duration> {
129        self.timeout
130    }
131
132    async fn call(&self, input: Value, ctx: &ToolCtx) -> Result<Value, ToolError> {
133        let typed: I = serde_json::from_value(input)
134            .map_err(|e| ToolError(format!("invalid input for `{}`: {e}", self.spec.name)))?;
135        let out = (self.f)(typed, ctx.clone()).await?;
136        serde_json::to_value(out).map_err(|e| ToolError(format!("could not serialize output: {e}")))
137    }
138}
cairn_runtime/tool.rs

cairn_runtime/
tool.rs
