diff --git a/roboco/services/gateway/rate_limit_tracker.py b/roboco/services/gateway/rate_limit_tracker.py index 44384b92..51e16101 100644 --- a/roboco/services/gateway/rate_limit_tracker.py +++ b/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: int = state.get("probe_failures", 0) + 1 - state["probe_failures"] = new_count - await r.set(self._key(), json.dumps(state)) - return new_count + new_count = await r.eval(_INCREMENT_PROBE_FAILURES, 1, self._key()) + return int(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() - state["probe_failures"] = 0 - await r.set(self._key(), json.dumps(state)) + await r.eval(_RESET_PROBE_FAILURES, 1, self._key()) # ------------------------------------------------------------------ # Class-level helpers diff --git a/tests/unit/services/test_rate_limit_tracker.py b/tests/unit/services/test_rate_limit_tracker.py index 330b0bb8..210ac862 100644 --- a/tests/unit/services/test_rate_limit_tracker.py +++ b/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 diff --git a/tests/unit/services/test_rate_limit_tracker_atomic.py b/tests/unit/services/test_rate_limit_tracker_atomic.py new file mode 100644 index 00000000..753d149b --- /dev/null +++ b/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