Skip to content

Commit 28f58f7

Browse files
jjeonsplunkclaude
andcommitted
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 <noreply@anthropic.com>
1 parent 6e31491 commit 28f58f7

2 files changed

Lines changed: 214 additions & 11 deletions

File tree

sdks/python/src/agent_control/client.py

Lines changed: 76 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
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

3739
class _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,

sdks/python/tests/test_client.py

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,101 @@ def handler(request: httpx.Request) -> httpx.Response:
226226
assert evaluation_api_key_headers == [None, None]
227227

228228

229+
@pytest.mark.asyncio
230+
async def test_runtime_evaluation_sends_token_on_custom_header() -> None:
231+
"""With a custom runtime header, the raw token rides that header (no Bearer)
232+
and Authorization is left free for a gateway's own JWT."""
233+
evaluation_runtime_headers: list[str | None] = []
234+
evaluation_authorization_headers: list[str | None] = []
235+
evaluation_api_key_headers: list[str | None] = []
236+
expires_at = (datetime.now(UTC) + timedelta(minutes=5)).isoformat()
237+
238+
def handler(request: httpx.Request) -> httpx.Response:
239+
if request.url.path == "/api/v1/auth/runtime-token-exchange":
240+
# The exchange itself must still authenticate with the API key,
241+
# even when the runtime token later rides a dedicated header.
242+
assert request.headers.get("X-API-Key") == "test-key"
243+
return httpx.Response(
244+
200,
245+
json={
246+
"token": "runtime-token",
247+
"expires_at": expires_at,
248+
"target_type": "log_stream",
249+
"target_id": "ls-1",
250+
"scopes": ["runtime.use"],
251+
},
252+
)
253+
254+
evaluation_runtime_headers.append(
255+
request.headers.get("X-Agent-Control-Runtime-Token")
256+
)
257+
evaluation_authorization_headers.append(request.headers.get("Authorization"))
258+
evaluation_api_key_headers.append(request.headers.get("X-API-Key"))
259+
return httpx.Response(200, json={"is_safe": True, "confidence": 1.0})
260+
261+
transport = httpx.MockTransport(handler)
262+
263+
async with AgentControlClient(
264+
base_url="https://agent-control.test",
265+
api_key="test-key",
266+
runtime_auth_mode="jwt",
267+
runtime_token_header="X-Agent-Control-Runtime-Token",
268+
transport=transport,
269+
) as client:
270+
response = await client.post_runtime_evaluation(
271+
json={"target_type": "log_stream", "target_id": "ls-1"},
272+
target_type="log_stream",
273+
target_id="ls-1",
274+
)
275+
276+
assert response.status_code == 200
277+
# Raw token on the dedicated header, no Bearer prefix.
278+
assert evaluation_runtime_headers == ["runtime-token"]
279+
# Authorization untouched, free for the gateway JWT.
280+
assert evaluation_authorization_headers == [None]
281+
# The runtime token is the sole credential: the API key must not ride
282+
# along just because Authorization happens to be free in custom-header mode.
283+
assert evaluation_api_key_headers == [None]
284+
285+
286+
@pytest.mark.asyncio
287+
async def test_runtime_token_header_env_var_overrides_default(
288+
monkeypatch: pytest.MonkeyPatch,
289+
) -> None:
290+
"""AGENT_CONTROL_RUNTIME_TOKEN_HEADER selects the send header."""
291+
monkeypatch.setenv(
292+
"AGENT_CONTROL_RUNTIME_TOKEN_HEADER", "X-Agent-Control-Runtime-Token"
293+
)
294+
client = AgentControlClient(base_url="https://agent-control.test")
295+
assert client._runtime_token_header == "X-Agent-Control-Runtime-Token"
296+
assert client._runtime_token_use_bearer is False
297+
298+
299+
def test_runtime_token_header_defaults_to_authorization() -> None:
300+
client = AgentControlClient(base_url="https://agent-control.test")
301+
assert client._runtime_token_header == "Authorization"
302+
assert client._runtime_token_use_bearer is True
303+
304+
305+
def test_runtime_token_header_rejects_blank() -> None:
306+
with pytest.raises(ValueError, match="runtime_token_header"):
307+
AgentControlClient(
308+
base_url="https://agent-control.test", runtime_token_header=" "
309+
)
310+
311+
312+
def test_runtime_token_header_whitespace_env_defaults_to_authorization(
313+
monkeypatch: pytest.MonkeyPatch,
314+
) -> None:
315+
"""A whitespace-only env var falls back to Authorization, matching the
316+
server's _resolve_runtime_token_header(); only an explicit blank param
317+
is a hard error."""
318+
monkeypatch.setenv("AGENT_CONTROL_RUNTIME_TOKEN_HEADER", " ")
319+
client = AgentControlClient(base_url="https://agent-control.test")
320+
assert client._runtime_token_header == "Authorization"
321+
assert client._runtime_token_use_bearer is True
322+
323+
229324
@pytest.mark.asyncio
230325
async def test_runtime_evaluation_single_flights_cold_cache_exchange() -> None:
231326
exchange_calls = 0
@@ -456,6 +551,49 @@ def handler(request: httpx.Request) -> httpx.Response:
456551
assert evaluation_api_key_headers == ["test-key", "test-key"]
457552

458553

554+
@pytest.mark.asyncio
555+
async def test_runtime_evaluation_auto_fallback_keeps_api_key_on_custom_header() -> None:
556+
"""Custom-header mode must still fall back to the API key when the
557+
exchange is unavailable: with no runtime token minted, the dedicated
558+
header is absent and X-API-Key must ride the evaluation request."""
559+
exchange_calls = 0
560+
evaluation_api_key_headers: list[str | None] = []
561+
evaluation_runtime_headers: list[str | None] = []
562+
563+
def handler(request: httpx.Request) -> httpx.Response:
564+
nonlocal exchange_calls
565+
if request.url.path == "/api/v1/auth/runtime-token-exchange":
566+
exchange_calls += 1
567+
return httpx.Response(503, json={"detail": "runtime auth disabled"})
568+
569+
evaluation_api_key_headers.append(request.headers.get("X-API-Key"))
570+
evaluation_runtime_headers.append(
571+
request.headers.get("X-Agent-Control-Runtime-Token")
572+
)
573+
return httpx.Response(200, json={"is_safe": True, "confidence": 1.0})
574+
575+
transport = httpx.MockTransport(handler)
576+
577+
async with AgentControlClient(
578+
base_url="https://agent-control.test",
579+
api_key="test-key",
580+
runtime_auth_mode="auto",
581+
runtime_token_header="X-Agent-Control-Runtime-Token",
582+
transport=transport,
583+
) as client:
584+
response = await client.post_runtime_evaluation(
585+
json={"target_type": "log_stream", "target_id": "ls-1"},
586+
target_type="log_stream",
587+
target_id="ls-1",
588+
)
589+
assert response.status_code == 200
590+
591+
assert exchange_calls == 1
592+
# No token minted -> dedicated header absent -> API key authenticates.
593+
assert evaluation_runtime_headers == [None]
594+
assert evaluation_api_key_headers == ["test-key"]
595+
596+
459597
@pytest.mark.asyncio
460598
async def test_runtime_evaluation_auto_404_fallback_recovers_after_ttl() -> None:
461599
exchange_calls = 0

0 commit comments

Comments
 (0)