[F100] atomic Redis probe-failure counter via server-side Lua

increment_probe_failures / reset_probe_failures did a non-atomic
get_state (GET) -> mutate -> set (SET) in Python. A concurrent
activate() re-park writes a FRESH episode blob (probe_failures: 0 +
fresh activated_at / retry_after / affected_agents / kind); if the
stale increment's SET landed after the fresh activate's SET, the stale
blob overwrote the fresh episode metadata AND un-reset the counter
(clobbering the new episode).

Redis single-threads a Lua EVAL, so a server-side read-modify-write
is indivisible: activate's SET is serialized entirely before or after
the script, never interleaved between the script's GET and SET. The
two scripts mutate ONLY probe_failures, so every other episode field
survives the bump. activate stays a single atomic SET (a fresh episode
resetting the counter to 0 is correct semantics).

roboco/services/gateway/rate_limit_tracker.py

@@ -17,6 +17,46 @@ import redis.asyncio as redis
 
 from roboco.config import settings
 
+# Server-side atomic read-modify-write scripts. Redis single-threads a Lua
+# ``EVAL``, so the GET → decode → mutate → SET inside one script is indivisible:
+# a concurrent ``activate()`` (a re-park) is serialized entirely before or after
+# the script, never interleaved between the script's GET and SET. Without this
+# the counter update was a non-atomic ``get_state`` → mutate → ``set`` in Python,
+# so a re-park's fresh episode blob (``probe_failures: 0`` + fresh
+# ``activated_at`` / ``retry_after`` / ``affected_agents`` / ``kind``) could be
+# clobbered by the stale increment writing back the OLD blob — un-resetting the
+# counter and overwriting the fresh episode metadata. The scripts mutate ONLY
+# ``probe_failures`` so every other episode field survives the bump.
+_INCREMENT_PROBE_FAILURES = """\
+-- roboco:increment_probe_failures
+local key = KEYS[1]
+local raw = redis.call('GET', key)
+if not raw then
+  redis.call('SET', key, cjson.encode({probe_failures = 1}))
+  return 1
+end
+local state = cjson.decode(raw)
+local cur = state['probe_failures']
+if cur == nil then cur = 0 end
+local new_count = cur + 1
+state['probe_failures'] = new_count
+redis.call('SET', key, cjson.encode(state))
+return new_count
+"""
+
+_RESET_PROBE_FAILURES = """\
+-- roboco:reset_probe_failures
+local key = KEYS[1]
+local raw = redis.call('GET', key)
+if not raw then
+  redis.call('SET', key, cjson.encode({probe_failures = 0}))
+  return
+end
+local state = cjson.decode(raw)
+state['probe_failures'] = 0
+redis.call('SET', key, cjson.encode(state))
+"""
+
 
 class RateLimitStateTracker:
     """Track rate-limit state for a single AI provider in Redis.
@@ -120,20 +160,22 @@ class RateLimitStateTracker:
         probes have failed since the rate limit was activated.  The
         orchestrator uses this to decide whether to keep waiting or give
         up entirely.
+
+        Atomic: the read-modify-write runs server-side as a Lua ``EVAL`` so a
+        concurrent ``activate()`` re-park cannot interleave between the GET and
+        SET and clobber the fresh episode blob with a stale one.
         """
         r = await self._conn()
-        state = await self.get_state()
+        new_count = await r.eval(_INCREMENT_PROBE_FAILURES, 1, self._key())
-        new_count: int = state.get("probe_failures", 0) + 1
+        return int(new_count)
-        state["probe_failures"] = new_count
-        await r.set(self._key(), json.dumps(state))
-        return new_count
 
     async def reset_probe_failures(self) -> None:
-        """Reset the probe-failure counter to 0."""
+        """Reset the probe-failure counter to 0.
+
+        Atomic: server-side Lua ``EVAL`` (see ``increment_probe_failures``).
+        """
         r = await self._conn()
-        state = await self.get_state()
+        await r.eval(_RESET_PROBE_FAILURES, 1, self._key())
-        state["probe_failures"] = 0
-        await r.set(self._key(), json.dumps(state))
 
     # ------------------------------------------------------------------
     # Class-level helpers

tests/unit/services/test_rate_limit_tracker.py

@@ -9,6 +9,7 @@ visible to a fresh instance.
 
 from __future__ import annotations
 
+import json
 from typing import Any
 from unittest.mock import AsyncMock
 
@@ -22,7 +23,8 @@ from roboco.services.gateway.rate_limit_tracker import RateLimitStateTracker
 def _make_redis_mock(initial_store: dict[str, Any] | None = None) -> AsyncMock:
     """Build an async Redis mock backed by a plain dict.
 
-    The mock supports ``get``, ``set``, and ``delete`` with the same
+    The mock supports ``get``, ``set``, ``delete`` and ``eval`` (server-side
+    Lua for the tracker's atomic probe-failure scripts) with the same
     semantics as the real redis.asyncio.Redis client.
     """
     # Use the dict AS-IS (no copy) so that two mocks sharing the same
