Skip to content

Ship CLAUDE_CODE_OAUTH_TOKEN to claude_sdk* workers (DEV-1517) - #17

Merged
ZmeiGorynych merged 4 commits into
mainfrom
egor/dev-1517-cloud-ship-claude_code_oauth_token-to-claude_sdk-workers
Jun 2, 2026
Merged

Ship CLAUDE_CODE_OAUTH_TOKEN to claude_sdk* workers (DEV-1517)#17
ZmeiGorynych merged 4 commits into
mainfrom
egor/dev-1517-cloud-ship-claude_code_oauth_token-to-claude_sdk-workers

Conversation

@ZmeiGorynych

@ZmeiGorynych ZmeiGorynych commented Jun 2, 2026

Copy link
Copy Markdown
Member

Summary

  • When CLAUDE_CODE_OAUTH_TOKEN is set locally and the framework is claude_sdk*, the driver ships it to workers as-is and omits ANTHROPIC_API_KEY from 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_KEY is renamed to BIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY so the user-sim (LiteLLM path) can still authenticate; usage._maybe_inject_anthropic_key injects it as an explicit api_key kwarg into anthropic/* calls.
  • ray_app._apply_actor_env_local pops ANTHROPIC_API_KEY from the in-process os.environ in local mode so the SDK cannot auto-discover it.
  • ray_app._assert_actor_oauth_invariant fires 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_keys validates the OAuth path: bad prefix → hard error; anthropic user-sim requires the dedicated key; slayer mode still requires OPENAI_API_KEY.
  • resubmit gracefully handles old manifests without a framework key (falls back to legacy path).

Pre-existing fixes included

  • test_strict.py: os.environ.setdefault leaked ANTHROPIC_API_KEY="fake-key" into subsequent tests causing a flaky litellm.AuthenticationError — fixed with monkeypatch.
  • test_ray_app.py: monkeypatch.delenv on absent keys is a no-op (doesn't record cleanup), so keys set by _apply_actor_env_local via os.environ.update weren't cleaned up — fixed with try/finally.
  • test_livesqlbench_audited_gold.py: KB/column-meaning citation tests hardcoded to museum entries only; credit/mental entries use different namespaces — added scope guard.
  • tests/sar_audit/test_verifier.py: regression test skips when audited_gold/credit/ data not yet present on this machine.

Test plan

  • Full non-integration suite: uv run --extra all --extra dev --extra pydantic-ai pytest → 1726 passed, 0 failed
  • test_claude_sdk_oauth_* in tests/cloud/test_prereqs.py — OAuth prereq validation
  • test_read_api_keys_oauth_* in tests/cloud/test_driver.py — secrets file content on OAuth path
  • test_apply_actor_env_local_* and test_assert_actor_oauth_invariant_* in tests/cloud/test_ray_app.py
  • test_maybe_inject_anthropic_key_* and test_acompletion_tracked_injects_api_key_* in tests/test_usage.py
  • Cloud smoke: claude_sdk_otf framework, single task, --detach, verify worker logs show OAuth auth

🤖 Generated with Claude Code

Summary by CodeRabbit

  • New Features

    • Added Claude SDK OAuth support with token validation and framework-aware key handling.
    • Inject Anthropic API key into model request calls during OAuth runs to preserve token/cost tracking.
  • Bug Fixes

    • Enforced OAuth/API-key environment invariants and strip conflicting ambient keys in local runs.
    • Deterministic fail-fast key validation, clearer remediation guidance, and consistent key selection across submit/resubmit.
  • Tests

    • Extensive new tests covering OAuth flows, env handling, prerequisite checks, and regressions.

…-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>
@linear

linear Bot commented Jun 2, 2026

Copy link
Copy Markdown

DEV-1517

@coderabbitai

coderabbitai Bot commented Jun 2, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: a9e22f20-8ccf-4bca-ae9a-c6009952b131

📥 Commits

Reviewing files that changed from the base of the PR and between 72db673 and 476b4e6.

📒 Files selected for processing (2)
  • src/bird_interact_agents/cloud/ray_app.py
  • tests/cloud/test_ray_app.py
🚧 Files skipped from review as they are similar to previous changes (2)
  • src/bird_interact_agents/cloud/ray_app.py
  • tests/cloud/test_ray_app.py

📝 Walkthrough

Walkthrough

This 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.

Changes

OAuth Support Implementation

Layer / File(s) Summary
OAuth Invariant Enforcement
src/bird_interact_agents/cloud/ray_app.py, src/bird_interact_agents/agents/claude_sdk*
New _assert_actor_oauth_invariant() enforces that CLAUDE_CODE_OAUTH_TOKEN and ANTHROPIC_API_KEY are not both set and that the OAuth token has the sk-ant-oat01- prefix. New _apply_actor_env_local() removes legacy ANTHROPIC_API_KEY from the ambient environment when using OAuth in local mode. Invariant is called at actor bootstrap in both _LocalActor and Ray WorkerActor. Agent modules receive inline comments documenting this validation.
Framework-aware Driver Key Selection
src/bird_interact_agents/cloud/driver.py
read_api_keys_from_local_env() gains optional framework parameter and implements an OAuth fast-path: when CLAUDE_CODE_OAUTH_TOKEN is present for claude_sdk* frameworks, it validates the token prefix, returns the OAuth token plus remapped Anthropic user-sim key (and provider-required vars for non-anthropic usersims), and fails fast with PrereqError and a deterministic remediation script on missing keys. submit() and resubmit() forward framework (resubmit reads it defensively from manifest).
Framework-aware Prereqs Key Validation
src/bird_interact_agents/cloud/prereqs.py
check_api_keys() accepts framework and branches for Claude SDK OAuth: validates token prefix and computes required keys based on user_sim_model (including OPENAI_API_KEY for slayer). Legacy path still requires keys from both agent_model and user_sim_model; missing-key reporting is sorted. The orchestrator check(args) forwards framework.
LiteLLM Anthropic Credentials Injection
src/bird_interact_agents/usage.py
New _maybe_inject_anthropic_key() reads a dedicated env var and injects api_key into kwargs for anthropic/* models when not already provided, avoiding mutation of os.environ. acompletion_tracked() calls this helper before invoking the retry-wrapped LiteLLM completion.
Ray App OAuth Environment Tests
tests/cloud/test_ray_app.py
New tests cover _assert_actor_oauth_invariant across no-OAuth, valid-token-only, co-presence error, and bad-prefix error scenarios; _load_secrets_file round-trips OAuth names; _with_actor_env propagates OAuth vars into runtime env; _apply_actor_env_local removes legacy keys on the OAuth path and preserves them on the legacy path.
Driver & Prereqs Tests
tests/cloud/test_driver.py, tests/cloud/test_prereqs.py
New tests validate read_api_keys_from_local_env OAuth path returns CLAUDE_CODE_OAUTH_TOKEN plus remapped Anthropic key, conditionally includes OPENAI_API_KEY for slayer/OpenAI usersim, omits legacy keys, and raises PrereqError on invalid prefix or missing keys. Prereqs tests cover claude_sdk OAuth vs legacy paths, pydantic_ai ignoring OAuth, and orchestrator forwarding of framework.
Usage Module Tests
tests/test_usage.py
Adds tests for _maybe_inject_anthropic_key behavior across model prefixes, env presence, caller-provided keys, and os.environ preservation; async tests assert acompletion_tracked passes api_key to _acompletion when appropriate.
Unrelated Test Contract Adjustments
tests/test_livesqlbench_audited_gold.py, tests/sar_audit/test_verifier.py, tests/test_strict.py
Audit file coverage now asserts at least (not exactly) museum_1..10 instances; selected_database derived from instance_id prefix; museum_9 pinned test replaced with existence + status check; citation tests limited to museum rows. Credit audit test is skipped if the sidecar is missing. Pydantic AI test switched to monkeypatch for env setup.

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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

Possibly related PRs

Poem

🐰 Hopping down the OAuth trail,

Tokens set and keys won't fail,
Actors check at bootstrap time,
LiteLLM gets the right key prime,
A little rabbit cheers: secure and hale!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 75.34% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title directly describes the main objective: shipping the CLAUDE_CODE_OAUTH_TOKEN environment variable to claude_sdk workers, and includes the tracking ID (DEV-1517). It clearly summarizes the primary change across the entire changeset.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch egor/dev-1517-cloud-ship-claude_code_oauth_token-to-claude_sdk-workers

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between af58457 and 864a2bb.

📒 Files selected for processing (14)
  • src/bird_interact_agents/agents/claude_sdk/agent.py
  • src/bird_interact_agents/agents/claude_sdk_otf/agent.py
  • src/bird_interact_agents/agents/claude_sdk_otf_ainteract/agent.py
  • src/bird_interact_agents/cloud/driver.py
  • src/bird_interact_agents/cloud/prereqs.py
  • src/bird_interact_agents/cloud/ray_app.py
  • src/bird_interact_agents/usage.py
  • tests/cloud/test_driver.py
  • tests/cloud/test_prereqs.py
  • tests/cloud/test_ray_app.py
  • tests/sar_audit/test_verifier.py
  • tests/test_livesqlbench_audited_gold.py
  • tests/test_strict.py
  • tests/test_usage.py

Comment thread src/bird_interact_agents/cloud/driver.py
Comment thread tests/test_strict.py Outdated
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>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Keep the missing-key remediation on the local env-var name.

This branch reads ANTHROPIC_API_KEY locally and only renames it for the worker, but the missing/remediation logic now surfaces BIRD_INTERACT_LITELLM_ANTHROPIC_API_KEY. On a resubmit failure, the suggested export ... 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

📥 Commits

Reviewing files that changed from the base of the PR and between 864a2bb and 499f2e3.

📒 Files selected for processing (4)
  • src/bird_interact_agents/cloud/driver.py
  • src/bird_interact_agents/cloud/ray_app.py
  • tests/cloud/test_driver.py
  • tests/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

ZmeiGorynych and others added 2 commits June 2, 2026 13:23
…-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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant