pf_effects/
proxy.rs

1// SPDX-License-Identifier: MIT
2//! `ToolProxy` — wraps a runtime's tool dispatch so every call passes through
3//! the effect ledger.
4//!
5//! Adapter authors register their tools by name once, then invoke through the
6//! proxy. The proxy hashes args, mints an idempotency key, runs the tool,
7//! hashes the result, appends to the ledger, returns the result.
8
9use std::collections::HashMap;
10use std::sync::{Arc, Mutex};
11
12use chrono::Utc;
13use serde_json::Value;
14
15use crate::ledger::{Ledger, SideEffectClass, args_hash, mint_idempotency_key};
16
17/// A dispatchable tool. Synchronous for the v1 proxy; the SDKs wrap this in
18/// `tokio::task::spawn_blocking` for async use.
19pub trait ToolHandler: Send + Sync {
20    /// The tool's declared side-effect class. Misclassification is a contract
21    /// violation by the tool author, not a ledger bug.
22    fn side_effect_class(&self) -> SideEffectClass;
23    /// Invoke the tool. Args and result are arbitrary JSON.
24    fn call(&self, args: &Value) -> pf_core::Result<Value>;
25}
26
27/// Holds tool registrations + the ledger + the secret. `Send + Sync` via
28/// inner `Mutex`; cheap to `clone()`.
29#[derive(Clone)]
30pub struct ToolProxy {
31    inner: Arc<ToolProxyInner>,
32}
33
34struct ToolProxyInner {
35    tools: Mutex<HashMap<String, Arc<dyn ToolHandler>>>,
36    ledger: Mutex<Ledger>,
37}
38
39impl ToolProxy {
40    /// Construct a proxy backed by the given (probably empty) ledger.
41    #[must_use]
42    pub fn new(ledger: Ledger) -> Self {
43        Self {
44            inner: Arc::new(ToolProxyInner {
45                tools: Mutex::new(HashMap::new()),
46                ledger: Mutex::new(ledger),
47            }),
48        }
49    }
50
51    /// Register a tool under `id`. Subsequent registrations under the same
52    /// `id` overwrite (last wins).
53    pub fn register(&self, id: impl Into<String>, handler: Arc<dyn ToolHandler>) {
54        self.inner.tools.lock().unwrap().insert(id.into(), handler);
55    }
56
57    /// Invoke a registered tool. The call is recorded in the ledger before
58    /// being returned to the caller.
59    pub fn invoke(&self, id: &str, args: &Value) -> pf_core::Result<Value> {
60        let handler = {
61            let tools = self.inner.tools.lock().unwrap();
62            tools
63                .get(id)
64                .cloned()
65                .ok_or_else(|| pf_core::Error::Integrity(format!("unregistered tool: {id}")))?
66        };
67        let arg_h = args_hash(args)?;
68        let key = mint_idempotency_key()?;
69        let result = handler.call(args)?;
70        let result_h = args_hash(&result)?;
71        self.inner.ledger.lock().unwrap().append(
72            Utc::now(),
73            id,
74            arg_h,
75            key,
76            result_h,
77            handler.side_effect_class(),
78        )?;
79        Ok(result)
80    }
81
82    /// Snapshot the current ledger (clones it so the caller can serialize
83    /// without holding the lock).
84    pub fn ledger_snapshot(&self) -> Ledger {
85        self.inner.ledger.lock().unwrap().clone()
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92    use crate::ledger::SessionSecret;
93    use serde_json::json;
94
95    struct AddTool;
96    impl ToolHandler for AddTool {
97        fn side_effect_class(&self) -> SideEffectClass {
98            SideEffectClass::Pure
99        }
100        fn call(&self, args: &Value) -> pf_core::Result<Value> {
101            let a = args.get("a").and_then(Value::as_i64).unwrap_or(0);
102            let b = args.get("b").and_then(Value::as_i64).unwrap_or(0);
103            Ok(json!({"sum": a + b}))
104        }
105    }
106
107    struct EmailTool;
108    impl ToolHandler for EmailTool {
109        fn side_effect_class(&self) -> SideEffectClass {
110            SideEffectClass::Irreversible
111        }
112        fn call(&self, _args: &Value) -> pf_core::Result<Value> {
113            Ok(json!({"sent": true}))
114        }
115    }
116
117    #[test]
118    fn invoke_records_in_ledger() {
119        let ledger = Ledger::new(SessionSecret::new(b"t".to_vec()));
120        let proxy = ToolProxy::new(ledger);
121        proxy.register("add", Arc::new(AddTool));
122        let r = proxy.invoke("add", &json!({"a": 2, "b": 40})).unwrap();
123        assert_eq!(r, json!({"sum": 42}));
124        let snap = proxy.ledger_snapshot();
125        assert_eq!(snap.entries().len(), 1);
126        assert_eq!(snap.entries()[0].tool_id, "add");
127        assert_eq!(snap.entries()[0].side_effect_class, SideEffectClass::Pure);
128        snap.verify().unwrap();
129    }
130
131    #[test]
132    fn unknown_tool_errors_cleanly() {
133        let proxy = ToolProxy::new(Ledger::new(SessionSecret::new(b"t".to_vec())));
134        let err = proxy.invoke("missing", &json!({})).unwrap_err();
135        assert!(matches!(err, pf_core::Error::Integrity(_)));
136    }
137
138    #[test]
139    fn mixed_classes_all_chain_correctly() {
140        let proxy = ToolProxy::new(Ledger::new(SessionSecret::new(b"t".to_vec())));
141        proxy.register("add", Arc::new(AddTool));
142        proxy.register("email", Arc::new(EmailTool));
143        for _ in 0..10 {
144            proxy.invoke("add", &json!({"a": 1, "b": 1})).unwrap();
145            proxy.invoke("email", &json!({"to": "x@y"})).unwrap();
146        }
147        let snap = proxy.ledger_snapshot();
148        assert_eq!(snap.entries().len(), 20);
149        snap.verify().unwrap();
150    }
151}
pf_effects/proxy.rs

pf_effects/
proxy.rs
