From 0b6d44890f49a1ec615d79f53cde4e9c8a1ae80c Mon Sep 17 00:00:00 2001 From: isharak7m <192635824+isharak7m@users.noreply.github.com> Date: Sun, 13 Sep 2026 18:46:27 +0530 Subject: [PATCH] test: add regression for token cache atomic swap race condition MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Exercises the concurrent reader/writer scenario that the atomic swap fix in app.py addresses. Uses a _SharedCache helper that mirrors the module-level _token_cache global — both reader and writer access the same .current reference, so the GIL-atomic swap is properly tested. 6 tests: swap correctness, concurrent readers (4 threads x 100 refreshes, zero empty reads), app.state sync, multiple prefixes, empty DB, and concurrent refresh from 4 threads. --- tests/test_token_cache_atomic_swap.py | 176 ++++++++++++++++++++++++++ 1 file changed, 176 insertions(+) create mode 100644 tests/test_token_cache_atomic_swap.py diff --git a/tests/test_token_cache_atomic_swap.py b/tests/test_token_cache_atomic_swap.py new file mode 100644 index 000000000..1e228a0fd --- /dev/null +++ b/tests/test_token_cache_atomic_swap.py @@ -0,0 +1,176 @@ +"""Token cache atomic swap prevents race condition during refresh. + +_refresh_token_cache() in app.py previously mutated the shared _token_cache +dict in two steps: .clear() then .update(). Between those calls the dict was +empty, so any concurrent reader (line 428) saw zero candidates and returned +401 for a valid token. + +The fix replaces the two-step mutation with an atomic reference swap +(_token_cache = dict(new_map)). Python's GIL makes the assignment atomic, +so readers always see either the old fully-populated dict or the new one. +""" +import threading +import time +from collections import defaultdict +from types import SimpleNamespace + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_token_row(prefix, token_id="t1", token_hash="h1", owner="admin", scopes="chat"): + return SimpleNamespace( + token_prefix=prefix, + id=token_id, + token_hash=token_hash, + owner=owner, + scopes=scopes, + is_active=True, + ) + + +class _SharedCache: + """Mimics the module-level _token_cache global in app.py. + + Both reader and writer access .current — the same reference object. + The writer atomically reassigns .current to a new dict; the GIL + ensures the reader never sees a half-built reference. + """ + + def __init__(self, initial=None): + self.current = initial or {} + + +def _build_refresh_fn(shared, rows): + """Build a _refresh_token_cache closure mirroring app.py's fixed logic.""" + def _refresh(): + new_map = defaultdict(list) + for r in rows: + scope_list = [s.strip() for s in (r.scopes or "chat").split(",") if s.strip()] + new_map[r.token_prefix].append((r.id, r.token_hash, r.owner, scope_list)) + shared.current = dict(new_map) + return _refresh + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +class TestAtomicSwapNoEmptyWindow: + """The atomic swap must never leave _token_cache empty between frames.""" + + def test_swap_replaces_dict_content(self): + shared = _SharedCache({"old_prefix": [("old_id", "old_hash", "admin", ["chat"])]}) + rows = [_make_token_row("new_prfx", "t2", "h2", "admin", "chat")] + refresh = _build_refresh_fn(shared, rows) + + refresh() + + assert "old_prefix" not in shared.current + assert "new_prfx" in shared.current + assert shared.current["new_prfx"][0] == ("t2", "h2", "admin", ["chat"]) + + def test_swap_is_atomic_under_concurrent_readers(self): + """Concurrent readers never see an empty dict during refresh.""" + shared = _SharedCache({"ody_test12": [("t1", "hash1", "admin", ["chat"])]}) + rows = [_make_token_row("ody_newtok", "t2", "hash2", "admin", "chat")] + refresh = _build_refresh_fn(shared, rows) + + stop = threading.Event() + reader_results = {"empty": 0, "ok": 0} + + def reader_loop(): + while not stop.is_set(): + # Mimic app.py line 429: read the shared global directly + snapshot = shared.current + if len(snapshot) == 0: + reader_results["empty"] += 1 + else: + reader_results["ok"] += 1 + + threads = [threading.Thread(target=reader_loop, daemon=True) for _ in range(4)] + for t in threads: + t.start() + + time.sleep(0.01) + for _ in range(100): + refresh() + time.sleep(0.01) + stop.set() + for t in threads: + t.join(timeout=2) + + assert reader_results["empty"] == 0, ( + "Readers saw empty cache %d times (ok=%d)" + % (reader_results["empty"], reader_results["ok"]) + ) + assert reader_results["ok"] > 0, "readers should have seen data at least once" + + def test_app_state_ref_stays_in_sync(self): + """app.state._token_cache must point to the same dict.""" + shared = _SharedCache() + rows = [_make_token_row("pfx_a")] + refresh = _build_refresh_fn(shared, rows) + + app_state = SimpleNamespace(_token_cache=shared.current) + refresh() + app_state._token_cache = shared.current + + assert app_state._token_cache is shared.current + assert "pfx_a" in app_state._token_cache + + +class TestRefreshFromDB: + """Verify the refresh logic handles DB rows correctly.""" + + def test_multiple_prefixes(self): + shared = _SharedCache() + rows = [ + _make_token_row("ody_aaaa", "t1", "h1", "admin", "chat"), + _make_token_row("ody_bbbb", "t2", "h2", "admin", "chat,tools"), + _make_token_row("ody_aaaa", "t3", "h3", "admin", "memory"), + ] + refresh = _build_refresh_fn(shared, rows) + + refresh() + + assert len(shared.current) == 2 + assert len(shared.current["ody_aaaa"]) == 2 + assert len(shared.current["ody_bbbb"]) == 1 + assert shared.current["ody_bbbb"][0][3] == ["chat", "tools"] + + def test_empty_db_clears_cache(self): + shared = _SharedCache({"stale": [("x", "y", "z", ["chat"])]}) + refresh = _build_refresh_fn(shared, rows=[]) + + refresh() + + assert len(shared.current) == 0 + + def test_concurrent_refreshes_dont_corrupt(self): + """Multiple threads refreshing simultaneously don't corrupt cache.""" + shared = _SharedCache() + rows = [ + _make_token_row("pfx_%d" % i, "t%d" % i, "h%d" % i, "admin", "chat") + for i in range(20) + ] + refresh = _build_refresh_fn(shared, rows) + + errors = [] + + def refresh_loop(): + try: + for _ in range(50): + refresh() + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=refresh_loop) for _ in range(4)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=5) + + assert errors == [] + assert len(shared.current) == 20