av_state/
lib.rs

1//! Session/agent state: atomic counters, token-velocity windows, rate limits,
2//! and action budgets (brief Modules B/D and the §8 "In-Memory State" layer).
3//!
4//! The core abstraction is [`StateStore`]: check-and-spend operations that are
5//! atomic even under concurrent tool calls from parallel sub-agents
6//! (silent-error class D13.9 — a budget must never over-spend by racing).
7//! `InMemoryStore` is the single-node reference; a Redis Cluster backend
8//! compiles behind the `redis` feature for multi-node deployments and is
9//! contract-tested against a live server when `AV_REDIS_URL` is set.
10//!
11//! All arithmetic here is checked — budget/monetary code must fail loudly,
12//! never wrap.
13
14pub mod budget;
15pub mod store;
16pub mod velocity;
17
18pub use budget::{ActionBudget, BudgetDecision, BudgetSpec};
19pub use store::{InMemoryStore, Refund, Spend, StateError, StateStore, TrySpendOutcome};
20pub use velocity::TokenVelocity;
21
22#[cfg(feature = "redis")]
23pub mod redis_store;
24
25/// Shared backend-agnostic `StateStore` contract.
26///
27/// EVOLUTION.md promises "new connectors must satisfy the same contract
28/// tests" — previously the in-memory and Redis suites asserted different
29/// things (none of the concurrency-adjacent semantics ran against Redis
30/// at all), so the implementations drifted (TTL, `remove_prefix`).
31/// Every backend's test suite MUST call this one function; backend-
32/// specific behavior (cluster slots, TTL) stays in the backend's own
33/// suite ON TOP of this contract, never instead of it.
34///
35/// `hash_tag` is interpolated into every key inside `{…}` so Redis
36/// Cluster callers keep all contract keys in one slot; in-memory
37/// callers can pass anything unique.
38#[doc(hidden)]
39#[allow(clippy::unwrap_used, clippy::panic, clippy::missing_panics_doc)]
40pub fn state_store_contract(store: &dyn StateStore, hash_tag: &str) {
41    let key = |name: &str| format!("contract:{{{hash_tag}}}:{name}");
42
43    // -- add / get roundtrip; absent key reads 0.
44    let counter = key("counter");
45    assert_eq!(store.get(&counter).unwrap(), 0, "absent key must read 0");
46    assert_eq!(store.add(&counter, 5).unwrap(), 5);
47    assert_eq!(store.add(&counter, 2).unwrap(), 7);
48    assert_eq!(store.get(&counter).unwrap(), 7);
49
50    // -- try_spend: exact-fit passes, one-over refuses, nothing partial.
51    let spend = key("spend");
52    assert!(store.try_spend(&spend, 6, 10).unwrap());
53    assert!(store.try_spend(&spend, 4, 10).unwrap(), "exact fit must pass");
54    assert!(!store.try_spend(&spend, 1, 10).unwrap(), "over-cap must refuse");
55    assert_eq!(
56        store.get(&spend).unwrap(),
57        10,
58        "refused spend must record nothing"
59    );
60
61    // -- try_spend_many: all-or-nothing across dimensions.
62    let dim_a = key("dim-a");
63    let dim_b = key("dim-b");
64    assert_eq!(
65        store
66            .try_spend_many(&[
67                Spend {
68                    key: dim_a.clone(),
69                    amount: 3,
70                    limit: 10
71                },
72                Spend {
73                    key: dim_b.clone(),
74                    amount: 20,
75                    limit: 10
76                },
77            ])
78            .unwrap(),
79        TrySpendOutcome::Refused { index: 1 },
80        "second dimension over-cap must be reported by index"
81    );
82    assert_eq!(
83        store.get(&dim_a).unwrap(),
84        0,
85        "failed multi-spend must commit NOTHING"
86    );
87    // dim_a: 3/10 -> remaining 7; dim_b: 4/10 -> remaining 6. Min = 6.
88    assert_eq!(
89        store
90            .try_spend_many(&[
91                Spend {
92                    key: dim_a.clone(),
93                    amount: 3,
94                    limit: 10
95                },
96                Spend {
97                    key: dim_b.clone(),
98                    amount: 4,
99                    limit: 10
100                },
101            ])
102            .unwrap(),
103        TrySpendOutcome::Committed {
104            post_commit_min_remaining: 6,
105        },
106        "R66 F3: try_spend_many must return post-commit min-headroom \
107         computed inside the atomic section; a subsequent get-based \
108         computation would race with concurrent remove_prefix/spend/refund"
109    );
110    assert_eq!(store.get(&dim_a).unwrap(), 3);
111    assert_eq!(store.get(&dim_b).unwrap(), 4);
112
113    // -- R69 F4 boundary coverage: empty spends slice and exact-fit
114    // to a fresh key. Both must match across backends; the
115    // RedisStore Lua has a `min_remaining == -1` empty-KEYS branch
116    // (redis_store.rs:73-76) that a client-side short-circuit
117    // covers today, so an incidental change to the short-circuit
118    // condition that let the empty branch reach Rust decode would
119    // silently break callers passing `try_spend_many(&[])`.
120    assert_eq!(
121        store.try_spend_many(&[]).unwrap(),
122        TrySpendOutcome::Committed {
123            post_commit_min_remaining: u64::MAX,
124        },
125        "empty spends slice must commit with u64::MAX headroom sentinel"
126    );
127    let boundary = key("boundary-exact-fit");
128    assert_eq!(
129        store
130            .try_spend_many(&[Spend {
131                key: boundary.clone(),
132                amount: 10,
133                limit: 10,
134            }])
135            .unwrap(),
136        TrySpendOutcome::Committed {
137            post_commit_min_remaining: 0,
138        },
139        "exact-fit spend must commit with 0 headroom, not refuse"
140    );
141    assert_eq!(store.get(&boundary).unwrap(), 10);
142
143    // -- try_spend_many refuses duplicate keys (API-misuse class).
144    assert!(
145        store
146            .try_spend_many(&[
147                Spend {
148                    key: dim_a.clone(),
149                    amount: 1,
150                    limit: 10
151                },
152                Spend {
153                    key: dim_a.clone(),
154                    amount: 1,
155                    limit: 10
156                },
157            ])
158            .is_err(),
159        "duplicate keys in one multi-spend must be refused, not double-counted"
160    );
161
162    // -- refund: subtracts, saturates at 0, never resurrects a removed key.
163    let refund = key("refund");
164    assert!(store.try_spend(&refund, 7, 10).unwrap());
165    store.refund(&refund, 3);
166    assert_eq!(store.get(&refund).unwrap(), 4, "refund must subtract");
167    store.refund(&refund, 100);
168    assert_eq!(
169        store.get(&refund).unwrap(),
170        0,
171        "refund must clamp at 0, never underflow"
172    );
173    store.remove(&refund);
174    store.refund(&refund, 5);
175    assert_eq!(
176        store.get(&refund).unwrap(),
177        0,
178        "refund after remove must not resurrect the counter"
179    );
180
181    // -- refund_many: same semantics as per-key refund, one atomic
182    // transaction. Every backend MUST subtract each amount, saturate
183    // at 0, and never resurrect a removed cell.
184    let m_a = key("refund_many_a");
185    let m_b = key("refund_many_b");
186    let m_c = key("refund_many_c");
187    assert!(store.try_spend(&m_a, 7, 10).unwrap());
188    assert!(store.try_spend(&m_b, 5, 10).unwrap());
189    // c intentionally never spent — refund_many on a fresh key must
190    // NOT resurrect it (same "never resurrect" invariant as `refund`).
191    store.refund_many(&[
192        Refund {
193            key: m_a.clone(),
194            amount: 3,
195        },
196        Refund {
197            key: m_b.clone(),
198            amount: 100,
199        },
200        Refund {
201            key: m_c.clone(),
202            amount: 1,
203        },
204    ]);
205    assert_eq!(
206        store.get(&m_a).unwrap(),
207        4,
208        "refund_many must subtract every listed amount"
209    );
210    assert_eq!(
211        store.get(&m_b).unwrap(),
212        0,
213        "refund_many must clamp at 0 like per-key refund"
214    );
215    assert_eq!(
216        store.get(&m_c).unwrap(),
217        0,
218        "refund_many on a fresh key must not resurrect it"
219    );
220    // Empty batch is a no-op.
221    store.refund_many(&[]);
222    assert_eq!(store.get(&m_a).unwrap(), 4);
223
224    // -- remove / remove_prefix: whole-session cleanup.
225    store.remove(&counter);
226    assert_eq!(store.get(&counter).unwrap(), 0);
227    store.remove_prefix(&key(""));
228    assert_eq!(
229        store.get(&spend).unwrap(),
230        0,
231        "remove_prefix must clear every key under the prefix (native TTL is NOT a substitute: a recycled session id would inherit the prior incarnation's counters)"
232    );
233    assert_eq!(store.get(&dim_a).unwrap(), 0);
234    assert_eq!(store.get(&dim_b).unwrap(), 0);
235}
av_state/lib.rs

av_state/
lib.rs