@@ -44,10 +46,36 @@ def _make_redis_mock(initial_store: dict[str, Any] | None = None) -> AsyncMock:
     async def _delete(key: str) -> int:
         return 1 if store.pop(key, None) is not None else 0
 
+    async def _eval(script: str, _numkeys: int, *keys_and_args: Any) -> Any:
+        # Mirror the tracker's two atomic Lua scripts (see rate_limit_tracker.py)
+        # so the counter update is observable in-process. Single-threaded tests
+        # get the same result production gets from Redis' single-threaded Lua.
+        key = keys_and_args[0]
+        raw = store.get(key)
+        text = raw.decode() if isinstance(raw, bytes) else (str(raw) if raw else None)
+        if "roboco:increment_probe_failures" in script:
+            if text is None:
+                state: dict[str, Any] = {"probe_failures": 1}
+            else:
+                state = json.loads(text)
+                state["probe_failures"] = state.get("probe_failures", 0) + 1
+            store[key] = json.dumps(state)
+            return state["probe_failures"]
+        if "roboco:reset_probe_failures" in script:
+            if text is None:
+                state = {"probe_failures": 0}
+            else:
+                state = json.loads(text)
+                state["probe_failures"] = 0
+            store[key] = json.dumps(state)
+            return None
+        raise AssertionError(f"unknown eval script: {script[:80]}")
+
     mock = AsyncMock()
     mock.get = AsyncMock(side_effect=_get)
     mock.set = AsyncMock(side_effect=_set)
     mock.delete = AsyncMock(side_effect=_delete)
+    mock.eval = AsyncMock(side_effect=_eval)
     # Stash the backing store so tests can inspect raw state
     mock._store = store
     return mock

tests/unit/services/test_rate_limit_tracker_atomic.py

