agentd/store/
mcp.rs

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

agentd/store/
mcp.rs
