Skip to main content

agent_store/
generation.rs

1//! Monotonic generation counter.
2//!
3//! A [`Generation`] is a named, strictly-increasing `u64` kept in a small meta
4//! table. It is the causal clock modulex-mcp already reasons with (its report
5//! identity / `last_generation`), lifted into the shared substrate so every
6//! consumer stamps rows the same way.
7//!
8//! **Wall-clock time is never a coordination primitive.** A generation only
9//! ever moves forward; [`Generation::set`] refuses to move it backward.
10
11use crate::backend::{as_u64, Backend, Value};
12use crate::error::{Result, StoreError};
13
14const META_TABLE: &str = "_agent_store_meta";
15
16/// A named monotonic counter.
17#[derive(Clone, Debug)]
18pub struct Generation {
19    key: String,
20}
21
22impl Generation {
23    /// Name a counter. Many independent counters can coexist in one database.
24    pub fn new(key: impl Into<String>) -> Self {
25        Self { key: key.into() }
26    }
27
28    /// Create the backing meta table if it does not exist. Idempotent.
29    pub fn ensure_schema(db: &dyn Backend) -> Result<()> {
30        db.exec(
31            &format!(
32                "CREATE TABLE IF NOT EXISTS {META_TABLE} (\
33                 key TEXT PRIMARY KEY, ival INTEGER NOT NULL)"
34            ),
35            &[],
36        )?;
37        Ok(())
38    }
39
40    /// The current value (0 if the counter has never been bumped).
41    pub fn current(&self, db: &dyn Backend) -> Result<u64> {
42        let rows = db.query(
43            &format!("SELECT ival FROM {META_TABLE} WHERE key = ?"),
44            &[Value::Text(self.key.clone())],
45        )?;
46        match rows.first() {
47            Some(row) => as_u64(&row[0]),
48            None => Ok(0),
49        }
50    }
51
52    /// Atomically increment and return the new value. The first bump yields 1.
53    pub fn bump(&self, db: &dyn Backend) -> Result<u64> {
54        // One atomic upsert-with-RETURNING (supported by both SQLite ≥3.35 and
55        // Postgres). `ival = ival + 1` references the existing row's value.
56        let rows = db.query(
57            &format!(
58                "INSERT INTO {META_TABLE} (key, ival) VALUES (?, 1) \
59                 ON CONFLICT(key) DO UPDATE SET ival = ival + 1 RETURNING ival"
60            ),
61            &[Value::Text(self.key.clone())],
62        )?;
63        let row = rows
64            .first()
65            .ok_or_else(|| StoreError::Backend("bump: RETURNING produced no row".into()))?;
66        as_u64(&row[0])
67    }
68
69    /// Set the counter to an explicit value, refusing to move it backward.
70    pub fn set(&self, db: &dyn Backend, value: u64) -> Result<()> {
71        let current = self.current(db)?;
72        if value < current {
73            return Err(StoreError::NonMonotonicGeneration {
74                key: self.key.clone(),
75                current,
76                attempted: value,
77            });
78        }
79        db.exec(
80            &format!(
81                "INSERT INTO {META_TABLE} (key, ival) VALUES (?, ?) \
82                 ON CONFLICT(key) DO UPDATE SET ival = excluded.ival"
83            ),
84            &[Value::Text(self.key.clone()), Value::Int(value as i64)],
85        )?;
86        Ok(())
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use crate::backend::SqliteBackend;
94
95    fn db() -> SqliteBackend {
96        let db = SqliteBackend::in_memory().unwrap();
97        Generation::ensure_schema(&db).unwrap();
98        db
99    }
100
101    #[test]
102    fn starts_at_zero_then_increments() {
103        let db = db();
104        let g = Generation::new("report");
105        assert_eq!(g.current(&db).unwrap(), 0);
106        assert_eq!(g.bump(&db).unwrap(), 1);
107        assert_eq!(g.bump(&db).unwrap(), 2);
108        assert_eq!(g.current(&db).unwrap(), 2);
109    }
110
111    #[test]
112    fn counters_are_independent() {
113        let db = db();
114        let a = Generation::new("a");
115        let b = Generation::new("b");
116        a.bump(&db).unwrap();
117        a.bump(&db).unwrap();
118        b.bump(&db).unwrap();
119        assert_eq!(a.current(&db).unwrap(), 2);
120        assert_eq!(b.current(&db).unwrap(), 1);
121    }
122
123    #[test]
124    fn set_advances_but_never_rewinds() {
125        let db = db();
126        let g = Generation::new("report");
127        g.set(&db, 10).unwrap();
128        assert_eq!(g.current(&db).unwrap(), 10);
129
130        // Regression: a backward set must be refused (monotonic contract).
131        let err = g.set(&db, 3).unwrap_err();
132        assert!(matches!(
133            err,
134            StoreError::NonMonotonicGeneration {
135                current: 10,
136                attempted: 3,
137                ..
138            }
139        ));
140        assert_eq!(g.current(&db).unwrap(), 10);
141    }
142}