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 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,