agentd/store/
mcp.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! The **MCP-mapped store**: the four store operations are `tools/call`s
3//! against a declared MCP server, with argument templates and result extraction
4//! from `store.mcp.{put,get,list,delete}`. The **default mapping is the
5//! checkpointer profile** (`state.put` / `state.get` / `state.list` /
6//! `state.delete`), so a server advertising those tools needs no mapping at all.
7//!
8//! Every call carries `_meta["agent/idempotency_key"] = "<key>#<seq>"` and
9//! `_meta["agent/instance"]`, and is bounded by the store timeout.
10
11use super::mapping::{self, Vars};
12use super::{KeySeq, PutOutcome, Store, StoreError};
13use crate::config::v2::{StoreMcp, StoreOp};
14use crate::wire::mcp::CallToolResult;
15use serde_json::{Value, json};
16use std::sync::Arc;
17use std::time::Duration;
18
19/// The one thing the adapter needs from a connected server: `tools/call`
20/// with per-call `_meta` and a timeout. Implemented by [`crate::mcp::client::McpClient`]
21/// and by test doubles.
22pub trait McpCall: Send + Sync {
23    fn call(
24        &self,
25        tool: &str,
26        args: Value,
27        meta: Value,
28        timeout: Duration,
29    ) -> Result<CallToolResult, String>;
30    fn server_name(&self) -> String;
31}
32
33impl McpCall for crate::mcp::client::McpClient {
34    fn call(
35        &self,
36        tool: &str,
37        args: Value,
38        meta: Value,
39        timeout: Duration,
40    ) -> Result<CallToolResult, String> {
41        self.call_tool_with_meta_within(tool, Some(args), meta, timeout)
42            .map_err(|e| e.to_string())
43    }
44    fn server_name(&self) -> String {
45        self.name().to_string()
46    }
47}
48
49/// The default checkpointer profile: `state.put` / `state.get` / `state.list` /
50/// `state.delete`, mapped onto the four store operations in that order.
51pub fn default_ops() -> (StoreOp, StoreOp, StoreOp, StoreOp) {
52    let op = |tool: &str, args: &str| StoreOp {
53        tool: tool.into(),
54        args: Some(args.into()),
55        ok: None,
56        conflict: None,
57        value: None,
58        keys: None,
59    };
60    (
61        StoreOp {
62            ok: Some("result.structuredContent.ok".into()),
63            conflict: Some("result.structuredContent.latest".into()),
64            ..op(
65                "state.put",
66                r#"{"key": "{key}", "seq": {seq}, "state": {envelope}}"#,
67            )
68        },
69        StoreOp {
70            value: Some("result.structuredContent.state".into()),
71            ..op("state.get", r#"{"key": "{key}", "seq": {seq}}"#)
72        },
73        StoreOp {
74            keys: Some("result.structuredContent.keys".into()),
75            ..op("state.list", r#"{"prefix": "{prefix}"}"#)
76        },
77        op("state.delete", r#"{"key": "{key}"}"#),
78    )
79}
80
81pub struct McpStore {
82    client: Arc<dyn McpCall>,
83    put: StoreOp,
84    get: StoreOp,
85    list: Option<StoreOp>,
86    delete: Option<StoreOp>,
87    timeout: Duration,
88    /// Filled from the key at call time (the key layout carries them).
89    prefix_hint: std::sync::Mutex<Option<(String, String)>>,
90}
91
92impl McpStore {
93    pub fn new(client: Arc<dyn McpCall>, cfg: StoreMcp, timeout: Duration) -> McpStore {
94        let (dput, dget, dlist, ddelete) = default_ops();
95        McpStore {
96            client,
97            put: cfg.put.unwrap_or(dput),
98            get: cfg.get.unwrap_or(dget),
99            list: Some(cfg.list.unwrap_or(dlist)),
100            delete: Some(cfg.delete.unwrap_or(ddelete)),
101            timeout,
102            prefix_hint: std::sync::Mutex::new(None),
103        }
104    }
105
106    /// Restrict to the ops a server actually advertises (called by the runtime
107    /// after `tools/list`): absent `list`/`delete` become `Unsupported`.
108    pub fn with_advertised(mut self, tools: &[String]) -> McpStore {
109        let has = |t: &str| tools.iter().any(|x| x == t);
110        if !self.list.as_ref().is_some_and(|op| has(&op.tool)) {
111            self.list = None;
112        }
113        if !self.delete.as_ref().is_some_and(|op| has(&op.tool)) {
114            self.delete = None;
115        }
116        self
117    }
118
119    fn vars(&self, key: &str, seq: Option<u64>, envelope: Option<&Value>) -> Vars {
120        // key = <prefix>/<instance>/<kind>/<id>
121        let mut parts = key.splitn(4, '/');
122        let prefix = parts.next().unwrap_or("");
123        let instance = parts.next().unwrap_or("");
124        let kind = parts.next().unwrap_or("");
125        let id = parts.next().unwrap_or("");
126        if let Ok(mut h) = self.prefix_hint.lock() {
127            *h = Some((prefix.to_string(), instance.to_string()));
128        }
129        mapping::store_vars(key, seq, prefix, instance, envelope, kind, id)
130    }
131
132    fn call(
133        &self,
134        op: &StoreOp,
135        vars: &Vars,
136        key: &str,
137        seq: Option<u64>,
138    ) -> Result<Value, StoreError> {
139        let args = match &op.args {
140            Some(t) => mapping::render_json(t, vars)
141                .map_err(|e| StoreError::Mapping(format!("{}: {e}", op.tool)))?,
142            None => json!({ "key": key }),
143        };
144        let idem = match seq {
145            Some(s) => format!("{key}#{s}"),
146            None => key.to_string(),
147        };
148        let meta = json!({
149            "agent/idempotency_key": idem,
150            "agent/instance": vars.get("instance").cloned().unwrap_or(Value::Null),
151        });
152        let res = self
153            .client
154            .call(&op.tool, args, meta, self.timeout)
155            .map_err(|e| {
156                StoreError::Io(format!(
157                    "{} on '{}': {e}",
158                    op.tool,
159                    self.client.server_name()
160                ))
161            })?;
162        Ok(result_ctx(&res))
163    }
164}
165
166/// The extraction context for a tool result: `structuredContent` (or the text
167/// content parsed as JSON), `isError`, `text`, `content`.
168pub fn result_ctx(res: &CallToolResult) -> Value {
169    let text = res.text();
170    let structured = res
171        .structured_content
172        .clone()
173        .or_else(|| serde_json::from_str::<Value>(&text).ok())
174        .unwrap_or(Value::Null);
175    json!({
176        "result": {
177            "structuredContent": structured,
178            "isError": res.is_error(),
179            "text": text,
180            "content": res.content,
181        }
182    })
183}
184
185impl Store for McpStore {
186    fn put(&self, key: &str, seq: u64, envelope: &Value) -> Result<PutOutcome, StoreError> {
187        let vars = self.vars(key, Some(seq), Some(envelope));
188        let ctx = self.call(&self.put, &vars, key, Some(seq))?;
189        let is_error = ctx["result"]["isError"].as_bool().unwrap_or(false);
190        if let Some(okx) = &self.put.ok
191            && let Some(v) = mapping::extract(okx, &ctx).map_err(|e| StoreError::Mapping(e.0))?
192            && mapping::truthy(&v)
193        {
194            return Ok(PutOutcome::Ok);
195        }
196        if let Some(cx) = &self.put.conflict
197            && let Some(v) = mapping::extract(cx, &ctx).map_err(|e| StoreError::Mapping(e.0))?
198            && !v.is_null()
199        {
200            return Ok(PutOutcome::Conflict {
201                latest_seq: v.as_u64(),
202            });
203        }
204        if is_error {
205            return Err(StoreError::Io(format!(
206                "{} failed: {}",
207                self.put.tool, ctx["result"]["text"]
208            )));
209        }
210        // No `ok` predicate configured and no error ⇒ success.
211        if self.put.ok.is_none() {
212            return Ok(PutOutcome::Ok);
213        }
214        Err(StoreError::Io(format!(
215            "{} not acknowledged: {}",
216            self.put.tool, ctx["result"]["text"]
217        )))
218    }
219
220    fn get(&self, key: &str, seq: Option<u64>) -> Result<Option<Value>, StoreError> {
221        let vars = self.vars(key, seq, None);
222        let ctx = self.call(&self.get, &vars, key, None)?;
223        if ctx["result"]["isError"].as_bool().unwrap_or(false) {
224            // A tool-domain error on a read is "absent" (the checkpointer
225            // profile answers `no such key` that way); transport errors are Io.
226            return Ok(None);
227        }
228        let v = match &self.get.value {
229            Some(x) => mapping::extract(x, &ctx).map_err(|e| StoreError::Mapping(e.0))?,
230            None => Some(ctx["result"]["structuredContent"].clone()),
231        };
232        Ok(v.filter(|v| !v.is_null()))
233    }
234
235    fn list(&self, prefix: &str) -> Result<Vec<KeySeq>, StoreError> {
236        let Some(op) = &self.list else {
237            return Err(StoreError::Unsupported("list"));
238        };
239        let mut vars = self.vars(prefix, None, None);
240        vars.insert("prefix".into(), Value::String(prefix.to_string()));
241        let ctx = self.call(op, &vars, prefix, None)?;
242        if ctx["result"]["isError"].as_bool().unwrap_or(false) {
243            return Err(StoreError::Io(format!(
244                "{} failed: {}",
245                op.tool, ctx["result"]["text"]
246            )));
247        }
248        let keys = match &op.keys {
249            Some(x) => mapping::extract(x, &ctx).map_err(|e| StoreError::Mapping(e.0))?,
250            None => Some(ctx["result"]["structuredContent"]["keys"].clone()),
251        };
252        let mut out = Vec::new();
253        if let Some(Value::Array(items)) = keys {
254            for it in items {
255                match it {
256                    Value::String(k) => out.push(KeySeq { key: k, seq: None }),
257                    Value::Object(o) => {
258                        if let Some(k) = o.get("key").and_then(Value::as_str) {
259                            out.push(KeySeq {
260                                key: k.to_string(),
261                                seq: o.get("seq").and_then(Value::as_u64),
262                            });
263                        }
264                    }
265                    _ => {}
266                }
267            }
268        }
269        Ok(out)
270    }
271
272    fn delete(&self, key: &str) -> Result<(), StoreError> {
273        let Some(op) = &self.delete else {
274            return Err(StoreError::Unsupported("delete"));
275        };
276        let vars = self.vars(key, None, None);
277        let ctx = self.call(op, &vars, key, None)?;
278        if ctx["result"]["isError"].as_bool().unwrap_or(false) {
279            return Err(StoreError::Io(format!(
280                "{} failed: {}",
281                op.tool, ctx["result"]["text"]
282            )));
283        }
284        Ok(())
285    }
286
287    fn kind(&self) -> &'static str {
288        "mcp"
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295    use std::collections::BTreeMap;
296    use std::sync::Mutex;
297
298    /// A scripted checkpointer server speaking the default profile, with
299    /// history + CAS — the shape `mcp/mock_http.rs` implements over HTTP.
300    #[derive(Default)]
301    struct FakeServer {
302        data: Mutex<BTreeMap<String, BTreeMap<u64, Value>>>,
303        calls: Mutex<Vec<(String, Value, Value)>>,
304        fail: Mutex<bool>,
305    }
306
307    fn ok(v: Value) -> CallToolResult {
308        CallToolResult {
309            content: vec![json!({"type": "text", "text": v.to_string()})],
310            is_error: Some(false),
311            structured_content: None, // text-JSON only, like the mock
312        }
313    }
314    fn err(msg: &str) -> CallToolResult {
315        CallToolResult {
316            content: vec![json!({"type": "text", "text": msg})],
317            is_error: Some(true),
318            structured_content: None,
319        }
320    }
321
322    impl McpCall for FakeServer {
323        fn call(
324            &self,
325            tool: &str,
326            args: Value,
327            meta: Value,
328            _t: Duration,
329        ) -> Result<CallToolResult, String> {
330            self.calls
331                .lock()
332                .unwrap()
333                .push((tool.to_string(), args.clone(), meta));
334            if *self.fail.lock().unwrap() {
335                return Err("connection reset".into());
336            }
337            let key = args["key"].as_str().unwrap_or("").to_string();
338            let mut data = self.data.lock().unwrap();
339            Ok(match tool {
340                "state.put" => {
341                    let seq = args["seq"].as_u64().unwrap_or(0);
342                    let hist = data.entry(key).or_default();
343                    let latest = hist.keys().next_back().copied().unwrap_or(0);
344                    if seq <= latest {
345                        ok(json!({"ok": false, "latest": latest}))
346                    } else {
347                        hist.insert(seq, args["state"].clone());
348                        ok(json!({"ok": true, "seq": seq}))
349                    }
350                }
351                "state.get" => match data.get(&key) {
352                    None => err("no such key"),
353                    Some(h) => {
354                        let picked = match args["seq"].as_u64() {
355                            Some(s) => h.get(&s),
356                            None => h.values().next_back(),
357                        };
358                        match picked {
359                            Some(v) => ok(json!({"state": v})),
360                            None => err("no such seq"),
361                        }
362                    }
363                },
364                "state.list" => {
365                    let prefix = args["prefix"].as_str().unwrap_or("");
366                    let keys: Vec<Value> = data
367                        .iter()
368                        .filter(|(k, _)| k.starts_with(prefix))
369                        .map(|(k, h)| json!({"key": k, "seq": h.keys().next_back()}))
370                        .collect();
371                    ok(json!({"keys": keys}))
372                }
373                "state.delete" => {
374                    data.remove(&key);
375                    ok(json!({"ok": true}))
376                }
377                _ => err("unknown tool"),
378            })
379        }
380        fn server_name(&self) -> String {
381            "fake".into()
382        }
383    }
384
385    fn store(server: Arc<FakeServer>) -> McpStore {
386        McpStore::new(
387            server,
388            StoreMcp {
389                server: "fake".into(),
390                put: None,
391                get: None,
392                list: None,
393                delete: None,
394            },
395            Duration::from_secs(1),
396        )
397    }
398
399    #[test]
400    fn default_profile_round_trips_with_cas_and_meta() {
401        let srv = Arc::new(FakeServer::default());
402        let s = store(srv.clone());
403        let env = json!({"v": 2, "kind": "run", "id": "1", "seq": 1, "state": {"x": 1}});
404        assert_eq!(s.put("agentd/i/run/1", 1, &env).unwrap(), PutOutcome::Ok);
405        assert_eq!(
406            s.put("agentd/i/run/1", 1, &env).unwrap(),
407            PutOutcome::Conflict {
408                latest_seq: Some(1)
409            }
410        );
411        assert_eq!(s.get("agentd/i/run/1", None).unwrap(), Some(env.clone()));
412        assert_eq!(
413            s.get("agentd/i/run/nope", None).unwrap(),
414            None,
415            "tool error on read = absent"
416        );
417        let l = s.list("agentd/i/").unwrap();
418        assert_eq!(l.len(), 1);
419        assert_eq!(l[0].key, "agentd/i/run/1");
420        assert_eq!(l[0].seq, Some(1));
421        s.delete("agentd/i/run/1").unwrap();
422        assert_eq!(s.get("agentd/i/run/1", None).unwrap(), None);
423        // The idempotency meta rode along.
424        {
425            let calls = srv.calls.lock().unwrap();
426            let (tool, args, meta) = &calls[0];
427            assert_eq!(tool, "state.put");
428            assert_eq!(args["key"], json!("agentd/i/run/1"));
429            assert_eq!(args["seq"], json!(1));
430            assert_eq!(meta["agent/idempotency_key"], json!("agentd/i/run/1#1"));
431            assert_eq!(meta["agent/instance"], json!("i"));
432        }
433        // Transport failure is Io.
434        *srv.fail.lock().unwrap() = true;
435        assert!(matches!(
436            s.get("agentd/i/run/1", None),
437            Err(StoreError::Io(_))
438        ));
439    }
440
441    #[test]
442    fn custom_mapping_and_advertised_ops() {
443        // A server whose put is `kv.set {k, version, doc}` returning {stored: true}
444        // and whose get is `kv.fetch {k}` returning {doc}.
445        struct Kv(Mutex<BTreeMap<String, (u64, Value)>>);
446        impl McpCall for Kv {
447            fn call(
448                &self,
449                tool: &str,
450                args: Value,
451                _m: Value,
452                _t: Duration,
453            ) -> Result<CallToolResult, String> {
454                let mut d = self.0.lock().unwrap();
455                Ok(match tool {
456                    "kv.set" => {
457                        let k = args["k"].as_str().unwrap().to_string();
458                        let v = args["version"].as_u64().unwrap();
459                        if d.get(&k).is_some_and(|(cur, _)| *cur >= v) {
460                            ok(json!({"stored": false, "current": d[&k].0}))
461                        } else {
462                            d.insert(k, (v, args["doc"].clone()));
463                            ok(json!({"stored": true}))
464                        }
465                    }
466                    "kv.fetch" => match d.get(args["k"].as_str().unwrap()) {
467                        Some((_, doc)) => ok(json!({"doc": doc})),
468                        None => ok(json!({"doc": null})),
469                    },
470                    _ => err("nope"),
471                })
472            }
473            fn server_name(&self) -> String {
474                "kv".into()
475            }
476        }
477        let cfg = StoreMcp {
478            server: "kv".into(),
479            put: Some(StoreOp {
480                tool: "kv.set".into(),
481                args: Some(r#"{"k": "{key}", "version": {seq}, "doc": {envelope}}"#.into()),
482                ok: Some("result.structuredContent.stored".into()),
483                conflict: Some("result.structuredContent.current".into()),
484                value: None,
485                keys: None,
486            }),
487            get: Some(StoreOp {
488                tool: "kv.fetch".into(),
489                args: Some(r#"{"k": "{key}"}"#.into()),
490                ok: None,
491                conflict: None,
492                value: Some("result.structuredContent.doc".into()),
493                keys: None,
494            }),
495            list: None,
496            delete: None,
497        };
498        let s = McpStore::new(
499            Arc::new(Kv(Mutex::new(BTreeMap::new()))),
500            cfg,
501            Duration::from_secs(1),
502        )
503        .with_advertised(&["kv.set".into(), "kv.fetch".into()]);
504        assert_eq!(
505            s.put("p/i/k/1", 3, &json!({"a": 1})).unwrap(),
506            PutOutcome::Ok
507        );
508        assert_eq!(
509            s.put("p/i/k/1", 3, &json!({"a": 2})).unwrap(),
510            PutOutcome::Conflict {
511                latest_seq: Some(3)
512            }
513        );
514        assert_eq!(s.get("p/i/k/1", None).unwrap(), Some(json!({"a": 1})));
515        assert_eq!(s.get("p/i/k/2", None).unwrap(), None);
516        assert!(matches!(s.list("p/"), Err(StoreError::Unsupported("list"))));
517        assert!(matches!(
518            s.delete("p/i/k/1"),
519            Err(StoreError::Unsupported("delete"))
520        ));
521    }
522}
agentd/store/mcp.rs

agentd/store/
mcp.rs
