Trait StateStore 

pub trait StateStore: Send + Sync {
    // Required methods
    fn add(&self, key: &str, delta: u64) -> Result<u64, StateError>;
    fn get(&self, key: &str) -> Result<u64, StateError>;
    fn try_spend_many(
        &self,
        spends: &[Spend],
    ) -> Result<TrySpendOutcome, StateError>;
    fn remove(&self, key: &str);

    // Provided methods
    fn fetch_max(&self, _key: &str, _value: u64) -> Result<u64, StateError> { ... }
    fn counter_ttl_secs(&self) -> Option<u64> { ... }
    fn try_spend(
        &self,
        key: &str,
        amount: u64,
        limit: u64,
    ) -> Result<bool, StateError> { ... }
    fn refund(&self, key: &str, amount: u64) { ... }
    fn refund_many(&self, refunds: &[Refund]) { ... }
    fn remove_prefix(&self, prefix: &str) { ... }
}

Atomic counter operations. Every mutation is atomic with respect to concurrent callers; try_spend is a single check-and-spend (never a read-then-write).

Counter lifetime

Backends with native expiry apply a TTL to every counter they touch: RedisStore refreshes a 24 h TTL (BUDGET_COUNTER_TTL_SECS in redis_store.rs) on each spend/add/refund, so a session idle past that window has its budget counters silently reset. InMemoryStore never expires. Callers must bound session lifetimes to the backend TTL (or clean up explicitly via remove/remove_prefix) so the two backends stay behaviorally aligned across dev and prod.

Required Methods

fn add(&self, key: &str, delta: u64) -> Result<u64, StateError>

Add delta to key, returning the new value.

fn get(&self, key: &str) -> Result<u64, StateError>

Read the current value of key (0 if absent).

fn try_spend_many( &self, spends: &[Spend], ) -> Result<TrySpendOutcome, StateError>

Atomically validate and commit every spend, or commit none.

Every Spend in spends must carry a distinct key; two entries for the same key would each observe the pre-commit value in the check phase and pass their independent limit checks, then the commit phase would sum them and blow through the cap. Duplicate keys return StateError::Backend (not Overflow), matching the API-misuse class.

On success returns TrySpendOutcome::Committed carrying the post-commit min-headroom across the spent keys — computed server-side in the SAME atomic section that committed the spend, so it cannot race with a concurrent spend/refund/ remove_prefix on the same keys. Callers that previously followed a successful spend with a get-per-dimension remaining_min computation to build BudgetDecision::Allowed { remaining } should use the returned value directly (R66 F3).

On refusal returns TrySpendOutcome::Refused with the index of the offending spend.

fn remove(&self, key: &str)

Remove a key (session cleanup).

Provided Methods

fn fetch_max(&self, _key: &str, _value: u64) -> Result<u64, StateError>

Atomically retain the greater of the stored value and value, returning the resulting maximum. Used for monotonic revocation cutoffs.

fn counter_ttl_secs(&self) -> Option<u64>

Backend counter TTL in seconds, if the backend expires counters natively. None means counters live for the process lifetime (in-memory). Exposing it makes the prod/dev divergence a VALUE callers can inspect (startup logs it; ops docs cite it) instead of a doc comment — a session active longer than this window has its budget counters silently reset against the expiring backend only.

fn try_spend( &self, key: &str, amount: u64, limit: u64, ) -> Result<bool, StateError>

Atomically spend amount from the remaining budget limit - spent(key). Returns Ok(true) and records the spend if the full amount fits, Ok(false) (recording nothing) otherwise.

Default impl routes through Self::try_spend_many; backends override only if they can express single-key check-and-spend more cheaply than the multi-key path.

fn refund(&self, key: &str, amount: u64)

Return a previously-spent amount to a counter. Saturating: if the stored value is below amount (e.g. a concurrent remove_prefix cleared it first, or another refund already covered part of the debt) the counter clamps at 0 rather than underflowing into “negative budget”. Backends must never propagate a refund error to the caller — the refund is best-effort compensation on a lost-race path where the primary verdict has already been decided. Default remove_prefix semantics apply: backends with native TTL expiry may fold this into their own cleanup if they prefer.

Introduced to close a concurrent-MCP budget double-spend. When two identical MCP requests race and one loses the atomic execution.claim(), the sandbox-gate spend is refunded so the budget counters reflect only the admitted work.

fn refund_many(&self, refunds: &[Refund])

Atomic multi-key refund. Twin to Self::try_spend_many.

Best-effort by contract (mirroring refund): backends swallow errors internally so a compensation failure never turns into a 5xx on the caller path. The atomicity guarantee is: every backend implementation MUST either apply every listed refund or none of them, so a concurrent try_tool_call / try_spend_many cannot observe partially- refunded state (fresh total_calls headroom but still- charged per-tool, etc.). The trait’s default falls back to per-key refund calls — safe for the single-key case and for backends that inherit the default, but atomicity of downstream callers relies on backends overriding this with a single-transaction implementation. InMemoryStore takes its transaction_lock once around the whole batch; RedisStore sends a single Lua script keyed on every counter (all sharing the session’s Redis Cluster hash tag, so cross-slot fan-out is impossible).

Duplicate keys in a single call are applied step-wise per entry with the same saturate-at-zero clamp as the single-key refund; the two backends match. In practice production callers (currently only refund_tool_call) never produce duplicate keys, so the step-wise-vs-summed distinction is only relevant to a bespoke caller that intentionally batches multiple entries against one counter.

fn remove_prefix(&self, prefix: &str)

Remove every key beginning with prefix (whole-session cleanup at finalization). Every backend must implement this: native expiry (e.g. Redis TTLs) is NOT a substitute, because a finalized session id can be recycled into a fresh session within the TTL window and would inherit the prior incarnation’s counters. In-process backends additionally need it or session-keyed counters accumulate for the process lifetime.

CLUSTER-MODE HAZARD FOR FUTURE CALLERS. The Redis Cluster implementation of remove_prefix routes SCAN and DEL to a SINGLE hash slot — the one derived from the prefix. This works only when every key under the prefix shares that slot, i.e. the prefix contains a Redis Cluster hash tag ({...}). Today ActionBudget::session_prefix wraps the digest in a {hash-tag} (see budget.rs), so all its keys land in the same slot and this is safe. A future caller that does NOT hash-tag its prefix would silently leave keys behind in other slots. Cross-slot SCAN is not supported by Redis Cluster — there is no correct implementation for a non-tagged prefix, so this remains a caller-side invariant.

Implementors

impl StateStore for RedisStore

impl StateStore for InMemoryStore