Ship CLAUDE_CODE_OAUTH_TOKEN to claude_sdk* workers (DEV-1517) - #17
Conversation
…-1517) When CLAUDE_CODE_OAUTH_TOKEN is present locally and the framework is claude_sdk*, read_api_keys_from_local_env() omits ANTHROPIC_API_KEY from the secrets file and renames it to BIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY so the user-sim (LiteLLM) can still authenticate while the Claude Agent SDK bills the user's claude.ai subscription via OAuth. Key pieces: - prereqs.check_api_keys: OAuth-aware validation (bad prefix → hard error; anthropic user-sim requires dedicated key; slayer still requires OPENAI) - driver.read_api_keys_from_local_env: OAuth path never ships ANTHROPIC_API_KEY - ray_app._apply_actor_env_local: pops ANTHROPIC_API_KEY from os.environ in local mode so the in-process SDK cannot discover it - ray_app._assert_actor_oauth_invariant: sanity check fired at actor init - usage._maybe_inject_anthropic_key: injects api_key kwarg into LiteLLM calls for anthropic/* models when the dedicated env var is present Full TDD test coverage across prereqs, driver, ray_app, and usage modules. Pre-existing test failures also fixed (test isolation leak in test_strict, livesqlbench audit scope guards, sar_audit skip when data absent). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThis PR adds framework-aware Claude SDK OAuth handling: prereqs and driver now support an OAuth fast-path, actor startup enforces OAuth/API-key invariants and applies local env cleanup, LiteLLM usage can inject Anthropic keys for OAuth runs, and comprehensive tests were added or adjusted. ChangesOAuth Support Implementation
Sequence Diagram(s)sequenceDiagram
participant DriverSubmit as driver.submit
participant ReadKeys as read_api_keys_from_local_env
participant Prereqs as check_api_keys
participant LocalActor as _LocalActor.__init__
participant WorkerActor as WorkerActor.__init__
participant Invariant as _assert_actor_oauth_invariant
participant LocalEnv as _apply_actor_env_local
participant AgentTask as Agent.run_task
participant Usage as _maybe_inject_anthropic_key
DriverSubmit->>ReadKeys: call with framework="claude_sdk"
ReadKeys->>Prereqs: (optional) validate token prefix & required keys
ReadKeys-->>DriverSubmit: return env_vars (CLAUDE_CODE_OAUTH_TOKEN, remapped keys)
DriverSubmit->>LocalEnv: apply actor_env_vars (local-mode)
LocalEnv->>Invariant: (when OAuth) remove ANTHROPIC_API_KEY from os.environ
LocalActor->>Invariant: validate env at bootstrap
WorkerActor->>Invariant: validate env at bootstrap
AgentTask->>Usage: before completion, call _maybe_inject_anthropic_key(model, kwargs)
Usage-->>AgentTask: kwargs possibly augmented with api_key
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/bird_interact_agents/cloud/driver.py`:
- Around line 125-145: The OAuth fast-path currently indexes os.environ[...] (in
the block guarded by prereqs._is_claude_sdk_framework(framework)) before running
the missing-key check, which can raise KeyError and it also skips validating the
CLAUDE_CODE_OAUTH_TOKEN format; change those direct env lookups to
os.environ.get(...) and treat absent or empty values as missing, and add a
validation step for CLAUDE_CODE_OAUTH_TOKEN (check prefix like "sk-ant-oat01-"
or other expected pattern) and treat an invalid token as missing so the same
missing list/messaging and PrereqError (the PrereqError raise in this block) is
used; update uses of user_sim_model, _required_api_keys, and the DEV-1468
OPENAI_API_KEY injection logic to populate result via .get and then compute
missing = [k for k, v in result.items() if not v or (k ==
"CLAUDE_CODE_OAUTH_TOKEN" and not valid_prefix(v))] before raising PrereqError
with remediation.
In `@tests/test_strict.py`:
- Around line 61-62: The test currently only checks membership of
"ANTHROPIC_API_KEY" in os.environ which treats an existing empty value as
present; change the guard to check the value instead (e.g., use
os.environ.get("ANTHROPIC_API_KEY") and treat None or "" as missing) and call
monkeypatch.setenv("ANTHROPIC_API_KEY", "fake-key") when the value is missing or
empty so the test no longer depends on an empty host environment variable.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: d235ee17-d8e3-41d5-9d9a-1a641707c2f9
📒 Files selected for processing (14)
src/bird_interact_agents/agents/claude_sdk/agent.pysrc/bird_interact_agents/agents/claude_sdk_otf/agent.pysrc/bird_interact_agents/agents/claude_sdk_otf_ainteract/agent.pysrc/bird_interact_agents/cloud/driver.pysrc/bird_interact_agents/cloud/prereqs.pysrc/bird_interact_agents/cloud/ray_app.pysrc/bird_interact_agents/usage.pytests/cloud/test_driver.pytests/cloud/test_prereqs.pytests/cloud/test_ray_app.pytests/sar_audit/test_verifier.pytests/test_livesqlbench_audited_gold.pytests/test_strict.pytests/test_usage.py
Group A: read_api_keys_from_local_env fails fast with PrereqError - Bare os.environ[k] accesses (lines 130/134/137) replaced with .get(k,"") so the existing missing-key guard fires on absent keys instead of raising raw KeyError (resubmit bypasses check_api_keys, making this a real crash). - Add sk-ant-oat01- prefix validation before building result so an invalid OAuth token is caught at submission time, not at actor __init__ retries. - Two new tests: bad-prefix and missing-usersim-key raise PrereqError. Group B: test_strict.py empty-string guard - `"ANTHROPIC_API_KEY" not in os.environ` → `not os.environ.get(...)` so an empty-string key doesn't skip the fake-key injection. Group C: ray_app._assert_actor_oauth_invariant diagnostic note - RuntimeError message now directs operator to the manifest secrets file rather than Ray restart logs, since the actor won't be restarted. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/bird_interact_agents/cloud/driver.py (1)
137-151:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winKeep the missing-key remediation on the local env-var name.
This branch reads
ANTHROPIC_API_KEYlocally and only renames it for the worker, but themissing/remediationlogic now surfacesBIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY. On a resubmit failure, the suggestedexport ...command points operators at the wrong variable.Suggested fix
result: dict[str, str] = {"CLAUDE_CODE_OAUTH_TOKEN": token} + missing_local: set[str] = set() # Rename the Anthropic key for user-sim; LiteLLM reads it via # _maybe_inject_anthropic_key in usage.acompletion_tracked. if user_sim_model.startswith("anthropic/"): - result["BIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY"] = os.environ.get("ANTHROPIC_API_KEY", "") + anthropic_key = os.environ.get("ANTHROPIC_API_KEY", "") + result["BIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY"] = anthropic_key + if not anthropic_key: + missing_local.add("ANTHROPIC_API_KEY") # Non-anthropic user-sim keys (OPENAI, CEREBRAS, GEMINI). for k in _required_api_keys(user_sim_model): if k != "ANTHROPIC_API_KEY": - result[k] = os.environ.get(k, "") + result[k] = os.environ.get(k, "") + if not result[k]: + missing_local.add(k) # DEV-1468: slayer embeddings always need OPENAI_API_KEY. if query_mode == "slayer": - result["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY", "") + result["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY", "") + if not result["OPENAI_API_KEY"]: + missing_local.add("OPENAI_API_KEY") # Fail fast on missing keys (resubmit has no prereq check). - missing = [k for k, v in result.items() if not v] + missing = sorted(missing_local) if missing: - cmds = "\n".join(f"export {k}=<your-key>" for k in sorted(missing)) + cmds = "\n".join(f"export {k}=<your-key>" for k in missing)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/bird_interact_agents/cloud/driver.py` around lines 137 - 151, The remediation lists the worker-side variable names (e.g., BIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY) instead of the local env-var names operators should export; update the missing-key detection to use the original/local env var names: when you populate result (including the special-case mapping of ANT HROPIC_API_KEY -> BIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY), keep a parallel list or mapping of the source/local env var names (e.g., "ANTHROPIC_API_KEY", plus all keys returned by _required_api_keys(user_sim_model) and "OPENAI_API_KEY" for query_mode == "slayer"), compute missing against that local-names list, and generate the remediation cmds and PrereqError using those local names rather than the worker-side names so the suggested export commands are correct.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/bird_interact_agents/cloud/driver.py`:
- Around line 137-151: The remediation lists the worker-side variable names
(e.g., BIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY) instead of the local env-var
names operators should export; update the missing-key detection to use the
original/local env var names: when you populate result (including the
special-case mapping of ANT HROPIC_API_KEY ->
BIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY), keep a parallel list or mapping of the
source/local env var names (e.g., "ANTHROPIC_API_KEY", plus all keys returned by
_required_api_keys(user_sim_model) and "OPENAI_API_KEY" for query_mode ==
"slayer"), compute missing against that local-names list, and generate the
remediation cmds and PrereqError using those local names rather than the
worker-side names so the suggested export commands are correct.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 55309152-e888-4b9d-b3fd-b2cf3ea8ad1d
📒 Files selected for processing (4)
src/bird_interact_agents/cloud/driver.pysrc/bird_interact_agents/cloud/ray_app.pytests/cloud/test_driver.pytests/test_strict.py
✅ Files skipped from review due to trivial changes (1)
- tests/test_strict.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/cloud/test_driver.py
- src/bird_interact_agents/cloud/ray_app.py
…-1517) The missing-key PrereqError was listing worker-side variable names (e.g. BIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY) in its remediation `export` commands instead of the local names operators should actually set (ANTHROPIC_API_KEY). Track a separate `missing_local` list of source env var names and use that for the error message. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…EV-1517)
The invariant check fired on ambient CLAUDE_CODE_OAUTH_TOKEN even for
pydantic_ai/agno/etc. local runs where the token is irrelevant — a developer
with the token in their shell would get a spurious RuntimeError. Scope the
guard to cfg["framework"].startswith("claude_sdk") and add a regression test.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Summary
CLAUDE_CODE_OAUTH_TOKENis set locally and the framework isclaude_sdk*, the driver ships it to workers as-is and omitsANTHROPIC_API_KEYfrom the secrets file, so the Claude Agent SDK authenticates via OAuth (billing the user's claude.ai subscription) rather than the console API key.ANTHROPIC_API_KEYis renamed toBIRD_INTERACT_LITELLM_ANTHROPIC_API_KEYso the user-sim (LiteLLM path) can still authenticate;usage._maybe_inject_anthropic_keyinjects it as an explicitapi_keykwarg intoanthropic/*calls.ray_app._apply_actor_env_localpopsANTHROPIC_API_KEYfrom the in-processos.environin local mode so the SDK cannot auto-discover it.ray_app._assert_actor_oauth_invariantfires at actor init (both cloud and local) as a safety net: raises if both vars co-exist, or if the token prefix is wrong.prereqs.check_api_keysvalidates the OAuth path: bad prefix → hard error; anthropic user-sim requires the dedicated key; slayer mode still requiresOPENAI_API_KEY.resubmitgracefully handles old manifests without aframeworkkey (falls back to legacy path).Pre-existing fixes included
test_strict.py:os.environ.setdefaultleakedANTHROPIC_API_KEY="fake-key"into subsequent tests causing a flakylitellm.AuthenticationError— fixed withmonkeypatch.test_ray_app.py:monkeypatch.delenvon absent keys is a no-op (doesn't record cleanup), so keys set by_apply_actor_env_localviaos.environ.updateweren't cleaned up — fixed withtry/finally.test_livesqlbench_audited_gold.py: KB/column-meaning citation tests hardcoded tomuseumentries only; credit/mental entries use different namespaces — added scope guard.tests/sar_audit/test_verifier.py: regression test skips whenaudited_gold/credit/data not yet present on this machine.Test plan
uv run --extra all --extra dev --extra pydantic-ai pytest→ 1726 passed, 0 failedtest_claude_sdk_oauth_*intests/cloud/test_prereqs.py— OAuth prereq validationtest_read_api_keys_oauth_*intests/cloud/test_driver.py— secrets file content on OAuth pathtest_apply_actor_env_local_*andtest_assert_actor_oauth_invariant_*intests/cloud/test_ray_app.pytest_maybe_inject_anthropic_key_*andtest_acompletion_tracked_injects_api_key_*intests/test_usage.pyclaude_sdk_otfframework, single task,--detach, verify worker logs show OAuth auth🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests