Skip to main content

agentd/store/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2//! The **state store** contract and adapters (RFC 0025 §2, §4).
3//!
4//! agentd's durability rests on four operations against a REMOTE store —
5//! `put(key, seq, envelope) / get(key[, seq]) / list(prefix) / delete(key)` —
6//! implemented by an adapter chosen in `store.kind`: [`mcp`] (any MCP server's
7//! tools, mapped), [`http`] (plain HTTP), or [`memory`] (in-process; tests and
8//! dev). `put` is a **compare-and-set on `seq`**: the stored seq must be lower,
9//! else `Conflict` — the split-brain guard every caller treats as fatal.
10//! agentd links no database client and defines no schema beyond the
11//! [`Envelope`].
12
13pub mod http;
14pub mod mapping;
15pub mod mcp;
16pub mod memory;
17
18use serde::{Deserialize, Serialize};
19use serde_json::Value;
20use std::sync::Arc;
21use std::time::Duration;
22
23/// The envelope major this build writes and accepts (RFC 0025 §3.2).
24pub const ENVELOPE_VERSION: u32 = 2;
25
26/// A versioned store record: `state` is the kind-specific payload; a tombstone
27/// carries `state: null`.
28#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
29pub struct Envelope {
30    pub v: u32,
31    pub kind: String,
32    pub id: String,
33    pub seq: u64,
34    /// Unix ms.
35    pub ts: u64,
36    pub instance: String,
37    #[serde(default, skip_serializing_if = "Option::is_none")]
38    pub hash: Option<String>,
39    pub state: Value,
40}
41
42impl Envelope {
43    pub fn new(
44        kind: &str,
45        id: &str,
46        seq: u64,
47        instance: &str,
48        hash: Option<String>,
49        state: Value,
50    ) -> Envelope {
51        Envelope {
52            v: ENVELOPE_VERSION,
53            kind: kind.to_string(),
54            id: id.to_string(),
55            seq,
56            ts: now_ms(),
57            instance: instance.to_string(),
58            hash,
59            state,
60        }
61    }
62
63    pub fn is_tombstone(&self) -> bool {
64        self.state.is_null()
65    }
66
67    pub fn to_value(&self) -> Value {
68        serde_json::to_value(self).unwrap_or(Value::Null)
69    }
70
71    /// Parse a stored value; refuses an unknown envelope major.
72    pub fn from_value(v: Value) -> Result<Envelope, StoreError> {
73        let env: Envelope = serde_json::from_value(v)
74            .map_err(|e| StoreError::Corrupt(format!("envelope does not parse: {e}")))?;
75        if env.v != ENVELOPE_VERSION {
76            return Err(StoreError::Corrupt(format!(
77                "envelope version {} is not supported (this build writes {})",
78                env.v, ENVELOPE_VERSION
79            )));
80        }
81        Ok(env)
82    }
83}
84
85/// `<prefix>/<instance>/<kind>/<id>` (RFC 0025 §3.1).
86pub fn key(prefix: &str, instance: &str, kind: &str, id: &str) -> String {
87    format!("{prefix}/{instance}/{kind}/{id}")
88}
89
90/// Split a key produced by [`key`] back into `(kind, id)` for the given
91/// prefix/instance; `None` for a foreign key.
92pub fn parse_key<'a>(prefix: &str, instance: &str, k: &'a str) -> Option<(&'a str, &'a str)> {
93    let rest = k.strip_prefix(&format!("{prefix}/{instance}/"))?;
94    let (kind, id) = rest.split_once('/')?;
95    Some((kind, id))
96}
97
98/// The outcome of a `put`.
99#[derive(Debug, Clone, PartialEq, Eq)]
100pub enum PutOutcome {
101    Ok,
102    /// Another writer owns the key (a stored seq ≥ ours). Fatal for the writer.
103    Conflict {
104        latest_seq: Option<u64>,
105    },
106}
107
108/// A `list` entry.
109#[derive(Debug, Clone, PartialEq, Eq)]
110pub struct KeySeq {
111    pub key: String,
112    pub seq: Option<u64>,
113}
114
115#[derive(Debug, Clone, PartialEq, Eq)]
116pub enum StoreError {
117    /// Transport / server failure (retryable at the caller's discretion).
118    Io(String),
119    /// The adapter does not implement this optional operation.
120    Unsupported(&'static str),
121    /// A mapping template or extraction failed (a config problem).
122    Mapping(String),
123    /// A stored record is unreadable.
124    Corrupt(String),
125    /// A `put` conflict surfaced as an error by a caller that treats it as fatal.
126    Conflict(String),
127}
128
129impl std::fmt::Display for StoreError {
130    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
131        match self {
132            StoreError::Io(m) => write!(f, "store i/o: {m}"),
133            StoreError::Unsupported(op) => {
134                write!(f, "store: {op} is not supported by this adapter")
135            }
136            StoreError::Mapping(m) => write!(f, "store mapping: {m}"),
137            StoreError::Corrupt(m) => write!(f, "store record: {m}"),
138            StoreError::Conflict(m) => write!(f, "store conflict: {m}"),
139        }
140    }
141}
142
143impl std::error::Error for StoreError {}
144
145/// The four-operation contract (RFC 0025 §2). Implementations are `Send +
146/// Sync` so the runtime's executor pool can call them; every operation is
147/// bounded by the adapter's timeout.
148pub trait Store: Send + Sync {
149    /// Compare-and-set write: `seq` must exceed the stored one.
150    fn put(&self, key: &str, seq: u64, envelope: &Value) -> Result<PutOutcome, StoreError>;
151    /// The latest record (or the pinned `seq` if the store keeps history).
152    fn get(&self, key: &str, seq: Option<u64>) -> Result<Option<Value>, StoreError>;
153    /// Keys under `prefix` (optional; `Unsupported` is a legal answer).
154    fn list(&self, prefix: &str) -> Result<Vec<KeySeq>, StoreError>;
155    /// Remove a key (optional; `Unsupported` ⇒ callers tombstone via `put`).
156    fn delete(&self, key: &str) -> Result<(), StoreError>;
157    /// The adapter kind (`mcp` / `http` / `memory`), for status and metrics.
158    fn kind(&self) -> &'static str;
159}
160
161/// A shared store handle.
162pub type SharedStore = Arc<dyn Store>;
163
164/// The store timeout class: the management timeout (RFC 0016 §10) unless the
165/// settings say otherwise.
166pub fn default_timeout() -> Duration {
167    crate::obs::health::management_timeout()
168}
169
170/// Bounded retry on `Io` errors (never on `Conflict`/`Mapping`/`Corrupt`).
171pub fn with_retry<T>(
172    mut op: impl FnMut() -> Result<T, StoreError>,
173    attempts: u32,
174) -> Result<T, StoreError> {
175    let mut last = None;
176    for n in 0..attempts.max(1) {
177        match op() {
178            Err(StoreError::Io(m)) => {
179                last = Some(StoreError::Io(m));
180                if n + 1 < attempts {
181                    std::thread::sleep(Duration::from_millis(50 * (n as u64 + 1)));
182                }
183            }
184            other => return other,
185        }
186    }
187    Err(last.unwrap_or(StoreError::Io("no attempts".into())))
188}
189
190pub(crate) fn now_ms() -> u64 {
191    std::time::SystemTime::now()
192        .duration_since(std::time::UNIX_EPOCH)
193        .map(|d| d.as_millis() as u64)
194        .unwrap_or(0)
195}
196
197/// Build the store an instance is configured with (RFC 0030 §3.5). `servers`
198/// resolves the `mcp` adapter's coordination server by name; `kind: none`
199/// yields no store (the caller decides whether that is allowed).
200pub fn open(
201    settings: &crate::config::v2::Store,
202    servers: &dyn Fn(&str) -> Option<Arc<dyn mcp::McpCall>>,
203) -> Result<Option<SharedStore>, StoreError> {
204    use crate::config::v2::StoreKind;
205    let timeout = settings
206        .timeout
207        .map(|d| d.0)
208        .unwrap_or_else(default_timeout);
209    match settings.kind {
210        StoreKind::None => Ok(None),
211        StoreKind::Memory => Ok(Some(Arc::new(memory::MemoryStore::new()))),
212        StoreKind::Mcp => {
213            let cfg = settings.mcp.as_ref().ok_or_else(|| {
214                StoreError::Mapping("store.kind is mcp but store.mcp is not set".into())
215            })?;
216            let client = servers(&cfg.server).ok_or_else(|| {
217                StoreError::Mapping(format!(
218                    "store.mcp.server '{}' is not a connected MCP server",
219                    cfg.server
220                ))
221            })?;
222            Ok(Some(Arc::new(mcp::McpStore::new(
223                client,
224                cfg.clone(),
225                timeout,
226            ))))
227        }
228        StoreKind::Http => {
229            let cfg = settings.http.as_ref().ok_or_else(|| {
230                StoreError::Mapping("store.kind is http but store.http is not set".into())
231            })?;
232            Ok(Some(Arc::new(http::HttpStore::new(cfg.clone(), timeout)?)))
233        }
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use serde_json::json;
241
242    #[test]
243    fn envelope_round_trips_and_refuses_unknown_major() {
244        let e = Envelope::new("run", "01J", 3, "inst", Some("abc".into()), json!({"a": 1}));
245        let v = e.to_value();
246        assert_eq!(v["v"], json!(2));
247        assert_eq!(v["seq"], json!(3));
248        let back = Envelope::from_value(v).unwrap();
249        assert_eq!(back, e);
250        let mut bad = e.to_value();
251        bad["v"] = json!(9);
252        assert!(matches!(
253            Envelope::from_value(bad),
254            Err(StoreError::Corrupt(_))
255        ));
256        assert!(!e.is_tombstone());
257        assert!(Envelope::new("run", "x", 1, "i", None, Value::Null).is_tombstone());
258    }
259
260    #[test]
261    fn keys_compose_and_parse() {
262        let k = key("agentd", "inst-0", "run", "01J");
263        assert_eq!(k, "agentd/inst-0/run/01J");
264        assert_eq!(parse_key("agentd", "inst-0", &k), Some(("run", "01J")));
265        assert_eq!(parse_key("agentd", "other", &k), None);
266        // Ids may contain slashes (a2a contextId…): kind is the first segment.
267        assert_eq!(
268            parse_key("agentd", "i", "agentd/i/task/a/b"),
269            Some(("task", "a/b"))
270        );
271    }
272
273    #[test]
274    fn retry_only_on_io() {
275        let mut n = 0;
276        let r: Result<(), StoreError> = with_retry(
277            || {
278                n += 1;
279                Err(StoreError::Io("down".into()))
280            },
281            3,
282        );
283        assert!(matches!(r, Err(StoreError::Io(_))));
284        assert_eq!(n, 3);
285        let mut m = 0;
286        let r: Result<(), StoreError> = with_retry(
287            || {
288                m += 1;
289                Err(StoreError::Mapping("bad".into()))
290            },
291            3,
292        );
293        assert!(matches!(r, Err(StoreError::Mapping(_))));
294        assert_eq!(m, 1);
295    }
296}