fix(gateway): move the budget counter at charge time, reconcile as a floor

apps/gateway/src/cache.rs

@@ -41,6 +41,18 @@ pub(crate) trait CacheStore: Send + Sync {
     /// Sets TTL only on first increment (new key / expired key).
     /// Returns the new count, or `None` on error (graceful fallback).
     async fn incr(&self, key: &str, ttl_secs: u64) -> Option<u64>;
+
+    /// Add `delta` to a signed counter, returning the new total.
+    ///
+    /// Distinct from [`incr`], which is a +1 rate-limit counter over `u64`.
+    /// Spend counters are signed nano-dollar totals and move by a per-request
+    /// amount, so they need their own primitive rather than a loop over `incr`.
+    ///
+    /// Atomic per key: the read-modify-write happens under the entry lock, so
+    /// concurrent charges against one budget accumulate instead of racing on a
+    /// read-then-set. That is the whole point of having it — a
+    /// `get_raw`/`set_raw` pair would lose every charge but the last.
+    async fn incr_by(&self, key: &str, delta: i64, ttl_secs: u64) -> Option<i64>;
 }
 
 /// Extension methods for typed get/set on any `CacheStore`.
@@ -151,6 +163,32 @@ impl CacheStore for InMemoryCacheStore {
         entry.data = count.to_string();
         Some(count)
     }
+
+    async fn incr_by(&self, key: &str, delta: i64, ttl_secs: u64) -> Option<i64> {
+        let now = Instant::now();
+        let ttl = Duration::from_secs(ttl_secs);
+
+        // `entry` holds the shard lock for the whole read-modify-write, so two
+        // concurrent charges on one budget cannot both read the same total.
+        let mut entry = self.map.entry(key.to_string()).or_insert(CachedEntry {
+            data: "0".to_string(),
+            expires_at: now + ttl,
+        });
+
+        if entry.expires_at <= now {
+            entry.data = "0".to_string();
+            entry.expires_at = now + ttl;
+        }
+
+        // An unparseable value is treated as 0 rather than propagating: this is
+        // a cache, and the durable BudgetSpend row is the floor that corrects it.
+        let total = entry.data.parse::<i64>().unwrap_or(0).saturating_add(delta);
+        entry.data = total.to_string();
+        // Refresh the window on write so an actively-charged period never
+        // expires mid-period.
+        entry.expires_at = now + ttl;
+        Some(total)
+    }
 }
 
 // ── Tests ────────────────────────────────────────────────────────────────
@@ -189,6 +227,49 @@ mod tests {
         assert!(result.is_none());
     }
 
+    #[tokio::test]
+    async fn incr_by_accumulates_concurrent_charges() {
+        // THE POINT OF incr_by. A get_raw/set_raw pair would let two concurrent
+        // charges both read 0 and both write their own delta, keeping only the
+        // last — which is exactly how spend used to slip past a cap.
+        let store = new_store();
+        let key = "budget:spent:sec:org:m:2026-08";
+
+        let tasks: Vec<_> = (0..64)
+            .map(|_| {
+                let store = Arc::clone(&store);
+                let key = key.to_string();
+                tokio::spawn(async move { store.incr_by(&key, 1_000, 60).await })
+            })
+            .collect();
+        for t in tasks {
+            t.await.unwrap();
+        }
+
+        let total: i64 = store.get_raw(key).await.unwrap().parse().unwrap();
+        assert_eq!(total, 64_000, "every concurrent charge must be counted");
+    }
+
+    #[tokio::test]
+    async fn incr_by_starts_from_zero_and_survives_a_set_floor() {
+        let store = new_store();
+        let key = "budget:spent:sec:org:total";
+
+        assert_eq!(store.incr_by(key, 250, 60).await, Some(250));
+        // A durable floor arriving from PostgreSQL, then more charges on top.
+        store.set_raw(key, "1000", 60).await;
+        assert_eq!(store.incr_by(key, 25, 60).await, Some(1025));
+    }
+
+    #[tokio::test]
+    async fn incr_by_treats_an_unparseable_value_as_zero() {
+        // The counter shares a namespace with JSON-valued caches; a stray value
+        // must not poison spend accounting. The durable row is the floor.
+        let store = new_store();
+        store.set_raw("budget:spent:x", "not-a-number", 60).await;
+        assert_eq!(store.incr_by("budget:spent:x", 7, 60).await, Some(7));
+    }
+
     #[tokio::test]
     async fn del_removes_entry() {
         let store = new_store();

apps/gateway/src/telemetry.rs

@@ -42,17 +42,51 @@ async fn record_spend(pool: &PgPool, cache: &dyn CacheStore, key: &SpendKey, del
     match crate::db::upsert_budget_spend(pool, secret_id, organization_id, period_key, delta).await
     {
         Ok(total) => {
+            // RECONCILE, don't overwrite. The cache was already incremented for
+            // this charge (`charge_cache` below), and may have been incremented
+            // again for a charge that lands in the NEXT batch. A blind
+            // `set_raw(total)` would roll those newer charges back and let spend
+            // through twice.
+            //
+            // The durable row is a FLOOR: raise the counter to it when the cache
+            // is behind (cold start, eviction, a lost increment), never lower it.
             let counter = crate::budget::counter_key(secret_id, organization_id, period_key);
-            cache
-                .set_raw(&counter, &total.to_string(), crate::budget::PERIOD_TTL)
-                .await;
+            let cached = cache
+                .get_raw(&counter)
+                .await
+                .and_then(|raw| raw.parse::<i64>().ok());
+            if cached.is_none_or(|c| c < total) {
+                cache
+                    .set_raw(&counter, &total.to_string(), crate::budget::PERIOD_TTL)
+                    .await;
+            }
         }
         Err(e) => {
             warn!(error = %e, secret_id = %secret_id, "budget: failed to record spend");
         }
     }
 }
 
+/// Apply a charge to the hot counter the moment it is drained from the channel,
+/// rather than waiting for the batch to reach PostgreSQL.
+///
+/// `pre_forward` reads this counter to decide whether to deny with 402. Before
+/// this, the counter only moved when the flush loop completed its upsert, so
+/// requests arriving between a charge and its flush all read the same stale
+/// total and were all admitted. The window was small — `collect_batch` returns
+/// as soon as one event arrives, so the 5s interval only applies when idle —
+/// but it was wide enough for concurrent requests against one budget.
+///
+/// Fail-open, like every other budget read: a cache miss returns None and the
+/// request proceeds. The durable floor in `record_spend` repairs the counter.
+async fn charge_cache(cache: &dyn CacheStore, key: &SpendKey, delta: i64) {
+    let (secret_id, organization_id, period_key) = key;
+    let counter = crate::budget::counter_key(secret_id, organization_id, period_key);
+    cache
+        .incr_by(&counter, delta, crate::budget::PERIOD_TTL)
+        .await;
+}
+
 async fn insert_batch(pool: &PgPool, events: &[RequestEvent]) -> Result<(), sqlx::Error> {
     let filtered: Vec<&RequestEvent> = events
         .iter()
@@ -181,7 +215,14 @@ async fn flush_loop(
             }
         }
 
-        // Persist the aggregated spend deltas (few keys per flush in practice).
+        // Move the hot counter FIRST, before the (slower) database round-trip:
+        // this is what `pre_forward` reads, so every microsecond it lags is a
+        // window in which a concurrent request sees a stale total.
+        for (key, delta) in &charges {
+            charge_cache(cache.as_ref(), key, *delta).await;
+        }
+
+        // Then persist, and reconcile the counter against the durable floor.
         for (key, delta) in &charges {
             record_spend(&pool, cache.as_ref(), key, *delta).await;
         }