2020_logger = logging .getLogger (__name__ )
2121
2222_RUNTIME_AUTH_MODE_ENV_VAR = "AGENT_CONTROL_RUNTIME_AUTH_MODE"
23+ _RUNTIME_TOKEN_HEADER_ENV_VAR = "AGENT_CONTROL_RUNTIME_TOKEN_HEADER"
24+ _DEFAULT_RUNTIME_TOKEN_HEADER = "Authorization"
2325_DEFAULT_RUNTIME_TOKEN_REFRESH_MARGIN_SECONDS = 30
2426_AUTO_RUNTIME_TOKEN_FALLBACK_STATUSES = {404 , 500 , 502 , 503 , 504 }
2527_GLOBAL_RUNTIME_TOKEN_FALLBACK_STATUSES = {404 }
@@ -35,19 +37,39 @@ def _runtime_cache_identity(api_key: str | None, api_key_header: str) -> str:
3537
3638
3739class _AgentControlAuth (httpx .Auth ):
38- """Attach local API-key credentials unless a request already has Bearer auth."""
40+ """Attach local API-key credentials unless a request already carries a token.
3941
40- def __init__ (self , api_key : str | None , header_name : str = "X-API-Key" ) -> None :
42+ The API key is suppressed when the request already presents a bearer
43+ credential on ``Authorization`` or a runtime token on its dedicated
44+ header, so a runtime-authenticated evaluation carries a single
45+ credential regardless of which header the runtime token rides.
46+ """
47+
48+ def __init__ (
49+ self ,
50+ api_key : str | None ,
51+ header_name : str = "X-API-Key" ,
52+ runtime_token_header : str | None = None ,
53+ ) -> None :
4154 self ._api_key = api_key
4255 self ._header_name = header_name
56+ self ._runtime_token_header = runtime_token_header
4357
4458 def auth_flow (
4559 self ,
4660 request : httpx .Request ,
4761 ) -> Generator [httpx .Request , httpx .Response , None ]:
48- if self ._api_key and "Authorization" not in request .headers :
49- if self ._header_name not in request .headers :
50- request .headers [self ._header_name ] = self ._api_key
62+ runtime_token_on_dedicated_header = (
63+ self ._runtime_token_header is not None
64+ and self ._runtime_token_header in request .headers
65+ )
66+ if (
67+ self ._api_key
68+ and "Authorization" not in request .headers
69+ and not runtime_token_on_dedicated_header
70+ and self ._header_name not in request .headers
71+ ):
72+ request .headers [self ._header_name ] = self ._api_key
5173 yield request
5274
5375
@@ -99,6 +121,7 @@ def __init__(
99121 api_key : str | None = None ,
100122 api_key_header : str | None = None ,
101123 runtime_auth_mode : RuntimeAuthMode | str | None = None ,
124+ runtime_token_header : str | None = None ,
102125 runtime_token_cache : RuntimeTokenCache | None = None ,
103126 runtime_token_refresh_margin_seconds : int = (_DEFAULT_RUNTIME_TOKEN_REFRESH_MARGIN_SECONDS ),
104127 transport : httpx .AsyncBaseTransport | None = None ,
@@ -121,6 +144,13 @@ def __init__(
121144 request auth when the exchange endpoint is unavailable. ``jwt``
122145 requires a successful exchange. ``api_key`` and ``none`` keep
123146 evaluation requests on the normal request-auth path.
147+ runtime_token_header: HTTP header name to send the runtime token
148+ on. Defaults to ``Authorization``; the
149+ AGENT_CONTROL_RUNTIME_TOKEN_HEADER environment variable
150+ overrides the default. Point this at a dedicated header (e.g.
151+ ``X-Agent-Control-Runtime-Token``) when the server runs behind
152+ a gateway that reserves ``Authorization`` for its own identity
153+ JWT. The server must be configured to read the same header.
124154 runtime_token_cache: Optional cache shared across client instances.
125155 runtime_token_refresh_margin_seconds: Refresh cached runtime tokens
126156 before this many seconds of validity remain.
@@ -140,6 +170,27 @@ def __init__(
140170 self ._runtime_cache_identity = _runtime_cache_identity (self ._api_key , self ._api_key_header )
141171 configured_runtime_mode = runtime_auth_mode or os .environ .get (_RUNTIME_AUTH_MODE_ENV_VAR )
142172 self ._runtime_auth_mode = normalize_runtime_auth_mode (configured_runtime_mode )
173+ if runtime_token_header is not None and not runtime_token_header .strip ():
174+ raise ValueError ("runtime_token_header must not be blank." )
175+ # A whitespace-only env var falls back to the default header, matching
176+ # the server's _resolve_runtime_token_header(); an explicit blank
177+ # *param* above is a hard error.
178+ env_runtime_token_header = os .environ .get (_RUNTIME_TOKEN_HEADER_ENV_VAR )
179+ if env_runtime_token_header is not None and not env_runtime_token_header .strip ():
180+ env_runtime_token_header = None
181+ self ._runtime_token_header = (
182+ runtime_token_header
183+ or env_runtime_token_header
184+ or _DEFAULT_RUNTIME_TOKEN_HEADER
185+ ).strip ()
186+ # On Authorization the server keeps the ``Bearer`` scheme prefix
187+ # (existing contract); a dedicated header carries the raw token so the
188+ # value no longer collides with a gateway's Authorization JWT.
189+ # Mirror of LocalJwtVerifyProvider._require_bearer on the server — the
190+ # send format and the verifier's expectation must stay in sync.
191+ self ._runtime_token_use_bearer = (
192+ self ._runtime_token_header .lower () == _DEFAULT_RUNTIME_TOKEN_HEADER .lower ()
193+ )
143194 if runtime_token_refresh_margin_seconds < 0 :
144195 raise ValueError ("runtime_token_refresh_margin_seconds must be >= 0." )
145196 self ._runtime_token_refresh_margin_seconds = runtime_token_refresh_margin_seconds
@@ -198,7 +249,13 @@ async def __aenter__(self) -> "AgentControlClient":
198249 base_url = self .base_url ,
199250 timeout = self .timeout ,
200251 headers = self ._get_headers (),
201- auth = _AgentControlAuth (self ._api_key , self ._api_key_header ),
252+ auth = _AgentControlAuth (
253+ self ._api_key ,
254+ self ._api_key_header ,
255+ runtime_token_header = (
256+ None if self ._runtime_token_use_bearer else self ._runtime_token_header
257+ ),
258+ ),
202259 transport = self ._transport ,
203260 event_hooks = {"response" : [self ._check_server_version ]},
204261 )
@@ -295,18 +352,28 @@ async def post_runtime_evaluation(
295352
296353 return response
297354
355+ def _format_runtime_token (self , token : str ) -> str :
356+ """Format a runtime token for the configured header.
357+
358+ Prefixes ``Bearer `` only on the ``Authorization`` header (existing
359+ server contract); a dedicated runtime header carries the raw token.
360+ """
361+ if self ._runtime_token_use_bearer :
362+ return f"Bearer { token } "
363+ return token
364+
298365 def _merge_runtime_headers (
299366 self ,
300367 headers : dict [str , str ] | None ,
301368 runtime_authorization : str | None ,
302369 ) -> dict [str , str ] | None :
303- """Merge caller headers with an optional Bearer token."""
370+ """Merge caller headers with an optional runtime token header ."""
304371 if headers is None and runtime_authorization is None :
305372 return None
306373
307374 merged = dict (headers or {})
308375 if runtime_authorization is not None :
309- merged ["Authorization" ] = runtime_authorization
376+ merged [self . _runtime_token_header ] = runtime_authorization
310377 return merged
311378
312379 async def _runtime_authorization (
@@ -350,7 +417,7 @@ async def _runtime_authorization(
350417 refresh_margin_seconds = self ._runtime_token_refresh_margin_seconds ,
351418 )
352419 if cached is not None :
353- return f"Bearer { cached .token } "
420+ return self . _format_runtime_token ( cached .token )
354421
355422 exchange_lock = self ._runtime_token_cache .exchange_lock (
356423 self .base_url ,
@@ -379,7 +446,7 @@ async def _runtime_authorization(
379446 refresh_margin_seconds = self ._runtime_token_refresh_margin_seconds ,
380447 )
381448 if cached is not None :
382- return f"Bearer { cached .token } "
449+ return self . _format_runtime_token ( cached .token )
383450
384451 token = await self ._exchange_runtime_token (
385452 target_type = target_type ,
@@ -388,7 +455,7 @@ async def _runtime_authorization(
388455 )
389456 if token is None :
390457 return None
391- return f"Bearer { token } "
458+ return self . _format_runtime_token ( token )
392459
393460 async def _exchange_runtime_token (
394461 self ,
0 commit comments