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,25 @@ 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+ self ._runtime_token_use_bearer = (
190+ self ._runtime_token_header .lower () == _DEFAULT_RUNTIME_TOKEN_HEADER .lower ()
191+ )
143192 if runtime_token_refresh_margin_seconds < 0 :
144193 raise ValueError ("runtime_token_refresh_margin_seconds must be >= 0." )
145194 self ._runtime_token_refresh_margin_seconds = runtime_token_refresh_margin_seconds
@@ -198,7 +247,13 @@ async def __aenter__(self) -> "AgentControlClient":
198247 base_url = self .base_url ,
199248 timeout = self .timeout ,
200249 headers = self ._get_headers (),
201- auth = _AgentControlAuth (self ._api_key , self ._api_key_header ),
250+ auth = _AgentControlAuth (
251+ self ._api_key ,
252+ self ._api_key_header ,
253+ runtime_token_header = (
254+ None if self ._runtime_token_use_bearer else self ._runtime_token_header
255+ ),
256+ ),
202257 transport = self ._transport ,
203258 event_hooks = {"response" : [self ._check_server_version ]},
204259 )
@@ -295,18 +350,28 @@ async def post_runtime_evaluation(
295350
296351 return response
297352
353+ def _format_runtime_token (self , token : str ) -> str :
354+ """Format a runtime token for the configured header.
355+
356+ Prefixes ``Bearer `` only on the ``Authorization`` header (existing
357+ server contract); a dedicated runtime header carries the raw token.
358+ """
359+ if self ._runtime_token_use_bearer :
360+ return f"Bearer { token } "
361+ return token
362+
298363 def _merge_runtime_headers (
299364 self ,
300365 headers : dict [str , str ] | None ,
301366 runtime_authorization : str | None ,
302367 ) -> dict [str , str ] | None :
303- """Merge caller headers with an optional Bearer token."""
368+ """Merge caller headers with an optional runtime token header ."""
304369 if headers is None and runtime_authorization is None :
305370 return None
306371
307372 merged = dict (headers or {})
308373 if runtime_authorization is not None :
309- merged ["Authorization" ] = runtime_authorization
374+ merged [self . _runtime_token_header ] = runtime_authorization
310375 return merged
311376
312377 async def _runtime_authorization (
@@ -350,7 +415,7 @@ async def _runtime_authorization(
350415 refresh_margin_seconds = self ._runtime_token_refresh_margin_seconds ,
351416 )
352417 if cached is not None :
353- return f"Bearer { cached .token } "
418+ return self . _format_runtime_token ( cached .token )
354419
355420 exchange_lock = self ._runtime_token_cache .exchange_lock (
356421 self .base_url ,
@@ -379,7 +444,7 @@ async def _runtime_authorization(
379444 refresh_margin_seconds = self ._runtime_token_refresh_margin_seconds ,
380445 )
381446 if cached is not None :
382- return f"Bearer { cached .token } "
447+ return self . _format_runtime_token ( cached .token )
383448
384449 token = await self ._exchange_runtime_token (
385450 target_type = target_type ,
@@ -388,7 +453,7 @@ async def _runtime_authorization(
388453 )
389454 if token is None :
390455 return None
391- return f"Bearer { token } "
456+ return self . _format_runtime_token ( token )
392457
393458 async def _exchange_runtime_token (
394459 self ,
0 commit comments