Skip to main content

agentd/store/
mod.rs

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