diff --git a/sdks/python/src/agent_control/__init__.py b/sdks/python/src/agent_control/__init__.py index f0d07520..267f134f 100644 --- a/sdks/python/src/agent_control/__init__.py +++ b/sdks/python/src/agent_control/__init__.py @@ -448,6 +448,7 @@ def init( server_url: str | None = None, api_key: str | None = None, api_key_header: str | None = None, + runtime_token_header: str | None = None, controls_file: str | None = None, steps: list[StepSchemaDict] | None = None, conflict_mode: Literal["strict", "overwrite"] = "overwrite", @@ -488,6 +489,12 @@ def init( api_key: Optional API key for authentication (defaults to AGENT_CONTROL_API_KEY env var) api_key_header: Optional HTTP header name for API key authentication (defaults to AGENT_CONTROL_API_KEY_HEADER env var or X-API-Key) + runtime_token_header: Optional HTTP header the runtime token is sent on + for evaluation requests (defaults to AGENT_CONTROL_RUNTIME_TOKEN_HEADER + env var or Authorization). 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. controls_file: Optional explicit path to controls.yaml (auto-discovered if not provided) steps: Optional list of step schemas for registration: [{"type": "tool", "name": "search", "input_schema": {...}, "output_schema": {...}}] @@ -588,6 +595,9 @@ async def handle(message: str): state.server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000' state.api_key = api_key state.api_key_header = resolved_api_key_header + # Stored raw (may be None); AgentControlClient resolves the env-var + # fallback and blank-handling so the rule stays in one place. + state.runtime_token_header = runtime_token_header state.runtime_token_cache.clear() state.target_type = target_type state.target_id = target_id @@ -746,6 +756,7 @@ def _reset_state() -> None: state.server_url = None state.api_key = None state.api_key_header = None + state.runtime_token_header = None state.runtime_token_cache.clear() state.target_type = None state.target_id = None diff --git a/sdks/python/src/agent_control/_state.py b/sdks/python/src/agent_control/_state.py index fe6e185a..610213bd 100644 --- a/sdks/python/src/agent_control/_state.py +++ b/sdks/python/src/agent_control/_state.py @@ -27,6 +27,7 @@ def __init__(self) -> None: self.server_url: str | None = None self.api_key: str | None = None self.api_key_header: str | None = None + self.runtime_token_header: str | None = None self.runtime_token_cache = RuntimeTokenCache() # Optional target context fixed at init() time; both fields are set # together or both remain None. diff --git a/sdks/python/src/agent_control/client.py b/sdks/python/src/agent_control/client.py index fcc0c398..1e3abd67 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,9 +37,21 @@ 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 the request already carries them. + + The API key is suppressed only when the request already presents a bearer + credential on ``Authorization`` or already carries the API key on its own + header. When the runtime token rides a dedicated header, the API key is + left in place: it may be the outer credential the request needs to + authenticate at a gateway before Agent Control verifies the runtime token, + and it rides a different header so there is no collision. + """ - def __init__(self, api_key: str | None, header_name: str = "X-API-Key") -> None: + def __init__( + self, + api_key: str | None, + header_name: str = "X-API-Key", + ) -> None: self._api_key = api_key self._header_name = header_name @@ -45,9 +59,12 @@ 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 + if ( + self._api_key + and "Authorization" not in request.headers + and self._header_name not in request.headers + ): + request.headers[self._header_name] = self._api_key yield request @@ -102,6 +119,7 @@ def __init__( runtime_token_cache: RuntimeTokenCache | None = None, runtime_token_refresh_margin_seconds: int = (_DEFAULT_RUNTIME_TOKEN_REFRESH_MARGIN_SECONDS), transport: httpx.AsyncBaseTransport | None = None, + runtime_token_header: str | None = None, ): """ Initialize the client. @@ -125,6 +143,13 @@ def __init__( runtime_token_refresh_margin_seconds: Refresh cached runtime tokens before this many seconds of validity remain. transport: Optional httpx transport, primarily for tests. + runtime_token_header: HTTP header the runtime token is sent 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. """ resolved_base_url = base_url or os.environ.get( self.BASE_URL_ENV_VAR, "http://localhost:8000" @@ -140,6 +165,24 @@ 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) + # Explicit blank param is a hard error; a blank env var falls back to + # the default (mirrors the server's _resolve_runtime_token_header). + if runtime_token_header is not None and not runtime_token_header.strip(): + raise ValueError("runtime_token_header must not be blank.") + 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() + # Bearer only on Authorization; a dedicated header carries the raw token + # so it can't collide with the gateway's Authorization JWT. Must stay in + # sync with LocalJwtVerifyProvider._require_bearer on the server. + 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 +241,10 @@ 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, + ), transport=self._transport, event_hooks={"response": [self._check_server_version]}, ) @@ -295,18 +341,22 @@ async def post_runtime_evaluation( return response + def _format_runtime_token(self, token: str) -> str: + """Sole place the Bearer prefix is applied (see __init__ for the rule).""" + return f"Bearer {token}" if self._runtime_token_use_bearer else 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 +400,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 +429,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 +438,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/src/agent_control/control_decorators.py b/sdks/python/src/agent_control/control_decorators.py index 6a6d3491..66e70917 100644 --- a/sdks/python/src/agent_control/control_decorators.py +++ b/sdks/python/src/agent_control/control_decorators.py @@ -280,6 +280,7 @@ async def _evaluate( base_url=server_url, api_key=state.api_key, api_key_header=state.api_key_header, + runtime_token_header=state.runtime_token_header, runtime_token_cache=state.runtime_token_cache, ) as client: # If we have controls, use local evaluation which handles both SDK and server controls diff --git a/sdks/python/src/agent_control/evaluation.py b/sdks/python/src/agent_control/evaluation.py index 767a3e02..e79b736b 100644 --- a/sdks/python/src/agent_control/evaluation.py +++ b/sdks/python/src/agent_control/evaluation.py @@ -560,6 +560,7 @@ async def evaluate_controls( base_url=state.server_url, api_key=state.api_key, api_key_header=state.api_key_header, + runtime_token_header=state.runtime_token_header, runtime_token_cache=state.runtime_token_cache, ) as client: return await check_evaluation_with_local( diff --git a/sdks/python/tests/test_client.py b/sdks/python/tests/test_client.py index 02c9c174..d06ceefe 100644 --- a/sdks/python/tests/test_client.py +++ b/sdks/python/tests/test_client.py @@ -1081,3 +1081,174 @@ async def test_check_server_version_ignores_missing_header() -> None: # Then: no warning is emitted mock_warning.assert_not_called() + + +# --------------------------------------------------------------------------- +# HYBIM-741: configurable runtime-token header (gateway Authorization collision) +# --------------------------------------------------------------------------- + + +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_param_selects_dedicated_header() -> None: + client = AgentControlClient( + base_url="https://agent-control.test", + runtime_token_header="X-Agent-Control-Runtime-Token", + ) + assert client._runtime_token_header == "X-Agent-Control-Runtime-Token" + assert client._runtime_token_use_bearer is False + + +def test_runtime_token_header_env_var_overrides_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + 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_rejects_blank_param() -> 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_falls_back( + monkeypatch: pytest.MonkeyPatch, +) -> None: + 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_sends_raw_token_on_custom_header() -> None: + """Custom header carries the raw token; Authorization stays free for the + gateway JWT; the API key is preserved on its own header as the outer + credential (it rides X-API-Key, so there is no collision).""" + seen: dict[str, str | None] = {} + expires_at = (datetime.now(UTC) + timedelta(minutes=5)).isoformat() + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/runtime-token-exchange"): + 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"], + }, + ) + seen["runtime"] = request.headers.get("X-Agent-Control-Runtime-Token") + seen["authorization"] = request.headers.get("Authorization") + seen["api_key"] = 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 + assert seen["runtime"] == "runtime-token" + assert seen["authorization"] is None + # The API key is retained on its own header: with the runtime token on a + # dedicated header, X-API-Key can still serve as the outer gateway + # credential without colliding with either the runtime token or the + # gateway's Authorization JWT. + assert seen["api_key"] == "test-key" + + +@pytest.mark.asyncio +async def test_runtime_evaluation_default_sends_bearer_on_authorization() -> None: + """Default (unset) behavior is unchanged: Bearer token on Authorization.""" + seen: dict[str, str | None] = {} + expires_at = (datetime.now(UTC) + timedelta(minutes=5)).isoformat() + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/runtime-token-exchange"): + return httpx.Response( + 200, + json={ + "token": "runtime-token", + "expires_at": expires_at, + "target_type": "log_stream", + "target_id": "ls-1", + "scopes": ["runtime.use"], + }, + ) + seen["authorization"] = request.headers.get("Authorization") + 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", + transport=transport, + ) as client: + await client.post_runtime_evaluation( + json={"target_type": "log_stream", "target_id": "ls-1"}, + target_type="log_stream", + target_id="ls-1", + ) + + assert seen["authorization"] == "Bearer runtime-token" + + +@pytest.mark.asyncio +async def test_runtime_evaluation_custom_header_fallback_keeps_api_key() -> 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 authenticate the evaluation request.""" + exchange_calls = 0 + seen: dict[str, str | None] = {} + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal exchange_calls + if request.url.path.endswith("/runtime-token-exchange"): + exchange_calls += 1 + return httpx.Response(503, json={"detail": "runtime auth disabled"}) + seen["runtime"] = request.headers.get("X-Agent-Control-Runtime-Token") + seen["api_key"] = request.headers.get("X-API-Key") + return httpx.Response(200, json={"is_safe": True, "confidence": 1.0}) + + transport = httpx.MockTransport(handler) + async with AgentControlClient( + base_url="https://agent-control.test", + api_key="test-key", + runtime_auth_mode="auto", + runtime_token_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 + assert seen["runtime"] is None + assert seen["api_key"] == "test-key" diff --git a/sdks/python/tests/test_init_validation.py b/sdks/python/tests/test_init_validation.py index c0ecea3d..d4054b27 100644 --- a/sdks/python/tests/test_init_validation.py +++ b/sdks/python/tests/test_init_validation.py @@ -1,11 +1,16 @@ """Validation tests for agent_control.init().""" -import agent_control +from unittest.mock import AsyncMock, patch +from uuid import uuid4 + import pytest from agent_control_models import ControlMatch as ModelControlMatch from agent_control_models import ControlScope as ModelControlScope from agent_control_models import EvaluatorResult as ModelEvaluatorResult +import agent_control +from agent_control._state import state + def test_init_rejects_invalid_agent_name() -> None: with pytest.raises(ValueError, match="at least 10 characters"): @@ -48,3 +53,50 @@ def test_init_exports_control_match() -> None: def test_init_exports_evaluator_result() -> None: assert agent_control.EvaluatorResult is ModelEvaluatorResult assert "EvaluatorResult" in agent_control.__all__ + + +def test_init_stores_runtime_token_header_in_state() -> None: + """HYBIM-741: init() retains runtime_token_header in session state so the + evaluation clients it creates ride the configured header (not only the + process-wide env var).""" + health_check_mock = AsyncMock(return_value={"status": "healthy"}) + register_agent_mock = AsyncMock(return_value={"created": True, "controls": []}) + try: + with patch( + "agent_control.__init__.AgentControlClient.health_check", + new=health_check_mock, + ), patch( + "agent_control.__init__.agents.register_agent", + new=register_agent_mock, + ): + agent_control.init( + agent_name=f"agent-{uuid4().hex[:12]}", + runtime_token_header="X-Agent-Control-Runtime-Token", + policy_refresh_interval_seconds=0, + ) + + assert state.runtime_token_header == "X-Agent-Control-Runtime-Token" + finally: + agent_control._reset_state() + + +def test_init_defaults_runtime_token_header_to_none() -> None: + """Unset: state carries None so the client falls back to env/default.""" + health_check_mock = AsyncMock(return_value={"status": "healthy"}) + register_agent_mock = AsyncMock(return_value={"created": True, "controls": []}) + try: + with patch( + "agent_control.__init__.AgentControlClient.health_check", + new=health_check_mock, + ), patch( + "agent_control.__init__.agents.register_agent", + new=register_agent_mock, + ): + agent_control.init( + agent_name=f"agent-{uuid4().hex[:12]}", + policy_refresh_interval_seconds=0, + ) + + assert state.runtime_token_header is None + finally: + agent_control._reset_state() diff --git a/server/src/agent_control_server/auth_framework/config.py b/server/src/agent_control_server/auth_framework/config.py index 06246a46..16d0a566 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,28 @@ 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: + """Header the runtime JWT verifier reads the token from (default ``Authorization``). + + Behind a gateway that overwrites ``Authorization`` with its own identity + JWT, set a dedicated header so the two tokens don't collide. Blank falls + back to the default. + """ + 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..786b63c5 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,15 @@ """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 the runtime 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 is configurable so the server can sit behind a gateway that +reserves ``Authorization`` for its own identity JWT (point the verifier at +a dedicated header to avoid the collision). """ from __future__ import annotations @@ -19,14 +23,28 @@ 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() + # Bearer only on Authorization; a dedicated header carries the raw token + # so it can't collide with the gateway's Authorization JWT. Must stay in + # sync with AgentControlClient._runtime_token_use_bearer in the SDK. + self._require_bearer = self._header_name.lower() == "authorization" async def authorize( self, @@ -79,18 +97,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 value. + 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..b05cce57 100644 --- a/server/tests/test_auth_framework.py +++ b/server/tests/test_auth_framework.py @@ -1764,3 +1764,134 @@ async def test_teardown_auth_clears_registry(): assert not _operation_authorizers with pytest.raises(RuntimeError, match="No RequestAuthorizer"): get_authorizer(Operation.CONTROL_BINDINGS_WRITE) + + +# --------------------------------------------------------------------------- +# HYBIM-741: configurable runtime-token header (gateway Authorization collision) +# --------------------------------------------------------------------------- + + +def _mint_runtime_use_token(): + from agent_control_server.auth_framework.runtime_token import mint_runtime_token + + token, _ = mint_runtime_token( + namespace_key="default", + actor_id="actor-1", + target_type="log_stream", + target_id="ls-1", + scopes=("runtime.use",), + secret=_TEST_SECRET, + ttl_seconds=60, + ) + return token + + +@pytest.mark.asyncio +async def test_local_jwt_default_reads_authorization_bearer(): + """Default behavior unchanged: token read as Bearer from Authorization.""" + provider = LocalJwtVerifyProvider(secret=_TEST_SECRET) + token = _mint_runtime_use_token() + principal = await provider.authorize( + _build_request(headers={"Authorization": f"Bearer {token}"}), + Operation.RUNTIME_USE, + context={"target_type": "log_stream", "target_id": "ls-1"}, + ) + assert principal.target_type == "log_stream" + assert principal.target_id == "ls-1" + + +@pytest.mark.asyncio +async def test_local_jwt_default_rejects_raw_token_on_authorization(): + """On Authorization the Bearer scheme stays mandatory (back-compat).""" + provider = LocalJwtVerifyProvider(secret=_TEST_SECRET) + token = _mint_runtime_use_token() + with pytest.raises(AuthenticationError): + await provider.authorize( + _build_request(headers={"Authorization": token}), + Operation.RUNTIME_USE, + context={"target_type": "log_stream", "target_id": "ls-1"}, + ) + + +@pytest.mark.asyncio +async def test_local_jwt_custom_header_reads_raw_token(): + """On a dedicated header the raw token is accepted and Authorization is + left free for the gateway's own identity JWT (no collision).""" + provider = LocalJwtVerifyProvider( + secret=_TEST_SECRET, header_name="X-Agent-Control-Runtime-Token" + ) + token = _mint_runtime_use_token() + principal = await provider.authorize( + _build_request( + headers={ + "X-Agent-Control-Runtime-Token": token, + "Authorization": "Bearer gateway-identity-jwt", + } + ), + Operation.RUNTIME_USE, + context={"target_type": "log_stream", "target_id": "ls-1"}, + ) + assert principal.target_id == "ls-1" + + +@pytest.mark.asyncio +async def test_local_jwt_custom_header_also_accepts_bearer_prefix(): + provider = LocalJwtVerifyProvider( + secret=_TEST_SECRET, header_name="X-Agent-Control-Runtime-Token" + ) + token = _mint_runtime_use_token() + principal = await provider.authorize( + _build_request(headers={"X-Agent-Control-Runtime-Token": f"Bearer {token}"}), + Operation.RUNTIME_USE, + context={"target_type": "log_stream", "target_id": "ls-1"}, + ) + assert principal.target_id == "ls-1" + + +@pytest.mark.asyncio +async def test_local_jwt_custom_header_missing_reports_that_header(): + provider = LocalJwtVerifyProvider( + secret=_TEST_SECRET, header_name="X-Agent-Control-Runtime-Token" + ) + with pytest.raises(AuthenticationError, match="X-Agent-Control-Runtime-Token"): + await provider.authorize( + _build_request(headers={"Authorization": "Bearer gateway-jwt"}), + Operation.RUNTIME_USE, + context={"target_type": "log_stream", "target_id": "ls-1"}, + ) + + +def test_local_jwt_rejects_blank_header_name(): + with pytest.raises(ValueError, match="header_name"): + LocalJwtVerifyProvider(secret=_TEST_SECRET, header_name=" ") + + +# --------------------------------------------------------------------------- +# HYBIM-741: AGENT_CONTROL_RUNTIME_TOKEN_HEADER env resolution (config wiring) +# --------------------------------------------------------------------------- + + +def test_resolve_runtime_token_header_defaults_to_authorization(monkeypatch): + from agent_control_server.auth_framework import config as auth_config + + monkeypatch.delenv("AGENT_CONTROL_RUNTIME_TOKEN_HEADER", raising=False) + assert auth_config._resolve_runtime_token_header() == "Authorization" + + +def test_resolve_runtime_token_header_reads_env(monkeypatch): + from agent_control_server.auth_framework import config as auth_config + + monkeypatch.setenv( + "AGENT_CONTROL_RUNTIME_TOKEN_HEADER", " X-Agent-Control-Runtime-Token " + ) + # Value is honored and trimmed. + assert ( + auth_config._resolve_runtime_token_header() == "X-Agent-Control-Runtime-Token" + ) + + +def test_resolve_runtime_token_header_blank_env_falls_back(monkeypatch): + from agent_control_server.auth_framework import config as auth_config + + monkeypatch.setenv("AGENT_CONTROL_RUNTIME_TOKEN_HEADER", " ") + assert auth_config._resolve_runtime_token_header() == "Authorization" diff --git a/server/tests/test_runtime_token_exchange_endpoint.py b/server/tests/test_runtime_token_exchange_endpoint.py index a59e9e85..2164e43a 100644 --- a/server/tests/test_runtime_token_exchange_endpoint.py +++ b/server/tests/test_runtime_token_exchange_endpoint.py @@ -9,6 +9,7 @@ from __future__ import annotations import logging +import uuid from datetime import UTC, datetime, timedelta import pytest @@ -250,6 +251,110 @@ def test_evaluation_rejects_runtime_jwt_for_wrong_target( assert response.json()["detail"] == "Runtime token target_id does not match the request." +def test_evaluation_accepts_runtime_jwt_on_configured_header_with_gateway_authorization( + client: TestClient, + runtime_config_enabled, +): + """HYBIM-741: end-to-end through /api/v1/evaluation with the runtime token + on a dedicated header while Authorization carries an (ignored) gateway JWT. + + Exercises the config wiring (LocalJwtVerifyProvider bound to a custom + header) and Operation.RUNTIME_USE routing together, not just the provider + in isolation: the runtime token rides X-Agent-Control-Runtime-Token and is + accepted, while the Authorization value is left for the gateway and does + not interfere. + """ + agent_name = f"agent-{uuid.uuid4().hex[:12]}" + register = client.post( + "/api/v1/agents/initAgent", + json={ + "agent": { + "agent_name": agent_name, + "agent_description": "test agent", + "agent_version": "1.0", + }, + "steps": [], + }, + ) + assert register.status_code == 200, register.text + + stub = _StubExchangeAuthorizer(actor_id="actor-rt", scopes=("runtime.use",)) + clear_authorizers() + set_authorizer(stub) + set_authorizer( + LocalJwtVerifyProvider( + secret=_TEST_SECRET, header_name="X-Agent-Control-Runtime-Token" + ), + 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={ + "X-Agent-Control-Runtime-Token": token, + "Authorization": "Bearer gateway-identity-jwt", + }, + json={ + "agent_name": agent_name, + "step": {"type": "llm", "name": "step", "input": "hello"}, + "stage": "pre", + "target_type": "log_stream", + "target_id": "ls-allowed", + }, + ) + + # The runtime token authorizes the request for its bound target; with no + # controls bound to that target the evaluation resolves to safe. + assert response.status_code == 200, response.text + assert response.json()["is_safe"] is True + + +def test_evaluation_rejects_runtime_jwt_on_wrong_header_when_custom_configured( + client: TestClient, + runtime_config_enabled, +): + """With a dedicated header configured, a token presented on Authorization + is ignored (it can't be smuggled past the gateway boundary): the verifier + reads only the configured header, so the request fails to authenticate.""" + stub = _StubExchangeAuthorizer(actor_id="actor-rt", scopes=("runtime.use",)) + clear_authorizers() + set_authorizer(stub) + set_authorizer( + LocalJwtVerifyProvider( + secret=_TEST_SECRET, header_name="X-Agent-Control-Runtime-Token" + ), + 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": "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,