feat: [Option A] configurable runtime-token header (server + SDK) [Need Discussion]

sdks/python/src/agent_control/client.py

@@ -20,6 +20,8 @@
 _logger = logging.getLogger(__name__)
 
 _RUNTIME_AUTH_MODE_ENV_VAR = "AGENT_CONTROL_RUNTIME_AUTH_MODE"
+_RUNTIME_TOKEN_HEADER_ENV_VAR = "AGENT_CONTROL_RUNTIME_TOKEN_HEADER"
+_DEFAULT_RUNTIME_TOKEN_HEADER = "Authorization"
 _DEFAULT_RUNTIME_TOKEN_REFRESH_MARGIN_SECONDS = 30
 _AUTO_RUNTIME_TOKEN_FALLBACK_STATUSES = {404, 500, 502, 503, 504}
 _GLOBAL_RUNTIME_TOKEN_FALLBACK_STATUSES = {404}
@@ -35,19 +37,39 @@ def _runtime_cache_identity(api_key: str | None, api_key_header: str) -> str:
 
 
 class _AgentControlAuth(httpx.Auth):
-    """Attach local API-key credentials unless a request already has Bearer auth."""
+    """Attach local API-key credentials unless a request already carries a token.
 
-    def __init__(self, api_key: str | None, header_name: str = "X-API-Key") -> None:
+    The API key is suppressed when the request already presents a bearer
+    credential on ``Authorization`` or a runtime token on its dedicated
+    header, so a runtime-authenticated evaluation carries a single
+    credential regardless of which header the runtime token rides.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None,
+        header_name: str = "X-API-Key",
+        runtime_token_header: str | None = None,
+    ) -> None:
         self._api_key = api_key
         self._header_name = header_name
+        self._runtime_token_header = runtime_token_header
 
     def auth_flow(
         self,
         request: httpx.Request,
     ) -> Generator[httpx.Request, httpx.Response, None]:
-        if self._api_key and "Authorization" not in request.headers:
-            if self._header_name not in request.headers:
-                request.headers[self._header_name] = self._api_key
+        runtime_token_on_dedicated_header = (
+            self._runtime_token_header is not None
+            and self._runtime_token_header in request.headers
+        )
+        if (
+            self._api_key
+            and "Authorization" not in request.headers
+            and not runtime_token_on_dedicated_header
+            and self._header_name not in request.headers
+        ):
+            request.headers[self._header_name] = self._api_key
         yield request
 
 
@@ -99,6 +121,7 @@ def __init__(
         api_key: str | None = None,
         api_key_header: str | None = None,
         runtime_auth_mode: RuntimeAuthMode | str | None = None,
+        runtime_token_header: str | None = None,
         runtime_token_cache: RuntimeTokenCache | None = None,
         runtime_token_refresh_margin_seconds: int = (_DEFAULT_RUNTIME_TOKEN_REFRESH_MARGIN_SECONDS),
         transport: httpx.AsyncBaseTransport | None = None,
@@ -121,6 +144,13 @@ def __init__(
                 request auth when the exchange endpoint is unavailable. ``jwt``
                 requires a successful exchange. ``api_key`` and ``none`` keep
                 evaluation requests on the normal request-auth path.
+            runtime_token_header: HTTP header name to send the runtime token
+                on. Defaults to ``Authorization``; the
+                AGENT_CONTROL_RUNTIME_TOKEN_HEADER environment variable
+                overrides the default. Point this at a dedicated header (e.g.
+                ``X-Agent-Control-Runtime-Token``) when the server runs behind
+                a gateway that reserves ``Authorization`` for its own identity
+                JWT. The server must be configured to read the same header.
             runtime_token_cache: Optional cache shared across client instances.
             runtime_token_refresh_margin_seconds: Refresh cached runtime tokens
                 before this many seconds of validity remain.
@@ -140,6 +170,27 @@ def __init__(
         self._runtime_cache_identity = _runtime_cache_identity(self._api_key, self._api_key_header)
         configured_runtime_mode = runtime_auth_mode or os.environ.get(_RUNTIME_AUTH_MODE_ENV_VAR)
         self._runtime_auth_mode = normalize_runtime_auth_mode(configured_runtime_mode)
+        if runtime_token_header is not None and not runtime_token_header.strip():
+            raise ValueError("runtime_token_header must not be blank.")
+        # A whitespace-only env var falls back to the default header, matching
+        # the server's _resolve_runtime_token_header(); an explicit blank
+        # *param* above is a hard error.
+        env_runtime_token_header = os.environ.get(_RUNTIME_TOKEN_HEADER_ENV_VAR)
+        if env_runtime_token_header is not None and not env_runtime_token_header.strip():
+            env_runtime_token_header = None
+        self._runtime_token_header = (
+            runtime_token_header
+            or env_runtime_token_header
+            or _DEFAULT_RUNTIME_TOKEN_HEADER
+        ).strip()
+        # On Authorization the server keeps the ``Bearer`` scheme prefix
+        # (existing contract); a dedicated header carries the raw token so the
+        # value no longer collides with a gateway's Authorization JWT.
+        # Mirror of LocalJwtVerifyProvider._require_bearer on the server — the
+        # send format and the verifier's expectation must stay in sync.
+        self._runtime_token_use_bearer = (
+            self._runtime_token_header.lower() == _DEFAULT_RUNTIME_TOKEN_HEADER.lower()
+        )
         if runtime_token_refresh_margin_seconds < 0:
             raise ValueError("runtime_token_refresh_margin_seconds must be >= 0.")
         self._runtime_token_refresh_margin_seconds = runtime_token_refresh_margin_seconds
@@ -198,7 +249,13 @@ async def __aenter__(self) -> "AgentControlClient":
             base_url=self.base_url,
             timeout=self.timeout,
             headers=self._get_headers(),
-            auth=_AgentControlAuth(self._api_key, self._api_key_header),
+            auth=_AgentControlAuth(
+                self._api_key,
+                self._api_key_header,
+                runtime_token_header=(
+                    None if self._runtime_token_use_bearer else self._runtime_token_header
+                ),
+            ),
             transport=self._transport,
             event_hooks={"response": [self._check_server_version]},
         )
@@ -295,18 +352,28 @@ async def post_runtime_evaluation(
 
         return response
 
+    def _format_runtime_token(self, token: str) -> str:
+        """Format a runtime token for the configured header.
+
+        Prefixes ``Bearer `` only on the ``Authorization`` header (existing
+        server contract); a dedicated runtime header carries the raw token.
+        """
+        if self._runtime_token_use_bearer:
+            return f"Bearer {token}"
+        return token
+
     def _merge_runtime_headers(
         self,
         headers: dict[str, str] | None,
         runtime_authorization: str | None,
     ) -> dict[str, str] | None:
-        """Merge caller headers with an optional Bearer token."""
+        """Merge caller headers with an optional runtime token header."""
         if headers is None and runtime_authorization is None:
             return None
 
         merged = dict(headers or {})
         if runtime_authorization is not None:
-            merged["Authorization"] = runtime_authorization
+            merged[self._runtime_token_header] = runtime_authorization
         return merged
 
     async def _runtime_authorization(
@@ -350,7 +417,7 @@ async def _runtime_authorization(
                 refresh_margin_seconds=self._runtime_token_refresh_margin_seconds,
             )
             if cached is not None:
-                return f"Bearer {cached.token}"
+                return self._format_runtime_token(cached.token)
 
         exchange_lock = self._runtime_token_cache.exchange_lock(
             self.base_url,
@@ -379,7 +446,7 @@ async def _runtime_authorization(
                     refresh_margin_seconds=self._runtime_token_refresh_margin_seconds,
                 )
                 if cached is not None:
-                    return f"Bearer {cached.token}"
+                    return self._format_runtime_token(cached.token)
 
             token = await self._exchange_runtime_token(
                 target_type=target_type,
@@ -388,7 +455,7 @@ async def _runtime_authorization(
             )
         if token is None:
             return None
-        return f"Bearer {token}"
+        return self._format_runtime_token(token)
 
     async def _exchange_runtime_token(
         self,

sdks/python/tests/test_client.py

@@ -226,6 +226,101 @@ def handler(request: httpx.Request) -> httpx.Response:
     assert evaluation_api_key_headers == [None, None]
 
 
+@pytest.mark.asyncio
+async def test_runtime_evaluation_sends_token_on_custom_header() -> None:
+    """With a custom runtime header, the raw token rides that header (no Bearer)
+    and Authorization is left free for a gateway's own JWT."""
+    evaluation_runtime_headers: list[str | None] = []
+    evaluation_authorization_headers: list[str | None] = []
+    evaluation_api_key_headers: list[str | None] = []
+    expires_at = (datetime.now(UTC) + timedelta(minutes=5)).isoformat()
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        if request.url.path == "/api/v1/auth/runtime-token-exchange":
+            # The exchange itself must still authenticate with the API key,
+            # even when the runtime token later rides a dedicated header.
+            assert request.headers.get("X-API-Key") == "test-key"
+            return httpx.Response(
+                200,
+                json={
+                    "token": "runtime-token",
+                    "expires_at": expires_at,
+                    "target_type": "log_stream",
+                    "target_id": "ls-1",
+                    "scopes": ["runtime.use"],
+                },
+            )
+
+        evaluation_runtime_headers.append(
+            request.headers.get("X-Agent-Control-Runtime-Token")
+        )
+        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="jwt",
+        runtime_token_header="X-Agent-Control-Runtime-Token",
+        transport=transport,
+    ) as client:
+        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
+    # Raw token on the dedicated header, no Bearer prefix.
+    assert evaluation_runtime_headers == ["runtime-token"]
+    # Authorization untouched, free for the gateway JWT.
+    assert evaluation_authorization_headers == [None]
+    # The runtime token is the sole credential: the API key must not ride
+    # along just because Authorization happens to be free in custom-header mode.
+    assert evaluation_api_key_headers == [None]
+
+
+@pytest.mark.asyncio
+async def test_runtime_token_header_env_var_overrides_default(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """AGENT_CONTROL_RUNTIME_TOKEN_HEADER selects the send header."""
+    monkeypatch.setenv(
+        "AGENT_CONTROL_RUNTIME_TOKEN_HEADER", "X-Agent-Control-Runtime-Token"
+    )
+    client = AgentControlClient(base_url="https://agent-control.test")
+    assert client._runtime_token_header == "X-Agent-Control-Runtime-Token"
+    assert client._runtime_token_use_bearer is False
+
+
+def test_runtime_token_header_defaults_to_authorization() -> None:
+    client = AgentControlClient(base_url="https://agent-control.test")
+    assert client._runtime_token_header == "Authorization"
+    assert client._runtime_token_use_bearer is True
+
+
+def test_runtime_token_header_rejects_blank() -> None:
+    with pytest.raises(ValueError, match="runtime_token_header"):
+        AgentControlClient(
+            base_url="https://agent-control.test", runtime_token_header="  "
+        )
+
+
+def test_runtime_token_header_whitespace_env_defaults_to_authorization(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A whitespace-only env var falls back to Authorization, matching the
+    server's _resolve_runtime_token_header(); only an explicit blank param
+    is a hard error."""
+    monkeypatch.setenv("AGENT_CONTROL_RUNTIME_TOKEN_HEADER", "   ")
+    client = AgentControlClient(base_url="https://agent-control.test")
+    assert client._runtime_token_header == "Authorization"
+    assert client._runtime_token_use_bearer is True
+
+
 @pytest.mark.asyncio
 async def test_runtime_evaluation_single_flights_cold_cache_exchange() -> None:
     exchange_calls = 0
@@ -456,6 +551,49 @@ 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_fallback_keeps_api_key_on_custom_header() -> None:
+    """Custom-header mode must still fall back to the API key when the
+    exchange is unavailable: with no runtime token minted, the dedicated
+    header is absent and X-API-Key must ride the evaluation request."""
+    exchange_calls = 0
+    evaluation_api_key_headers: list[str | None] = []
+    evaluation_runtime_headers: list[str | None] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal exchange_calls
+        if request.url.path == "/api/v1/auth/runtime-token-exchange":
+            exchange_calls += 1
+            return httpx.Response(503, json={"detail": "runtime auth disabled"})
+
+        evaluation_api_key_headers.append(request.headers.get("X-API-Key"))
+        evaluation_runtime_headers.append(
+            request.headers.get("X-Agent-Control-Runtime-Token")
+        )
+        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_header="X-Agent-Control-Runtime-Token",
+        transport=transport,
+    ) as client:
+        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 == 1
+    # No token minted -> dedicated header absent -> API key authenticates.
+    assert evaluation_runtime_headers == [None]
+    assert evaluation_api_key_headers == ["test-key"]
+
+
 @pytest.mark.asyncio
 async def test_runtime_evaluation_auto_404_fallback_recovers_after_ttl() -> None:
     exchange_calls = 0

server/src/agent_control_server/auth_framework/config.py

@@ -21,6 +21,10 @@
   The ``runtime.token_exchange`` operation continues to flow through
   the default authorizer because the exchange itself is shaped like a
   management call (forward credential, get grant).
+  ``AGENT_CONTROL_RUNTIME_TOKEN_HEADER`` (default ``Authorization``)
+  selects which request header the ``jwt`` verifier reads the runtime
+  token from, so the server can run behind a gateway that reserves
+  ``Authorization`` for its own downstream identity JWT.
 """
 
 from __future__ import annotations
@@ -39,6 +43,7 @@
     NoAuthProvider,
 )
 from .providers.http_upstream import HttpUpstreamConfig
+from .providers.local_jwt import DEFAULT_RUNTIME_TOKEN_HEADER
 
 _logger = get_logger(__name__)
 
@@ -60,6 +65,7 @@
 _RUNTIME_MODE_ENV = "AGENT_CONTROL_RUNTIME_AUTH_MODE"
 _RUNTIME_TOKEN_SECRET_ENV = "AGENT_CONTROL_RUNTIME_TOKEN_SECRET"
 _RUNTIME_TOKEN_TTL_ENV = "AGENT_CONTROL_RUNTIME_TOKEN_TTL_SECONDS"
+_RUNTIME_TOKEN_HEADER_ENV = "AGENT_CONTROL_RUNTIME_TOKEN_HEADER"
 _DEFAULT_RUNTIME_TOKEN_TTL_SECONDS = 300
 # HS256 needs at least 256 bits (32 bytes) of secret material to be safe
 # against brute force; reject anything shorter so production deployments
@@ -378,12 +384,30 @@ def _build_runtime_provider(
     if mode == "jwt":
         if config is None:
             raise RuntimeError(f"{_RUNTIME_MODE_ENV}=jwt but runtime auth config is missing.")
-        return LocalJwtVerifyProvider(secret=config.secret)
+        return LocalJwtVerifyProvider(
+            secret=config.secret,
+            header_name=_resolve_runtime_token_header(),
+        )
     raise RuntimeError(
         f"Unknown runtime auth mode {mode!r}; expected 'none', 'api_key', or 'jwt'."
     )
 
 
+def _resolve_runtime_token_header() -> str:
+    """Return the header the runtime JWT verifier reads the Bearer token from.
+
+    Defaults to ``Authorization``. Deployments behind a gateway that
+    overwrites ``Authorization`` with its own downstream identity JWT can
+    point the verifier at a dedicated header (e.g.
+    ``X-Agent-Control-Runtime-Token``) so the two tokens no longer
+    collide on the hot evaluation path.
+    """
+    raw = os.environ.get(_RUNTIME_TOKEN_HEADER_ENV)
+    if raw is None or not raw.strip():
+        return DEFAULT_RUNTIME_TOKEN_HEADER
+    return raw.strip()
+
+
 def _load_runtime_auth_config(*, require_secret: bool = False) -> RuntimeAuthConfig | None:
     """Parse, validate, and return the runtime-auth config from env.