From 03f26810d2a3a0a2755646c889ca84530b9e287b Mon Sep 17 00:00:00 2001 From: Joshua Jeon Date: Wed, 15 Jul 2026 09:41:56 -0700 Subject: [PATCH 1/2] feat(auth): make the runtime-token header configurable Let the runtime-token JWT verifier read its token from a configurable request header (env AGENT_CONTROL_RUNTIME_TOKEN_HEADER), defaulting to Authorization so existing deployments are unaffected. Why: when Agent Control runs behind an API gateway that reserves the Authorization header for its own downstream identity JWT, the gateway overwrites the runtime token on the hot evaluation path and the verifier fails. Pointing the verifier at a dedicated header (e.g. X-Agent-Control-Runtime-Token) lets the two tokens coexist. - LocalJwtVerifyProvider gains a header_name param. On Authorization the Bearer scheme prefix stays required (back-compat); on a dedicated header the raw token is accepted (Bearer optional). The token is still signature-verified, scope-checked, and target-bound after extraction, so the header choice cannot bypass verification. - config.py resolves the header via _resolve_runtime_token_header(); a whitespace-only env value falls back to the default. - Tests: custom-header raw/Bearer acceptance, case-insensitive lookup, no fallback to Authorization, whitespace/blank handling, and app-level E2E through /api/v1/evaluation (gateway JWT on Authorization coexists with the runtime token on a dedicated header). Co-Authored-By: Claude Opus 4.7 --- .../auth_framework/config.py | 26 ++- .../auth_framework/providers/local_jwt.py | 61 +++-- server/tests/test_auth_framework.py | 219 +++++++++++++++++- .../test_runtime_token_exchange_endpoint.py | 126 ++++++++++ 4 files changed, 415 insertions(+), 17 deletions(-) diff --git a/server/src/agent_control_server/auth_framework/config.py b/server/src/agent_control_server/auth_framework/config.py index 06246a46..9fbab903 100644 --- a/server/src/agent_control_server/auth_framework/config.py +++ b/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. diff --git a/server/src/agent_control_server/auth_framework/providers/local_jwt.py b/server/src/agent_control_server/auth_framework/providers/local_jwt.py index 3f39e6fd..ab407631 100644 --- a/server/src/agent_control_server/auth_framework/providers/local_jwt.py +++ b/server/src/agent_control_server/auth_framework/providers/local_jwt.py @@ -1,11 +1,17 @@ """Authorizer that verifies a locally-minted runtime token. -Wired to the runtime resolution path. Reads a Bearer token from the -``Authorization`` header, verifies the signature against the runtime -secret, checks the token's scope covers the requested operation, and -returns a :class:`Principal` carrying the bound target. When a -``context_builder`` on the dependency must surface matching -``target_type`` / ``target_id`` values for target-bound tokens. +Wired to the runtime resolution path. Reads a Bearer token from a +configurable header (``Authorization`` by default), verifies the +signature against the runtime secret, checks the token's scope covers +the requested operation, and returns a :class:`Principal` carrying the +bound target. When a ``context_builder`` on the dependency must surface +matching ``target_type`` / ``target_id`` values for target-bound tokens. + +The header the runtime token arrives on is configurable so the server +can sit behind a gateway that reserves ``Authorization`` for its own +downstream identity JWT. When the token rides a dedicated header (e.g. +``X-Agent-Control-Runtime-Token``), the ``Bearer`` scheme prefix is +optional; on ``Authorization`` it stays required for back-compat. """ from __future__ import annotations @@ -19,14 +25,30 @@ from ..core import Operation, Principal, RequestAuthorizer from ..runtime_token import RuntimeTokenError, verify_runtime_token +DEFAULT_RUNTIME_TOKEN_HEADER = "Authorization" + class LocalJwtVerifyProvider(RequestAuthorizer): """Verifies a runtime Bearer token and emits a target-bound :class:`Principal`.""" - def __init__(self, *, secret: str) -> None: + def __init__( + self, + *, + secret: str, + header_name: str = DEFAULT_RUNTIME_TOKEN_HEADER, + ) -> None: if not secret: raise ValueError("LocalJwtVerifyProvider requires a non-empty secret.") + if not header_name or not header_name.strip(): + raise ValueError("LocalJwtVerifyProvider requires a non-empty header_name.") self._secret = secret + self._header_name = header_name.strip() + # The Bearer scheme is mandatory on Authorization (existing SDK + # contract) but optional on a dedicated runtime-token header, + # where the raw token is the whole value. + # Mirror of AgentControlClient._runtime_token_use_bearer in the SDK — + # the verifier's expectation and the send format must stay in sync. + self._require_bearer = self._header_name.lower() == "authorization" async def authorize( self, @@ -79,18 +101,29 @@ async def authorize( ) def _extract_bearer_token(self, request: Request) -> str: - header = request.headers.get("Authorization") + header = request.headers.get(self._header_name) if not header: raise AuthenticationError( error_code=ErrorCode.AUTH_MISSING_KEY, - detail="Missing Authorization header.", + detail=f"Missing {self._header_name} header.", hint="Present a Bearer runtime token.", ) - scheme, _, value = header.partition(" ") - if scheme.lower() != "bearer" or not value: + scheme, sep, value = header.partition(" ") + if sep and scheme.lower() == "bearer": + token = value.strip() + elif self._require_bearer: raise AuthenticationError( error_code=ErrorCode.AUTH_MISSING_KEY, - detail="Authorization header must be a Bearer token.", - hint="Format: ``Authorization: Bearer ``.", + detail=f"{self._header_name} header must be a Bearer token.", + hint=f"Format: ``{self._header_name}: Bearer ``.", + ) + else: + # Dedicated runtime-token header: accept the raw token. + token = header.strip() + if not token: + raise AuthenticationError( + error_code=ErrorCode.AUTH_MISSING_KEY, + detail=f"{self._header_name} header is empty.", + hint="Present a Bearer runtime token.", ) - return value.strip() + return token diff --git a/server/tests/test_auth_framework.py b/server/tests/test_auth_framework.py index c3514fba..78210442 100644 --- a/server/tests/test_auth_framework.py +++ b/server/tests/test_auth_framework.py @@ -7,6 +7,8 @@ import httpx import pytest +from starlette.datastructures import Headers + from agent_control_server.auth_framework.core import ( Operation, Principal, @@ -43,9 +45,14 @@ def _build_request( headers: dict[str, str] | None = None, cookies: dict[str, str] | None = None, ): - """Build a minimal Starlette-compatible request mock.""" + """Build a minimal Starlette-compatible request mock. + + Uses the real :class:`starlette.datastructures.Headers` so header + lookups are case-insensitive exactly as they are at runtime; a plain + dict would be case-sensitive and mask casing bugs. + """ request = MagicMock() - request.headers = headers or {} + request.headers = Headers(headers=headers or {}) request.cookies = cookies or {} return request @@ -875,6 +882,7 @@ def test_runtime_token_rejects_empty_required_claims(kwargs, message): def test_runtime_token_rejects_management_token_passed_to_runtime_verify(): """A token without ``domain=runtime`` must be rejected by runtime verify.""" import jwt + from agent_control_server.auth_framework.runtime_token import ( RuntimeTokenError, verify_runtime_token, @@ -975,6 +983,158 @@ async def test_local_jwt_provider_rejects_non_bearer_authorization(): await provider.authorize(request, Operation.RUNTIME_USE) +@pytest.mark.asyncio +async def test_local_jwt_provider_reads_configured_custom_header(): + """A dedicated runtime-token header carries the raw token (no Bearer).""" + from agent_control_server.auth_framework.providers import LocalJwtVerifyProvider + from agent_control_server.auth_framework.runtime_token import ( + mint_runtime_token, + ) + + token, _ = mint_runtime_token( + namespace_key="default", + actor_id="actor-7", + target_type="log_stream", + target_id="ls-42", + scopes=("runtime.use",), + secret=_TEST_SECRET, + ttl_seconds=60, + ) + provider = LocalJwtVerifyProvider( + secret=_TEST_SECRET, header_name="X-Agent-Control-Runtime-Token" + ) + request = _build_request(headers={"X-Agent-Control-Runtime-Token": token}) + + principal = await provider.authorize( + request, + Operation.RUNTIME_USE, + context={"target_type": "log_stream", "target_id": "ls-42"}, + ) + + assert principal.target_id == "ls-42" + + +@pytest.mark.asyncio +async def test_local_jwt_provider_custom_header_accepts_bearer_prefix(): + """Custom header still accepts an explicit Bearer prefix.""" + from agent_control_server.auth_framework.providers import LocalJwtVerifyProvider + from agent_control_server.auth_framework.runtime_token import ( + mint_runtime_token, + ) + + token, _ = mint_runtime_token( + namespace_key="default", + actor_id="a", + target_type="log_stream", + target_id="ls", + scopes=("runtime.use",), + secret=_TEST_SECRET, + ttl_seconds=60, + ) + provider = LocalJwtVerifyProvider( + secret=_TEST_SECRET, header_name="X-Agent-Control-Runtime-Token" + ) + request = _build_request( + headers={"X-Agent-Control-Runtime-Token": f"Bearer {token}"} + ) + + principal = await provider.authorize( + request, + Operation.RUNTIME_USE, + context={"target_type": "log_stream", "target_id": "ls"}, + ) + assert principal.caller_id == "a" + + +@pytest.mark.asyncio +async def test_local_jwt_provider_custom_header_does_not_fall_back_to_authorization(): + """A valid token on Authorization must NOT satisfy a custom-header verifier. + + Uses a genuinely valid runtime token on ``Authorization`` so the test + can only pass if the verifier reads the configured custom header + exclusively; a wrongful fallback to ``Authorization`` would let the + valid token through. The error must name the custom header. + """ + from agent_control_server.auth_framework.providers import LocalJwtVerifyProvider + from agent_control_server.auth_framework.runtime_token import ( + mint_runtime_token, + ) + from agent_control_server.errors import AuthenticationError + + token, _ = mint_runtime_token( + namespace_key="default", + actor_id="a", + target_type="log_stream", + target_id="ls", + scopes=("runtime.use",), + secret=_TEST_SECRET, + ttl_seconds=60, + ) + provider = LocalJwtVerifyProvider( + secret=_TEST_SECRET, header_name="X-Agent-Control-Runtime-Token" + ) + request = _build_request(headers={"Authorization": f"Bearer {token}"}) + with pytest.raises(AuthenticationError) as exc_info: + await provider.authorize(request, Operation.RUNTIME_USE) + assert "X-Agent-Control-Runtime-Token" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_local_jwt_provider_custom_header_whitespace_only_raises_401(): + """A whitespace-only custom header is treated as an empty token.""" + from agent_control_server.auth_framework.providers import LocalJwtVerifyProvider + from agent_control_server.errors import AuthenticationError + + provider = LocalJwtVerifyProvider( + secret=_TEST_SECRET, header_name="X-Agent-Control-Runtime-Token" + ) + request = _build_request(headers={"X-Agent-Control-Runtime-Token": " "}) + with pytest.raises(AuthenticationError): + await provider.authorize(request, Operation.RUNTIME_USE) + + +@pytest.mark.asyncio +async def test_local_jwt_provider_reads_custom_header_case_insensitively(): + """The verifier finds the token even when the client sends lower-case. + + HTTP header names are case-insensitive (HTTP/2 mandates lower-case), + so a verifier configured with ``X-Agent-Control-Runtime-Token`` must + still read a request that arrives as ``x-agent-control-runtime-token``. + """ + from agent_control_server.auth_framework.providers import LocalJwtVerifyProvider + from agent_control_server.auth_framework.runtime_token import ( + mint_runtime_token, + ) + + token, _ = mint_runtime_token( + namespace_key="default", + actor_id="a", + target_type="log_stream", + target_id="ls", + scopes=("runtime.use",), + secret=_TEST_SECRET, + ttl_seconds=60, + ) + provider = LocalJwtVerifyProvider( + secret=_TEST_SECRET, header_name="X-Agent-Control-Runtime-Token" + ) + request = _build_request(headers={"x-agent-control-runtime-token": token}) + + principal = await provider.authorize( + request, + Operation.RUNTIME_USE, + context={"target_type": "log_stream", "target_id": "ls"}, + ) + assert principal.caller_id == "a" + + +def test_local_jwt_provider_rejects_blank_header_name(): + from agent_control_server.auth_framework.providers import LocalJwtVerifyProvider + + with pytest.raises(ValueError, match="header_name"): + LocalJwtVerifyProvider(secret=_TEST_SECRET, header_name=" ") + + @pytest.mark.asyncio async def test_local_jwt_provider_carries_token_namespace_to_principal(): """Tokens minted in a non-default namespace must verify under that namespace.""" @@ -1722,6 +1882,61 @@ def test_configure_runtime_jwt_requires_secret(monkeypatch): auth_config.configure_auth_from_env() +def test_configure_runtime_jwt_defaults_to_authorization_header(monkeypatch): + from agent_control_server.auth_framework import config as auth_config + from agent_control_server.auth_framework.providers import ( + LocalJwtVerifyProvider, + ) + + clear_authorizers() + monkeypatch.delenv("AGENT_CONTROL_RUNTIME_TOKEN_HEADER", raising=False) + monkeypatch.setenv("AGENT_CONTROL_RUNTIME_TOKEN_SECRET", _TEST_SECRET) + auth_config.configure_auth_from_env() + + provider = get_authorizer(Operation.RUNTIME_USE) + assert isinstance(provider, LocalJwtVerifyProvider) + assert provider._header_name == "Authorization" + + +def test_configure_runtime_jwt_honors_custom_token_header(monkeypatch): + from agent_control_server.auth_framework import config as auth_config + from agent_control_server.auth_framework.providers import ( + LocalJwtVerifyProvider, + ) + + clear_authorizers() + monkeypatch.setenv("AGENT_CONTROL_RUNTIME_TOKEN_SECRET", _TEST_SECRET) + monkeypatch.setenv( + "AGENT_CONTROL_RUNTIME_TOKEN_HEADER", "X-Agent-Control-Runtime-Token" + ) + auth_config.configure_auth_from_env() + + provider = get_authorizer(Operation.RUNTIME_USE) + assert isinstance(provider, LocalJwtVerifyProvider) + assert provider._header_name == "X-Agent-Control-Runtime-Token" + assert provider._require_bearer is False + + +def test_configure_runtime_jwt_whitespace_only_header_env_defaults_to_authorization( + monkeypatch, +): + """A whitespace-only header env var falls back to the default header.""" + from agent_control_server.auth_framework import config as auth_config + from agent_control_server.auth_framework.providers import ( + LocalJwtVerifyProvider, + ) + + clear_authorizers() + monkeypatch.setenv("AGENT_CONTROL_RUNTIME_TOKEN_SECRET", _TEST_SECRET) + monkeypatch.setenv("AGENT_CONTROL_RUNTIME_TOKEN_HEADER", " ") + auth_config.configure_auth_from_env() + + provider = get_authorizer(Operation.RUNTIME_USE) + assert isinstance(provider, LocalJwtVerifyProvider) + assert provider._header_name == "Authorization" + assert provider._require_bearer is True + + def test_configure_then_reconfigure_clears_runtime_override(monkeypatch): """Reconfiguring without a runtime secret must drop the override.""" from agent_control_server.auth_framework import config as auth_config diff --git a/server/tests/test_runtime_token_exchange_endpoint.py b/server/tests/test_runtime_token_exchange_endpoint.py index a59e9e85..1b597a2b 100644 --- a/server/tests/test_runtime_token_exchange_endpoint.py +++ b/server/tests/test_runtime_token_exchange_endpoint.py @@ -250,6 +250,132 @@ def test_evaluation_rejects_runtime_jwt_for_wrong_target( assert response.json()["detail"] == "Runtime token target_id does not match the request." +_RUNTIME_TOKEN_HEADER = "X-Agent-Control-Runtime-Token" + + +def test_evaluation_accepts_runtime_jwt_on_custom_header( + client: TestClient, + runtime_config_enabled, +): + """Option A end-to-end: verifier on a custom header authorizes /evaluation. + + Mints a runtime token, then drives the real /api/v1/evaluation route + with the token on ``X-Agent-Control-Runtime-Token`` while the default + ``X-API-Key`` client header rides along untouched. Auth must pass; the + request then 404s on agent lookup, which is the auth-passed signal + (a 401/403 would mean the custom header was not read). + """ + stub = _StubExchangeAuthorizer(actor_id="actor-rt", scopes=("runtime.use",)) + clear_authorizers() + set_authorizer(stub) + set_authorizer( + LocalJwtVerifyProvider(secret=_TEST_SECRET, header_name=_RUNTIME_TOKEN_HEADER), + operation=Operation.RUNTIME_USE, + ) + + exchange = client.post( + "/api/v1/auth/runtime-token-exchange", + json={"target_type": "log_stream", "target_id": "ls-allowed"}, + ) + assert exchange.status_code == 200, exchange.text + token = exchange.json()["token"] + + response = client.post( + "/api/v1/evaluation", + headers={_RUNTIME_TOKEN_HEADER: token}, + json={ + "agent_name": "no-such-agent", + "step": {"type": "llm", "name": "step", "input": "hello"}, + "stage": "pre", + "target_type": "log_stream", + "target_id": "ls-allowed", + }, + ) + + assert response.status_code == 404, response.text + assert response.json()["error_code"] == "AGENT_NOT_FOUND" + + +def test_evaluation_custom_header_ignores_authorization_slot( + client: TestClient, + runtime_config_enabled, +): + """The collision fix: gateway JWT on Authorization + runtime token on the + custom header coexist. The verifier reads only the custom header, so a + (bogus) Authorization value is inert and the runtime token authorizes. + """ + stub = _StubExchangeAuthorizer(actor_id="actor-rt", scopes=("runtime.use",)) + clear_authorizers() + set_authorizer(stub) + set_authorizer( + LocalJwtVerifyProvider(secret=_TEST_SECRET, header_name=_RUNTIME_TOKEN_HEADER), + operation=Operation.RUNTIME_USE, + ) + + exchange = client.post( + "/api/v1/auth/runtime-token-exchange", + json={"target_type": "log_stream", "target_id": "ls-allowed"}, + ) + assert exchange.status_code == 200, exchange.text + token = exchange.json()["token"] + + response = client.post( + "/api/v1/evaluation", + headers={ + # Simulates the O11y gateway's downstream identity JWT. + "Authorization": "Bearer gateway-downstream-identity-jwt", + _RUNTIME_TOKEN_HEADER: token, + }, + json={ + "agent_name": "no-such-agent", + "step": {"type": "llm", "name": "step", "input": "hello"}, + "stage": "pre", + "target_type": "log_stream", + "target_id": "ls-allowed", + }, + ) + + assert response.status_code == 404, response.text + assert response.json()["error_code"] == "AGENT_NOT_FOUND" + + +def test_evaluation_custom_header_verifier_rejects_token_on_authorization( + client: TestClient, + runtime_config_enabled, +): + """When the verifier reads a custom header, a runtime token on + Authorization must not authorize (no fallback to Authorization). + """ + stub = _StubExchangeAuthorizer(actor_id="actor-rt", scopes=("runtime.use",)) + clear_authorizers() + set_authorizer(stub) + set_authorizer( + LocalJwtVerifyProvider(secret=_TEST_SECRET, header_name=_RUNTIME_TOKEN_HEADER), + operation=Operation.RUNTIME_USE, + ) + + exchange = client.post( + "/api/v1/auth/runtime-token-exchange", + json={"target_type": "log_stream", "target_id": "ls-allowed"}, + ) + assert exchange.status_code == 200, exchange.text + token = exchange.json()["token"] + + response = client.post( + "/api/v1/evaluation", + headers={"Authorization": f"Bearer {token}"}, + json={ + "agent_name": "no-such-agent", + "step": {"type": "llm", "name": "step", "input": "hello"}, + "stage": "pre", + "target_type": "log_stream", + "target_id": "ls-allowed", + }, + ) + + assert response.status_code == 401, response.text + + def test_evaluation_rejects_runtime_jwt_without_bound_target_context( client: TestClient, runtime_config_enabled, From 48b69ac22413e6e243da2374cecf861c5a5d2197 Mon Sep 17 00:00:00 2001 From: Joshua Jeon Date: Wed, 15 Jul 2026 09:41:56 -0700 Subject: [PATCH 2/2] feat(sdk): send the runtime token on the configurable header Mirror the server change in the Python SDK so Option A works end to end. AgentControlClient gains a runtime_token_header param (same env, default Authorization). The Bearer prefix is applied only on Authorization; a dedicated header carries the raw token. - _merge_runtime_headers sends the token on the configured header; _format_runtime_token owns the Bearer-prefix rule. - The runtime token stays the sole credential on an evaluation request: _AgentControlAuth suppresses X-API-Key when a runtime token is present on its dedicated header. The auto-fallback path (no token minted) and the token-exchange POST both still carry X-API-Key, so nothing becomes unauthenticated. - Blank-header handling matches the server: a whitespace-only env value falls back to the default; an explicit blank param is a hard error. - Tests cover custom-header send (raw, no Bearer; Authorization and X-API-Key both absent), env override, defaults, blank/whitespace handling, custom-header auto-fallback keeping X-API-Key, and exchange auth in custom-header mode. Co-Authored-By: Claude Opus 4.7 --- sdks/python/src/agent_control/client.py | 89 +++++++++++++-- sdks/python/tests/test_client.py | 138 ++++++++++++++++++++++++ 2 files changed, 216 insertions(+), 11 deletions(-) diff --git a/sdks/python/src/agent_control/client.py b/sdks/python/src/agent_control/client.py index fcc0c398..3e0827d4 100644 --- a/sdks/python/src/agent_control/client.py +++ b/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, diff --git a/sdks/python/tests/test_client.py b/sdks/python/tests/test_client.py index 02c9c174..ab34df53 100644 --- a/sdks/python/tests/test_client.py +++ b/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