fix(sdk): bound runtime auth fallback cache

abhinav-galileo · abhinav-galileo · commit 4e0e0d112218 · 2026-05-15T14:41:57.000+05:30
diff --git a/sdks/python/src/agent_control/client.py b/sdks/python/src/agent_control/client.py
@@ -371,6 +371,12 @@ async def _exchange_runtime_token(
             )
         except httpx.RequestError:
             if self._runtime_auth_mode == "auto" and allow_auto_fallback:
+                _logger.debug(
+                    "Runtime token exchange request failed; falling back to normal request auth "
+                    "for %s/%s.",
+                    target_type,
+                    target_id,
+                )
                 self._runtime_token_cache.mark_jwt_unavailable(
                     server_url=self.base_url,
                     target_type=target_type,
@@ -384,6 +390,13 @@ async def _exchange_runtime_token(
             and allow_auto_fallback
             and response.status_code in _AUTO_RUNTIME_TOKEN_FALLBACK_STATUSES
         ):
+            _logger.debug(
+                "Runtime token exchange returned HTTP %s; falling back to normal request auth "
+                "for %s/%s.",
+                response.status_code,
+                target_type,
+                target_id,
+            )
             self._runtime_token_cache.mark_jwt_unavailable(
                 server_url=self.base_url,
                 target_type=target_type,
diff --git a/sdks/python/src/agent_control/runtime_auth.py b/sdks/python/src/agent_control/runtime_auth.py
@@ -12,7 +12,7 @@
 RuntimeAuthMode = Literal["auto", "none", "api_key", "jwt"]
 
 _TokenKey = tuple[str, str, str]
-_LockKey = tuple[str, str, str, int]
+_LockKey = tuple[str, str, str, asyncio.AbstractEventLoop]
 _DEFAULT_MAX_CACHE_ENTRIES = 256
 _DEFAULT_JWT_UNAVAILABLE_TTL_SECONDS = 30
 
@@ -37,12 +37,21 @@ def is_fresh(self, *, refresh_margin_seconds: int) -> bool:
 class RuntimeTokenCache:
     """Thread-safe runtime token cache keyed by server and target."""
 
-    def __init__(self, *, max_entries: int = _DEFAULT_MAX_CACHE_ENTRIES) -> None:
+    def __init__(
+        self,
+        *,
+        max_entries: int = _DEFAULT_MAX_CACHE_ENTRIES,
+        jwt_unavailable_ttl_seconds: int = _DEFAULT_JWT_UNAVAILABLE_TTL_SECONDS,
+    ) -> None:
         if max_entries < 1:
             raise ValueError("max_entries must be >= 1.")
+        if jwt_unavailable_ttl_seconds < 0:
+            raise ValueError("jwt_unavailable_ttl_seconds must be >= 0.")
         self._max_entries = max_entries
+        self._jwt_unavailable_ttl_seconds = jwt_unavailable_ttl_seconds
         self._tokens: dict[_TokenKey, RuntimeToken] = {}
-        self._jwt_unavailable = False
+        self._jwt_unavailable_until: datetime | None = None
+        self._jwt_unavailable_servers: dict[str, datetime] = {}
         self._jwt_unavailable_targets: dict[_TokenKey, datetime] = {}
         self._exchange_locks: dict[_LockKey, asyncio.Lock] = {}
         self._lock = threading.Lock()
@@ -75,6 +84,7 @@ def set(self, token: RuntimeToken) -> None:
                 self._tokens.pop(oldest_key, None)
                 self._jwt_unavailable_targets.pop(oldest_key, None)
             self._tokens[key] = token
+            self._jwt_unavailable_servers.pop(token.server_url, None)
             self._jwt_unavailable_targets.pop(key, None)
 
     def remove(self, server_url: str, target_type: str, target_id: str) -> None:
@@ -89,13 +99,30 @@ def mark_jwt_unavailable(
         target_type: str | None = None,
         target_id: str | None = None,
         globally: bool = False,
-        ttl_seconds: int = _DEFAULT_JWT_UNAVAILABLE_TTL_SECONDS,
+        ttl_seconds: int | None = None,
     ) -> None:
         """Record that JWT runtime auth should not be attempted."""
+        ttl = self._jwt_unavailable_ttl_seconds if ttl_seconds is None else ttl_seconds
+        if ttl < 0:
+            raise ValueError("ttl_seconds must be >= 0.")
+        expires_at = datetime.now(UTC) + timedelta(seconds=ttl)
         with self._lock:
+            self._prune_expired_unavailable_markers_locked()
             if globally:
-                self._jwt_unavailable = True
-                self._tokens.clear()
+                if server_url is None:
+                    self._jwt_unavailable_until = expires_at
+                    self._tokens.clear()
+                    self._jwt_unavailable_servers.clear()
+                    self._jwt_unavailable_targets.clear()
+                    return
+
+                if (
+                    server_url not in self._jwt_unavailable_servers
+                    and len(self._jwt_unavailable_servers) >= self._max_entries
+                ):
+                    self._jwt_unavailable_servers.pop(next(iter(self._jwt_unavailable_servers)))
+                self._jwt_unavailable_servers[server_url] = expires_at
+                self._drop_server_entries_locked(server_url)
                 return
             if server_url is not None and target_type is not None and target_id is not None:
                 key = (server_url, target_type, target_id)
@@ -104,42 +131,70 @@ def mark_jwt_unavailable(
                     and len(self._jwt_unavailable_targets) >= self._max_entries
                 ):
                     self._jwt_unavailable_targets.pop(next(iter(self._jwt_unavailable_targets)))
-                self._jwt_unavailable_targets[key] = datetime.now(UTC) + timedelta(
-                    seconds=ttl_seconds
-                )
+                self._jwt_unavailable_targets[key] = expires_at
                 self._tokens.pop(key, None)
 
     def is_jwt_unavailable(self, server_url: str, target_type: str, target_id: str) -> bool:
         """Return whether JWT exchange is known unavailable for the target."""
         key = (server_url, target_type, target_id)
         with self._lock:
-            if self._jwt_unavailable:
+            self._prune_expired_unavailable_markers_locked()
+            if self._jwt_unavailable_until is not None:
                 return True
-            expires_at = self._jwt_unavailable_targets.get(key)
-            if expires_at is None:
-                return False
-            if expires_at > datetime.now(UTC):
+            if server_url in self._jwt_unavailable_servers:
                 return True
-            self._jwt_unavailable_targets.pop(key, None)
-            return False
+            expires_at = self._jwt_unavailable_targets.get(key)
+            return expires_at is not None
 
     def clear(self) -> None:
         """Clear every cached token and fallback marker."""
         with self._lock:
             self._tokens.clear()
-            self._jwt_unavailable = False
+            self._jwt_unavailable_until = None
+            self._jwt_unavailable_servers.clear()
             self._jwt_unavailable_targets.clear()
 
     def exchange_lock(self, server_url: str, target_type: str, target_id: str) -> asyncio.Lock:
         """Return the async exchange lock for one server and target."""
-        key = (server_url, target_type, target_id, id(asyncio.get_running_loop()))
+        key = (server_url, target_type, target_id, asyncio.get_running_loop())
         with self._lock:
             lock = self._exchange_locks.get(key)
             if lock is None:
+                if len(self._exchange_locks) >= self._max_entries:
+                    self._evict_idle_exchange_lock_locked()
                 lock = asyncio.Lock()
                 self._exchange_locks[key] = lock
+            else:
+                # Preserve insertion order as a simple LRU for idle-lock eviction.
+                self._exchange_locks.pop(key)
+                self._exchange_locks[key] = lock
             return lock
 
+    def _drop_server_entries_locked(self, server_url: str) -> None:
+        for key in list(self._tokens):
+            if key[0] == server_url:
+                self._tokens.pop(key, None)
+        for key in list(self._jwt_unavailable_targets):
+            if key[0] == server_url:
+                self._jwt_unavailable_targets.pop(key, None)
+
+    def _prune_expired_unavailable_markers_locked(self) -> None:
+        now = datetime.now(UTC)
+        if self._jwt_unavailable_until is not None and self._jwt_unavailable_until <= now:
+            self._jwt_unavailable_until = None
+        for server_url, expires_at in list(self._jwt_unavailable_servers.items()):
+            if expires_at <= now:
+                self._jwt_unavailable_servers.pop(server_url, None)
+        for key, expires_at in list(self._jwt_unavailable_targets.items()):
+            if expires_at <= now:
+                self._jwt_unavailable_targets.pop(key, None)
+
+    def _evict_idle_exchange_lock_locked(self) -> None:
+        for key, lock in list(self._exchange_locks.items()):
+            if not lock.locked():
+                self._exchange_locks.pop(key, None)
+                return
+
 
 def normalize_runtime_auth_mode(raw: str | None) -> RuntimeAuthMode:
     """Normalize configured SDK runtime auth mode."""
diff --git a/sdks/python/tests/test_client.py b/sdks/python/tests/test_client.py
@@ -369,6 +369,57 @@ def handler(request: httpx.Request) -> httpx.Response:
     assert evaluation_api_key_headers == ["test-key", "test-key"]
 
 
+@pytest.mark.asyncio
+async def test_runtime_evaluation_auto_404_fallback_recovers_after_ttl() -> None:
+    exchange_calls = 0
+    evaluation_authorization_headers: list[str | None] = []
+    evaluation_api_key_headers: list[str | None] = []
+    expires_at = (datetime.now(UTC) + timedelta(minutes=5)).isoformat()
+    cache = RuntimeTokenCache(jwt_unavailable_ttl_seconds=0)
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal exchange_calls
+        if request.url.path == "/api/v1/auth/runtime-token-exchange":
+            exchange_calls += 1
+            if exchange_calls == 1:
+                return httpx.Response(404, json={"detail": "not found"})
+            return httpx.Response(
+                200,
+                json={
+                    "token": "runtime-token",
+                    "expires_at": expires_at,
+                    "target_type": "log_stream",
+                    "target_id": "ls-1",
+                    "scopes": ["runtime.use"],
+                },
+            )
+
+        evaluation_authorization_headers.append(request.headers.get("Authorization"))
+        evaluation_api_key_headers.append(request.headers.get("X-API-Key"))
+        return httpx.Response(200, json={"is_safe": True, "confidence": 1.0})
+
+    transport = httpx.MockTransport(handler)
+
+    async with AgentControlClient(
+        base_url="https://agent-control.test",
+        api_key="test-key",
+        runtime_auth_mode="auto",
+        runtime_token_cache=cache,
+        transport=transport,
+    ) as client:
+        for _ in range(2):
+            response = await client.post_runtime_evaluation(
+                json={"target_type": "log_stream", "target_id": "ls-1"},
+                target_type="log_stream",
+                target_id="ls-1",
+            )
+            assert response.status_code == 200
+
+    assert exchange_calls == 2
+    assert evaluation_authorization_headers == [None, "Bearer runtime-token"]
+    assert evaluation_api_key_headers == ["test-key", None]
+
+
 @pytest.mark.asyncio
 async def test_runtime_evaluation_auto_503_fallback_is_target_scoped() -> None:
     exchange_targets: list[str] = []
diff --git a/sdks/python/tests/test_runtime_auth.py b/sdks/python/tests/test_runtime_auth.py
@@ -98,22 +98,41 @@ def test_runtime_token_cache_target_unavailable_marker_expires() -> None:
     assert not cache.is_jwt_unavailable("https://server-a.test", "log_stream", "ls-1")
 
 
-def test_runtime_token_cache_global_unavailable_clears_cache() -> None:
+def test_runtime_token_cache_server_unavailable_clears_server_cache() -> None:
     cache = RuntimeTokenCache()
     cache.set(_runtime_token())
+    token_2 = _runtime_token(
+        token="token-2",
+        server_url="https://server-b.test",
+        target_id="ls-2",
+    )
+    cache.set(token_2)
 
-    cache.mark_jwt_unavailable(globally=True)
+    cache.mark_jwt_unavailable(server_url="https://server-a.test", globally=True)
 
     assert cache.is_jwt_unavailable("https://server-a.test", "log_stream", "ls-1")
+    assert not cache.is_jwt_unavailable("https://server-b.test", "log_stream", "ls-2")
     assert (
         cache.get("https://server-a.test", "log_stream", "ls-1", refresh_margin_seconds=0) is None
     )
+    assert (
+        cache.get("https://server-b.test", "log_stream", "ls-2", refresh_margin_seconds=0)
+        == token_2
+    )
 
     cache.clear()
 
     assert not cache.is_jwt_unavailable("https://server-a.test", "log_stream", "ls-1")
 
 
+def test_runtime_token_cache_server_unavailable_marker_expires() -> None:
+    cache = RuntimeTokenCache(jwt_unavailable_ttl_seconds=0)
+
+    cache.mark_jwt_unavailable(server_url="https://server-a.test", globally=True)
+
+    assert not cache.is_jwt_unavailable("https://server-a.test", "log_stream", "ls-1")
+
+
 def test_runtime_token_cache_remove_drops_one_token() -> None:
     cache = RuntimeTokenCache()
     cache.set(_runtime_token(target_id="ls-1"))
@@ -159,6 +178,16 @@ async def test_runtime_token_cache_token_eviction_preserves_exchange_lock() -> N
     assert cache.exchange_lock("https://server-a.test", "log_stream", "ls-1") is lock
 
 
+@pytest.mark.asyncio
+async def test_runtime_token_cache_evicts_idle_exchange_locks() -> None:
+    cache = RuntimeTokenCache(max_entries=1)
+    first = cache.exchange_lock("https://server-a.test", "log_stream", "ls-1")
+
+    cache.exchange_lock("https://server-a.test", "log_stream", "ls-2")
+
+    assert cache.exchange_lock("https://server-a.test", "log_stream", "ls-1") is not first
+
+
 def test_runtime_token_cache_exchange_locks_are_loop_scoped(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
@@ -180,6 +209,18 @@ def test_runtime_token_cache_rejects_empty_capacity() -> None:
         RuntimeTokenCache(max_entries=0)
 
 
+def test_runtime_token_cache_rejects_negative_unavailable_ttl() -> None:
+    with pytest.raises(ValueError, match="jwt_unavailable_ttl_seconds"):
+        RuntimeTokenCache(jwt_unavailable_ttl_seconds=-1)
+
+
+def test_runtime_token_cache_rejects_negative_marker_ttl() -> None:
+    cache = RuntimeTokenCache()
+
+    with pytest.raises(ValueError, match="ttl_seconds"):
+        cache.mark_jwt_unavailable(ttl_seconds=-1)
+
+
 @pytest.mark.parametrize(
     ("raw", "expected"),
     [