@@ -0,0 +1,184 @@
+"""The probe-failure counter update must be a single atomic Redis op.
+
+``increment_probe_failures`` / ``reset_probe_failures`` used to do a non-atomic
+read-modify-write: ``get_state`` (GET) → mutate the dict → ``SET`` the whole
+blob back. A concurrent ``activate()`` (a re-park — grok 429 re-park, a 529
+overload re-park) writes a FRESH episode blob (``probe_failures: 0`` + fresh
+``activated_at`` / ``retry_after`` / ``affected_agents`` / ``kind``). If the
+stale ``increment``'s SET lands AFTER the fresh ``activate``'s SET, the stale
+blob overwrites the fresh episode metadata AND un-resets the counter (writes
+back the old ``probe_failures`` + old metadata) — clobbering the new episode.
+
+Redis single-threads a Lua ``EVAL``, so a server-side read-modify-write is
+indivisible: ``activate``'s ``SET`` is serialized entirely before or entirely
+after the script — never interleaved between the script's GET and SET. These
+tests pin the WIRING (the increment/reset go through ``eval``, a single atomic
+server-side call, NOT a separate ``get``+``set`` pair) and the field-preservation
+(the script decodes, mutates ONLY ``probe_failures``, re-encodes — every other
+episode field survives the bump).
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+from unittest.mock import AsyncMock
+
+import pytest
+from roboco.services.gateway.rate_limit_tracker import RateLimitStateTracker
+
+
+def _make_redis_mock(initial_store: dict[str, Any] | None = None) -> AsyncMock:
+    """Async Redis mock backed by a dict, with server-side ``eval`` for the
+    tracker's two Lua scripts.
+
+    The ``eval`` impl mirrors the Lua (GET → decode → mutate only
+    ``probe_failures`` → SET) so a single-threaded test observes the same result
+    production gets from Redis' atomic Lua execution. Real concurrency cannot be
+    simulated with a mock; the atomicity guarantee in production is Redis'
+    single-threaded Lua, which these tests pin by asserting the tracker routes
+    through ``eval`` (one atomic op) rather than a separate ``get``+``set``.
+    """
+    store: dict[str, Any] = initial_store if initial_store is not None else {}
+
+    async def _get(key: str) -> bytes | None:
+        val = store.get(key)
+        if val is None:
+            return None
+        if isinstance(val, bytes):
+            return val
+        return str(val).encode()
+
+    async def _set(key: str, value: Any) -> None:
+        store[key] = value
+
+    async def _delete(key: str) -> int:
+        return 1 if store.pop(key, None) is not None else 0
+
+    async def _eval(script: str, _numkeys: int, *keys_and_args: Any) -> Any:
+        key = keys_and_args[0]
+        raw = store.get(key)
+        text = raw.decode() if isinstance(raw, bytes) else (str(raw) if raw else None)
+        if "roboco:increment_probe_failures" in script:
+            if text is None:
+                state: dict[str, Any] = {"probe_failures": 1}
+            else:
+                state = json.loads(text)
+                state["probe_failures"] = state.get("probe_failures", 0) + 1
+            store[key] = json.dumps(state)
+            return state["probe_failures"]
+        if "roboco:reset_probe_failures" in script:
+            if text is None:
+                state = {"probe_failures": 0}
+            else:
+                state = json.loads(text)
+                state["probe_failures"] = 0
+            store[key] = json.dumps(state)
+            return None
+        raise AssertionError(f"unknown eval script: {script[:80]}")
+
+    mock = AsyncMock()
+    mock.get = AsyncMock(side_effect=_get)
+    mock.set = AsyncMock(side_effect=_set)
+    mock.delete = AsyncMock(side_effect=_delete)
+    mock.eval = AsyncMock(side_effect=_eval)
+    mock._store = store
+    return mock
+
+
+def _make_tracker(redis_mock: AsyncMock) -> RateLimitStateTracker:
+    tracker = RateLimitStateTracker(provider="anthropic", redis_url="redis://unused")
+    tracker._redis = redis_mock
+    return tracker
+
+
+@pytest.mark.asyncio
+async def test_increment_uses_atomic_eval_not_separate_get_set() -> None:
+    """The increment must route through ``eval`` (one atomic server-side op) and
+    must NOT issue a separate ``set`` for the read-modify-write — the separate
+    SET is exactly the non-atomic write a concurrent ``activate`` can clobber."""
+    mock = _make_redis_mock()
+    tracker = _make_tracker(mock)
+    await tracker.activate(retry_after=60.0, affected_agents=["be-dev-1"])
+    # activate issued the only legitimate SET; reset call counts so the increment
+    # path's commands are isolated.
+    mock.set.reset_mock()
+    mock.get.reset_mock()
+    mock.eval.reset_mock()
+
+    count = await tracker.increment_probe_failures()
+
+    assert count == 1
+    mock.eval.assert_awaited_once()
+    # The atomic op is server-side — the tracker must not issue its own SET
+    # (a separate SET is the non-atomic write a racing activate clobbers).
+    mock.set.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_reset_uses_atomic_eval_not_separate_get_set() -> None:
+    mock = _make_redis_mock()
+    tracker = _make_tracker(mock)
+    await tracker.activate()
+    await tracker.increment_probe_failures()
+    mock.set.reset_mock()
+    mock.get.reset_mock()
+    mock.eval.reset_mock()
+
+    await tracker.reset_probe_failures()
+
+    mock.eval.assert_awaited_once()
+    mock.set.assert_not_awaited()
+    assert (await tracker.get_state())["probe_failures"] == 0
+
+
+@pytest.mark.asyncio
+async def test_increment_preserves_episode_metadata() -> None:
+    """The atomic script decodes, mutates ONLY ``probe_failures``, and re-encodes
+    — every other episode field (rate_limited / kind / activated_at / retry_after
+    / affected_agents) survives the bump. This is the property a non-atomic
+    GET+SET that read a STALE blob would violate under a concurrent activate."""
+    retry_after = 120.0
+    bumps = 2
+    mock = _make_redis_mock()
+    tracker = _make_tracker(mock)
+    await tracker.activate(
+        retry_after=retry_after,
+        affected_agents=["be-dev-1", "fe-dev-1"],
+        kind="overloaded",
+    )
+    before = await tracker.get_state()
+
+    for _ in range(bumps):
+        await tracker.increment_probe_failures()
+
+    after = await tracker.get_state()
+    assert after["probe_failures"] == bumps
+    # Episode metadata untouched by the counter bump.
+    assert after["rate_limited"] is before["rate_limited"] is True
+    assert after["kind"] == "overloaded"
+    assert after["retry_after"] == retry_after
+    assert after["affected_agents"] == ["be-dev-1", "fe-dev-1"]
+    assert after["activated_at"] == before["activated_at"]
+
+
+@pytest.mark.asyncio
+async def test_increment_accumulates_across_calls() -> None:
+    mock = _make_redis_mock()
+    tracker = _make_tracker(mock)
+    await tracker.activate()
+    total = 4
+    counts = [await tracker.increment_probe_failures() for _ in range(total)]
+    assert counts == [1, 2, 3, total]
+    assert (await tracker.get_state())["probe_failures"] == total
+
+
+@pytest.mark.asyncio
+async def test_reset_zeroes_after_increments() -> None:
+    mock = _make_redis_mock()
+    tracker = _make_tracker(mock)
+    await tracker.activate()
+    for _ in range(3):
+        await tracker.increment_probe_failures()
+    await tracker.reset_probe_failures()
+    assert (await tracker.get_state())["probe_failures"] == 0