fix(sdk): harden runtime token auth

Author: abhinav-galileo

sdks/python/src/agent_control/__init__.py

@@ -542,6 +542,11 @@ async def handle(message: str):
         raise ValueError(
             "target_type and target_id must be supplied together."
         )
+    resolved_api_key_header = (
+        api_key_header
+        or os.getenv(AgentControlClient.API_KEY_HEADER_ENV_VAR)
+        or AgentControlClient.DEFAULT_API_KEY_HEADER
+    )
 
     # Re-init behavior: always stop the existing refresh loop before mutating
     # shared agent/session globals.
@@ -569,7 +574,7 @@ async def handle(message: str):
         state.current_agent = next_agent
         state.server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000'
         state.api_key = api_key
-        state.api_key_header = api_key_header
+        state.api_key_header = resolved_api_key_header
         state.runtime_token_cache.clear()
         state.target_type = target_type
         state.target_id = target_id

sdks/python/src/agent_control/client.py

@@ -20,8 +20,8 @@
 
 _RUNTIME_AUTH_MODE_ENV_VAR = "AGENT_CONTROL_RUNTIME_AUTH_MODE"
 _DEFAULT_RUNTIME_TOKEN_REFRESH_MARGIN_SECONDS = 30
-_AUTO_RUNTIME_TOKEN_FALLBACK_STATUSES = {404, 503}
-_GLOBAL_RUNTIME_TOKEN_FALLBACK_STATUSES = {404, 503}
+_AUTO_RUNTIME_TOKEN_FALLBACK_STATUSES = {404, 500, 502, 503, 504}
+_GLOBAL_RUNTIME_TOKEN_FALLBACK_STATUSES = {404}
 
 
 class _AgentControlAuth(httpx.Auth):
@@ -248,8 +248,10 @@ async def post_runtime_evaluation(
             headers=request_headers,
         )
 
-        if response.status_code == 401 and runtime_authorization is not None:
+        if _should_refresh_runtime_token(response) and runtime_authorization is not None:
             await response.aread()
+            if target_type is not None and target_id is not None:
+                self._runtime_token_cache.remove(self.base_url, target_type, target_id)
             runtime_authorization = await self._runtime_authorization(
                 target_type=target_type,
                 target_id=target_id,
@@ -262,6 +264,13 @@ async def post_runtime_evaluation(
                 json=json,
                 headers=request_headers,
             )
+            if (
+                _should_refresh_runtime_token(response)
+                and target_type is not None
+                and target_id is not None
+            ):
+                await response.aread()
+                self._runtime_token_cache.remove(self.base_url, target_type, target_id)
 
         return response
 
@@ -320,6 +329,14 @@ async def _runtime_authorization(
             target_id,
         )
         async with exchange_lock:
+            if (
+                self._runtime_auth_mode == "auto"
+                and not force_refresh
+                and self._runtime_token_cache.is_jwt_unavailable(
+                    self.base_url, target_type, target_id
+                )
+            ):
+                return None
             if not force_refresh:
                 cached = self._runtime_token_cache.get(
                     self.base_url,
@@ -347,10 +364,20 @@ async def _exchange_runtime_token(
         allow_auto_fallback: bool = True,
     ) -> str | None:
         """Exchange the configured credential for a target-bound runtime token."""
-        response = await self.http_client.post(
-            "/api/v1/auth/runtime-token-exchange",
-            json={"target_type": target_type, "target_id": target_id},
-        )
+        try:
+            response = await self.http_client.post(
+                "/api/v1/auth/runtime-token-exchange",
+                json={"target_type": target_type, "target_id": target_id},
+            )
+        except httpx.RequestError:
+            if self._runtime_auth_mode == "auto" and allow_auto_fallback:
+                self._runtime_token_cache.mark_jwt_unavailable(
+                    server_url=self.base_url,
+                    target_type=target_type,
+                    target_id=target_id,
+                )
+                return None
+            raise
 
         if (
             self._runtime_auth_mode == "auto"
@@ -373,5 +400,18 @@ async def _exchange_runtime_token(
             cast(dict[str, object], payload),
             server_url=self.base_url,
         )
+        if token.target_type != target_type or token.target_id != target_id:
+            raise RuntimeError(
+                "Runtime token exchange response target did not match the requested target."
+            )
         self._runtime_token_cache.set(token)
         return token.token
+
+
+def _should_refresh_runtime_token(response: httpx.Response) -> bool:
+    if response.status_code == 401:
+        return True
+    if response.status_code != 403:
+        return False
+    authenticate = response.headers.get("WWW-Authenticate", "")
+    return "invalid_token" in authenticate.lower()

sdks/python/src/agent_control/control_decorators.py

@@ -40,7 +40,11 @@ async def chat(message: str) -> str:
 
 from agent_control import AgentControlClient
 from agent_control._state import state
-from agent_control.evaluation import _resolve_session_target, check_evaluation_with_local
+from agent_control.evaluation import (
+    _post_evaluation_request,
+    _resolve_session_target,
+    check_evaluation_with_local,
+)
 from agent_control.observability import (
     get_logger,
     log_control_evaluation,
@@ -388,10 +392,12 @@ async def _evaluate(
             payload["target_type"] = target_type
             payload["target_id"] = target_id
 
-        response = await client.http_client.post(
-            "/api/v1/evaluation",
-            json=payload,
+        response = await _post_evaluation_request(
+            client,
+            request_payload=payload,
             headers=headers,
+            target_type=target_type,
+            target_id=target_id,
         )
         response.raise_for_status()
         result_dict: dict[str, Any] = response.json()

sdks/python/src/agent_control/evaluation.py

@@ -229,7 +229,9 @@ async def _post_evaluation_request(
 ) -> httpx.Response:
     """Send an evaluation request, using runtime auth when the client supports it."""
     runtime_post = None
-    if (target_type is not None and target_id is not None) or client.runtime_auth_mode == "jwt":
+    if (target_type is not None and target_id is not None) or getattr(
+        client, "runtime_auth_mode", "auto"
+    ) == "jwt":
         runtime_post = _runtime_post_evaluation(client)
     if runtime_post is not None:
         return await runtime_post(

sdks/python/src/agent_control/runtime_auth.py

@@ -5,21 +5,23 @@
 import asyncio
 import threading
 from collections.abc import Mapping, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import UTC, datetime, timedelta
 from typing import Literal
 
 RuntimeAuthMode = Literal["auto", "none", "api_key", "jwt"]
 
 _TokenKey = tuple[str, str, str]
+_LockKey = tuple[str, str, str, int]
 _DEFAULT_MAX_CACHE_ENTRIES = 256
+_DEFAULT_JWT_UNAVAILABLE_TTL_SECONDS = 30
 
 
 @dataclass(frozen=True)
 class RuntimeToken:
     """Short-lived runtime token bound to one target."""
 
-    token: str
+    token: str = field(repr=False)
     expires_at: datetime
     server_url: str
     target_type: str
@@ -41,8 +43,8 @@ def __init__(self, *, max_entries: int = _DEFAULT_MAX_CACHE_ENTRIES) -> None:
         self._max_entries = max_entries
         self._tokens: dict[_TokenKey, RuntimeToken] = {}
         self._jwt_unavailable = False
-        self._jwt_unavailable_targets: set[_TokenKey] = set()
-        self._exchange_locks: dict[_TokenKey, asyncio.Lock] = {}
+        self._jwt_unavailable_targets: dict[_TokenKey, datetime] = {}
+        self._exchange_locks: dict[_LockKey, asyncio.Lock] = {}
         self._lock = threading.Lock()
 
     def get(
@@ -71,10 +73,9 @@ def set(self, token: RuntimeToken) -> None:
             if key not in self._tokens and len(self._tokens) >= self._max_entries:
                 oldest_key = next(iter(self._tokens))
                 self._tokens.pop(oldest_key, None)
-                self._jwt_unavailable_targets.discard(oldest_key)
-                self._exchange_locks.pop(oldest_key, None)
+                self._jwt_unavailable_targets.pop(oldest_key, None)
             self._tokens[key] = token
-            self._jwt_unavailable_targets.discard(key)
+            self._jwt_unavailable_targets.pop(key, None)
 
     def remove(self, server_url: str, target_type: str, target_id: str) -> None:
         """Drop the cached token for one target."""
@@ -88,6 +89,7 @@ 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,
     ) -> None:
         """Record that JWT runtime auth should not be attempted."""
         with self._lock:
@@ -101,28 +103,36 @@ def mark_jwt_unavailable(
                     key not in self._jwt_unavailable_targets
                     and len(self._jwt_unavailable_targets) >= self._max_entries
                 ):
-                    evicted_key = self._jwt_unavailable_targets.pop()
-                    self._exchange_locks.pop(evicted_key, None)
-                self._jwt_unavailable_targets.add(key)
+                    self._jwt_unavailable_targets.pop(next(iter(self._jwt_unavailable_targets)))
+                self._jwt_unavailable_targets[key] = datetime.now(UTC) + timedelta(
+                    seconds=ttl_seconds
+                )
                 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:
-            return self._jwt_unavailable or key in self._jwt_unavailable_targets
+            if self._jwt_unavailable:
+                return True
+            expires_at = self._jwt_unavailable_targets.get(key)
+            if expires_at is None:
+                return False
+            if expires_at > datetime.now(UTC):
+                return True
+            self._jwt_unavailable_targets.pop(key, None)
+            return False
 
     def clear(self) -> None:
         """Clear every cached token and fallback marker."""
         with self._lock:
             self._tokens.clear()
             self._jwt_unavailable = False
             self._jwt_unavailable_targets.clear()
-            self._exchange_locks.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)
+        key = (server_url, target_type, target_id, id(asyncio.get_running_loop()))
         with self._lock:
             lock = self._exchange_locks.get(key)
             if lock is None: