feat(sdk): add more out of box controls - #257
Open
karansohi wants to merge 5 commits into
Open
Conversation
Added 6 luna out of the box, verified and tested. Added 2 more controls, 1 regex ( for prompt injection ) & one json control.
karansohi
changed the title
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
…orization collision [HYBIM-866] (#258) _Recreated from #253 on an upstream branch (not a fork) so CI secrets (SPEAKEASY_API_KEY) are available. Same commits, signed. Original PR: #253._ --- ## Problem When Agent Control runs behind the O11y gateway, the gateway overwrites `Authorization` with its own downstream identity JWT — clobbering AC's runtime-eval token on the hot path (HYBIM-866, under epic HYBIM-741). The runtime token and the gateway's identity JWT both want the `Authorization` header. ## Fix Make the runtime token ride a **configurable header** on both sides, selected by `AGENT_CONTROL_RUNTIME_TOKEN_HEADER` (default `Authorization`). Behind the gateway, point both sides at a dedicated header (e.g. `X-Agent-Control-Runtime-Token`); the runtime token rides that header while the gateway keeps `Authorization` for its identity JWT — no collision. - **`Authorization`**: keeps the mandatory `Bearer` scheme (existing contract). - **Dedicated header**: carries the **raw** token (no `Bearer` prefix). - **Default unchanged**: with the env unset, behavior is byte-identical to today. ### Server (verify side) - `auth_framework/providers/local_jwt.py`: `LocalJwtVerifyProvider` takes a `header_name` (default `Authorization`) and reads the token from it. `Bearer` required only on `Authorization`; raw token accepted on a dedicated header. - `auth_framework/config.py`: `_resolve_runtime_token_header()` reads `AGENT_CONTROL_RUNTIME_TOKEN_HEADER` (blank → default), passed into the provider when runtime mode is `jwt`. ### SDK (send side) - `sdks/python/.../client.py`: `AgentControlClient` gains a `runtime_token_header` param (+ same env var). Sends the runtime token on the configured header — raw on a dedicated header, `Bearer` on `Authorization` (single `_format_runtime_token` helper is the sole authority for that rule). - The API key is **preserved as the outer gateway credential**: when the runtime token rides a dedicated header, the API key still rides its own header (`X-API-Key`), so a request can authenticate at the gateway while the runtime JWT is verified by Agent Control. The existing same-header guard prevents any collision. - High-level SDK: `agent_control.init()` accepts `runtime_token_header`, stores it in session state, and threads it into the evaluation clients (`evaluation.py`, `control_decorators.py`); it is cleared on reset. ## Configuration (behind the gateway) ``` # server AGENT_CONTROL_RUNTIME_AUTH_MODE=jwt AGENT_CONTROL_RUNTIME_TOKEN_SECRET=<secret> AGENT_CONTROL_RUNTIME_TOKEN_HEADER=X-Agent-Control-Runtime-Token # SDK (must match the server header) AgentControlClient(..., runtime_token_header="X-Agent-Control-Runtime-Token") # or agent_control.init(..., runtime_token_header="X-Agent-Control-Runtime-Token") # or AGENT_CONTROL_RUNTIME_TOKEN_HEADER=X-Agent-Control-Runtime-Token ``` ## Tests - **Server** (`test_auth_framework.py`): default Bearer path; default rejects raw on `Authorization`; dedicated header reads raw token and coexists with a gateway `Authorization` JWT; Bearer also accepted on dedicated header; missing-header error names the configured header; blank `header_name` rejected; env resolver (unset → default, set → trimmed, whitespace → default). - **Server, app-level** (`test_runtime_token_exchange_endpoint.py`): end-to-end through `/api/v1/evaluation` exercising config wiring + `Operation.RUNTIME_USE` routing — runtime token on a dedicated header with an outer `Authorization` gateway JWT is accepted; a token presented on `Authorization` is rejected (401) when a custom header is configured. - **SDK** (`test_client.py`): header resolution (param/env/default, blank rejected, whitespace-env fallback); raw token on dedicated header with `Authorization` free while the API key is preserved on its own header; default sends `Bearer` on `Authorization`; auto-mode fallback keeps the API key when the exchange is unavailable. - **SDK, high-level** (`test_init_validation.py`): `init()` stores `runtime_token_header` in session state; defaults to `None` when unset; validates the header up front (blank and bad field-name rejected before state is mutated); a positional call through `target_id` proves the new param is appended last and does not shift any existing slot (`controls_file` onward). ## Live validation on lab0 (through the real O11y gateway) Validated the branch on lab0 behind the actual O11y api-gateway (the gateway that owns `Authorization` for its own identity JWT — the real collision scenario). - **Deploy:** built the branch server image, pushed it to a personal `docker-test.repo.splunkdev.net/user-<name>/agent-control:<tag>` namespace (a path a personal Artifactory token can write and lab0 can pull), then `kubectl set image` the lab0 `agent-control` deploy at it with `AGENT_CONTROL_RUNTIME_TOKEN_HEADER=X-Agent-Control-Runtime-Token`. Rolled back to the released image afterward. - **Full flow returns 200, with a real org and log_stream.** `runtime-token-exchange` returns 200 and mints a real server-issued runtime JWT for the log_stream; the evaluation call carries that token on the **custom header** through the gateway and returns **200** with a real result (`is_safe: true`). - **The collision is handled.** The same token presented on `Authorization` is rejected with 401 ("Missing X-Agent-Control-Runtime-Token"), because the server reads only the dedicated header. So even when the gateway overwrites `Authorization` with its own identity JWT, the runtime JWT survives on `X-Agent-Control-Runtime-Token`. - **Correction from an earlier draft of this PR:** an initial run hit `502` (upstream `422`) on the exchange, which looked like a capability/provisioning gap. That was a test-input error, not a real gap: the request used a made-up `target_id` that was not a real log_stream, so the exchange could not resolve it. With a real log_stream id the whole flow returns 200. Nothing is blocked on provisioning. - **Not exercised in this run:** a control actually firing (steer/deny) and control spans landing in Galileo (no control was bound to the log_stream, so the eval ran clean). That is separate from the header change and was already shown on the devstack. ## Security notes - **Header isolation**: with a dedicated header configured, the verifier reads only that header — a token presented on `Authorization` is ignored, so it can't be smuggled past the gateway boundary. - **No token leakage**: auth error messages reference the header name only, never the token value. - Signature / scope / target-binding checks (`verify_runtime_token`) are unchanged. ## Backward compatibility Safe. No backward compatibility issue. The change is default-preserving. `runtime_token_header` defaults to `Authorization` on both the SDK and server, and an unset env var falls back to that default. With no config, behavior is byte-identical to today: the SDK still sends `Bearer <token>` on `Authorization`, and the server still requires the Bearer scheme there. All four default-path branches were traced to confirm. One behavior change on the default path, and it's an improvement: a malformed `Authorization: Bearer ` with only trailing whitespace used to return an empty token and fail deep in signature verification. Now it's rejected up front with a clean `AUTH_MISSING_KEY`. No valid request changes. ### Operational notes (expected tradeoffs, not design risks) - **Two-sided header contract.** Setting `AGENT_CONTROL_RUNTIME_TOKEN_HEADER` on only one side breaks runtime auth (401 on every eval). This is intrinsic to any configurable-header feature, not a regression, and it fails closed: a mismatch denies auth, never bypasses it. Neither side can validate the other's value since they run as separate processes, so the intended safeguard is a startup log of the resolved header on both sides, making a mismatch a quick log diff rather than a debugging session. - **API-key retention on a dedicated header.** `_AgentControlAuth` now leaves the API key in place when the runtime token rides a dedicated header, so the key can act as the outer gateway credential. This is deliberate and required by the two-credential model: the API key authenticates at the gateway, and the runtime JWT is verified by Agent Control. The default path is unchanged. It is called out only because it is the one change on every request's auth path, which makes it the right spot to focus review. ## Notes / scope - Backwards compatible: unset env → identical behavior. No change to the API-key or `none` runtime modes. - Design follows the O11y api-service precedent: support both auth methods, gateway stays neutral, opt-in, default preserved. - Server-side tests require the repo's Postgres test fixture (run in CI); the SDK suite runs standalone. - Validated end-to-end on a devstack: with the SDK sending the token on `X-Agent-Control-Runtime-Token` and the server reading the same header, runtime JWT exchange, control steering, and control-span ingestion all worked. This confirms the SDK and server agree on the custom header. It does not yet exercise the gateway-collision path (the devstack has no O11y gateway overwriting `Authorization`). - Gateway-collision validation on lab0 is done and green end-to-end (see "Live validation on lab0" above): with a real org and log_stream, `runtime-token-exchange` returns 200, the evaluation call rides `X-Agent-Control-Runtime-Token` through the real O11y gateway and returns 200 (`is_safe: true`), and the same token on `Authorization` is rejected 401. Control firing and span ingestion were not exercised in this run (no control bound to the log_stream), and are separate from the header change. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Added 6 luna out of the box, verified and tested.
Added 2 more controls, 1 regex ( for prompt injection ) & one json control.
Summary
Scope
Risk and Rollout
Testing
make check(or explained why not)Checklist