diff --git a/.egg-state/agent-outputs/2769-architect-output.json b/.egg-state/agent-outputs/2769-architect-output.json new file mode 100644 index 0000000000..f94dc8e479 --- /dev/null +++ b/.egg-state/agent-outputs/2769-architect-output.json @@ -0,0 +1,488 @@ +{ + "issue": 2769, + "phase": "plan", + "role": "architect", + "schema_version": "2", + "title": "Architecture: non-Claude models per agent via LiteLLM proxy", + "summary": "Design for an additive gateway-side upstream router that lets each agent run on either Anthropic (existing path, untouched) or a LiteLLM proxy that translates to OpenAI-compatible backends (hosted Qwen first per cq-6). Per-agent selection is driven by per-agent session metadata recorded at spawn time (cq-2). Configuration is per-role on PipelineConfig with a repositories.yaml default (cq-3). LiteLLM lives as a separate Deployment+Service in egg-system (cq-1). Failure policy is fail-closed (cq-8). Tool-strip in private mode stays uniform regardless of upstream (cq-9). The opus[1m] Claude defaults are left alone (cq-11). For LiteLLM-bound requests, the gateway rewrites the request body's `model` field from the Claude alias Claude Code sent (e.g. 'opus') to the LiteLLM-side natural model name (e.g. 'qwen-2.5-coder') — chosen explicitly as 'Semantics B' over the alternative ('Semantics A' — pass body unchanged, require operator to key LiteLLM's `model_list` on Claude aliases). v2 changes: pinned Semantics B (the wire-level body rewrite for both /v1/messages and /v1/messages/count_tokens) with file:line evidence; added R8 (tokenizer mismatch — the reason Semantics A was rejected); added AC-10 (repositories.yaml absence handling); added k8s/base/gateway-deployment.yaml to slice-1 production_code (LITELLM_BASE_URL env var); flagged key-validation requirement on PipelineConfig.agent_models; surfaced observable cost-tracker mis-pricing as documented-and-out-of-scope; concrete NetworkPolicy egress guidance (operator-supplied per backend); cross-checked config/repo_config.py is not under #2261 decomposition; expanded slice-1 SSE test to include a client-disconnect mid-stream case against LiteLLM.", + "v2_change_log": [ + "v1 left ambiguous how `agent_model_litellm` reaches the wire. v2 commits to Semantics B (gateway rewrites the request body's `model` field on LiteLLM-bound requests). See `recommended_approach.litellm_dispatch_semantics` below. This single decision drives the slice-1 test surface, the LiteLLM Deployment sample model_list, and the operator-facing documentation, so it is now pinned at the architect tier (reviewer_plan NACK item #1).", + "Added slice-1 component 'Body model rewriter for LiteLLM-bound requests' with the exact insertion point (gateway/gateway.py between credential injection at line 9767 and client.send at line 9827, plus a sibling rewrite in proxy_count_tokens between the credential injection at line 10028 and client.post at line 10035).", + "Bumped AC-2 (slice-1) to assert body mutation on the LiteLLM path (body['model'] changes from 'opus' to 'qwen-2.5-coder') for BOTH /v1/messages and /v1/messages/count_tokens.", + "Added R8: tokenizer mismatch under Semantics A — Claude vs Qwen token boundaries differ; compaction math driven by count_tokens would be wrong under Semantics A. Rejection of Semantics A.", + "Added AC-10: get_default_model(repo) returns None silently when config/repositories.yaml is absent or has no default_model entry — matches the existing get_repo_setting behavior. config/repositories.yaml is operator-supplied at runtime (only repositories.yaml.example is in-repo).", + "Added k8s/base/gateway-deployment.yaml to slice-1 production_code: declare `LITELLM_BASE_URL` env var on the gateway container (default in-cluster Service DNS). Without it, the seam's swap-out flexibility (refine feedback Q3) is unreachable.", + "Flagged PipelineConfig.agent_models key-validation as a non-negotiable design requirement (Pydantic v2 field_validator) — planner picks the shape (`dict[AgentRole,str]` with validator vs explicit per-role fields like the overseer-model precedent), but typo-tolerance is forbidden.", + "Added the cost-tracker mis-pricing observable behavior to out_of_scope_explicitly so downstream operators are not surprised.", + "Made NetworkPolicy egress concrete: operator-supplied per backend choice via a kustomize overlay (not hard-coded in k8s/base/network-policies.yaml). Example hosts named: api.together.xyz, api.fireworks.ai, openrouter.ai.", + "Cross-checked config/repo_config.py: 836 lines, not in any #2261 decomposition row (verified against scripts/file-size-allowlist.yaml and CLAUDE.md tables). Helper placement near get_repo_setting (line 248) is safe.", + "Expanded slice-1 SSE acceptance test (AC-2 part (e)) to include a 'client disconnect mid-stream against the LiteLLM upstream' case, asserting _SSEAccumulator partial-flush parity with the Anthropic path." + ], + "recommended_approach": { + "option_id": "A", + "name": "Gateway-side UpstreamRegistry keyed by per-agent session metadata; LiteLLM as a sibling Deployment in egg-system; gateway rewrites body.model on LiteLLM-bound requests (Semantics B)", + "summary": "Decompose get_anthropic_client() into a tiny UpstreamRegistry with two named upstreams ('anthropic', 'litellm'). proxy_anthropic_messages() and proxy_count_tokens() resolve the upstream per request via the existing IP-keyed session lookup. The orchestrator records the per-agent upstream + LiteLLM-side natural model name on register_session (mirroring how it already records phase, agent_role, mode). _inject_anthropic_credentials becomes upstream-aware: LiteLLM-bound requests get the LiteLLM master key (from secrets.env, mirroring the existing ANTHROPIC_API_KEY pattern) instead of the Anthropic OAuth/API credential. **For LiteLLM-bound requests, the gateway also rewrites the request body's `model` field** from the Claude alias Claude Code put there (e.g. 'opus' — what Claude Code's compaction math expects to see) to the LiteLLM-side natural model name (e.g. 'qwen-2.5-coder' — what the operator configured in LiteLLM's model_list). This mutation is the ONE additional gateway-side responsibility on the LiteLLM path. The SSE accumulator, stream-resilience retries, and tool-strip are upstream-agnostic and untouched. LiteLLM runs as a separate Deployment+Service (egg-system namespace), reached over cluster DNS at http://litellm.egg-system.svc.cluster.local:; agent pods cannot reach it directly (squid allowlist excludes it; only the gateway calls it).", + "litellm_dispatch_semantics": { + "choice": "Semantics B — gateway rewrites body.model before forwarding to LiteLLM", + "what_changes_on_the_wire": "Claude Code emits body={..., 'model': 'opus', ...}. On the Anthropic upstream, the body is forwarded byte-unchanged (today's behavior). On the LiteLLM upstream, between credential injection and client.send, the gateway parses the body JSON, replaces body['model'] with session.agent_model_litellm (e.g. 'qwen-2.5-coder'), and re-serializes. The Authorization header is the LiteLLM Bearer token (not the Anthropic credential).", + "why_this_choice": [ + "Decouples 'what Claude Code sees in its --model flag and its compaction math' (recognized Claude alias) from 'what reaches LiteLLM's model_list lookup' (the natural LiteLLM model name). This is exactly the seam the issue's compaction-mitigation note calls for.", + "Avoids the tokenizer-mismatch failure mode of Semantics A: if LiteLLM's count_tokens (the route Claude Code hits to drive compaction) ran against the Claude alias, LiteLLM would compute tokens with the WRONG tokenizer (Anthropic's tokenizer vs Qwen's BPE — different boundaries → different counts → compaction triggers at the wrong moment → possible over-length requests that hard-fail). With Semantics B, the body rewrite extends to /v1/messages/count_tokens, so LiteLLM tokenizes against the actual backend model.", + "Lets operators configure LiteLLM's model_list with NATURAL model names (the keys every LiteLLM example uses: `model_name: 'qwen-2.5-coder'`). Semantics A would force the constraint 'LiteLLM model_list keys must equal Claude aliases the agents will send' — a non-obvious operator footgun that an operator-supplied config drift can break silently.", + "Session.agent_model_litellm becomes a functional, required-at-request-time field, not advisory metadata. This is simpler to reason about and tests assert it directly drives the wire." + ], + "what_we_explicitly_rejected_semantics_a": { + "name": "Semantics A — body forwarded byte-unchanged; LiteLLM dispatches on Claude-alias keys in its model_list", + "rejected_because": "Tokenizer-mismatch on count_tokens (R8) is unfixable under A without partial-rewrite gymnastics; operator's LiteLLM model_list config becomes brittle (a typo in either the agent's --model alias or the LiteLLM key silently routes to the wrong model); Session.agent_model_litellm becomes purely advisory at request time, which is harder to reason about than 'this field decides what reaches the wire'." + } + }, + "why_recommended": [ + "Satisfies the issue's hard constraint that the Claude path must not regress: anthropic upstream is the existing singleton client with the existing base_url, existing credential injection, existing SSE accumulator, existing retry policy. The Anthropic path's body is unchanged. Routing is additive; rewriting only happens on the LiteLLM branch.", + "Inert-by-default: no agent is configured for LiteLLM in shipping defaults, so no LiteLLM request ever fires until an operator opts a role in via PipelineConfig.agent_models or repositories.yaml.", + "Decouples 'model name Claude Code sees' (recognized alias, keeps compaction math sane per the issue's primary risk note) from 'model name the LiteLLM upstream dispatches on' (the natural LiteLLM model name). Option C (route on body) cannot do this; option B (LiteLLM-fronts-everything) adds the LiteLLM blast-radius to the Claude path and is forbidden by cq-1/refine.", + "Per-agent granularity falls out naturally: the gateway already keys per-request policy by IP-based session lookup (session_mode is the existing precedent — gateway/gateway.py:9775). Adding upstream + litellm_model fields to the same Session and the existing /api/v1/sessions/create payload is a single additive change.", + "Matches the existing operational grain: gateway is already the per-request policy point; a sibling Deployment in egg-system is the standard pattern (gateway, orchestrator); secrets.env is the existing credential-mount pattern." + ], + "rejected_options_summary": { + "B_litellm_fronts_everything": "Forces all Claude traffic through LiteLLM. Adds a hop + supply-chain blast radius + restart-cycle dependency to every existing Claude request. Violates cq-1 (refine) and the issue's no-regression-on-Claude constraint.", + "C_route_on_request_body_model": "Smallest gateway change, but directly conflicts with the compaction-math mitigation: Claude Code must see a recognized Claude alias even when the backend is Qwen, so model-in-body cannot be the routing key. Also couples upstream selection to model-name conventions (silent break if Anthropic ships a non-claude-* model or LiteLLM exposes a claude-aliased Qwen).", + "D_egg_agent_sdk_path_for_non_claude": "Sidesteps Claude Code's compaction math by bypassing Claude Code entirely. Forbidden by cq-5 (refine) — operator chose 'keep Claude Code harness for non-Claude models too'." + } + }, + "component_breakdown": { + "slices": [ + { + "id": "slice-1", + "name": "Gateway upstream router + LiteLLM topology (no-op by default)", + "depends_on": [], + "summary": "Ship the gateway-side abstraction and the LiteLLM Deployment+Service. Claude path unchanged. With no agent configured for LiteLLM, the LiteLLM upstream is reachable from the gateway but no agent traffic ever reaches it. Buildable, lintable, unit-testable without a live LiteLLM endpoint.", + "components": [ + { + "name": "UpstreamRegistry (new gateway module)", + "where": "gateway/_upstream_registry.py (new file, picked up by gateway.py's barrel via the slice-14 decomposition pattern; until slice-14 lands, may live inline in gateway/gateway.py near get_anthropic_client at line 9320 — planner to pick)", + "responsibility": "Holds a dict of upstream name -> (base_url, httpx.Client singleton, credential-resolution policy). Two upstreams at landing: 'anthropic' (preserves today's https://api.anthropic.com client, preserving timeouts/limits exactly) and 'litellm' (base_url from env LITELLM_BASE_URL, default http://litellm.egg-system.svc.cluster.local:4000). Exposes get_client(upstream_name) -> httpx.Client. Designed so a future swap-out of LiteLLM is one new upstream entry (refine feedback Q3).", + "evidence": "gateway/gateway.py:9316-9329 (current singleton); gateway/gateway.py:9789 (sole call site in proxy_anthropic_messages); gateway/gateway.py:10032 (sole call site in proxy_count_tokens). Both call sites are replaced with get_client(resolved_upstream)." + }, + { + "name": "Session-driven upstream resolver (new helper, gateway-side)", + "where": "gateway/gateway.py near proxy_anthropic_messages (current line 9752) — small inline helper or a sibling _resolve_upstream(session) -> str", + "responsibility": "Given the IP-resolved Session (already fetched at gateway/gateway.py:9775 for session_mode), return the upstream name. Today's session attributes drive policy (session.mode → tool-strip); the new attribute session.agent_upstream / session.agent_model_litellm drives upstream selection. Default = 'anthropic' (no-op for Claude agents).", + "evidence": "gateway/gateway.py:9775 calls session_manager.get_session_by_ip(); the resolver is the same lookup with one extra attribute read." + }, + { + "name": "Body model rewriter for LiteLLM-bound requests (NEW, v2 — directly addresses reviewer_plan NACK #1)", + "where": "gateway/gateway.py — TWO insertion points: (a) in proxy_anthropic_messages, between _inject_anthropic_credentials at line 9767 and _filter_blocked_tools at line 9778 (i.e. immediately after credential injection but before the body is consumed by the tool-strip and the streaming branch at line 9827). (b) in proxy_count_tokens, between _inject_anthropic_credentials at line 10028 and client.post at line 10035.", + "responsibility": "When the resolved upstream is 'litellm' and session.agent_model_litellm is set: parse request_body JSON, replace body['model'] with session.agent_model_litellm, re-serialize (json.dumps with separators=(',', ':') to keep the wire-format compact). On JSON-parse failure (defensive — should never happen for /v1/messages or /v1/messages/count_tokens): leave the body unchanged and emit a logger.warning so the operator can investigate. On the Anthropic upstream: body unchanged (today's behavior preserved byte-for-byte). The mutation only touches the single 'model' key; all other body fields (messages, tools, system, max_tokens, etc.) are preserved.", + "evidence": "Current handler reads body at gateway/gateway.py:9771 (request.get_data()) and applies _filter_blocked_tools at :9778 — the rewriter sits between credential injection and tool-strip so the tool-strip sees the rewritten body (preserves the invariant that everything downstream sees one canonical body). Current /v1/messages/count_tokens at gateway/gateway.py:10019-10044 reads body at :10038 (request.get_data() inside the client.post call) — the rewriter performs request.get_data() once at the top and threads the rewritten bytes through. NEW symbol; no existing implementation to cite." + }, + { + "name": "Upstream-aware credential injection", + "where": "gateway/gateway.py:9355 (_inject_anthropic_credentials) — extend with upstream parameter", + "responsibility": "Today the function calls credentials_manager.get_credential() unconditionally. New behavior: when upstream=='litellm', read LITELLM_MASTER_KEY from the AnthropicCredentialsManager / a sibling LiteLLMCredentialsManager (planner to decide whether to widen AnthropicCredentialsManager or add a sibling — both are minimally invasive; the cleaner cut is a sibling LiteLLMCredentialsManager keyed on the secrets.env LITELLM_MASTER_KEY var) and inject as 'Authorization: Bearer '. Claude path unchanged.", + "evidence": "gateway/anthropic_credentials.py:36-49 (AnthropicCredential.header_name/header_value already accommodates both x-api-key and Authorization-Bearer shapes); gateway/anthropic_credentials.py:94-226 (existing manager pattern is the template for a sibling). Sibling manager must replicate the mtime-based cache-invalidation pattern at gateway/anthropic_credentials.py:126 so LITELLM_MASTER_KEY rotation works without a gateway restart (#R6)." + }, + { + "name": "Session-storage extensions for per-agent upstream/model", + "where": "gateway/session_manager.py:288-328 (Session dataclass) + register_session signature (gateway/session_manager.py:548) + /api/v1/sessions/create payload acceptance (gateway/gateway.py:8560-8575)", + "responsibility": "Add two optional Session fields: agent_upstream: Literal['anthropic','litellm']|None and agent_model_litellm: str|None (e.g. 'qwen-2.5-coder' — the natural LiteLLM model name the operator put in model_list; under Semantics B this is what the gateway writes into body['model'] on the LiteLLM path). Persisted in to_dict_for_persistence/from_persistence (gateway/session_manager.py:338-405). Backward-compat: missing fields default to None → 'anthropic'.", + "evidence": "Session dataclass at gateway/session_manager.py:288; the agent_role field landed via the same pattern (gateway/session_manager.py:314)." + }, + { + "name": "Orchestrator → gateway session-create payload extension", + "where": "orchestrator/gateway_client.py:602 (register_session) — add agent_upstream/agent_model_litellm params; orchestrator/kubernetes_spawner.py:735 (call site)", + "responsibility": "Plumb new optional fields through gateway_client and into the spawn-time session_info = self.gateway.register_session(...) call at kubernetes_spawner.py:735. The spawner derives the values from a new helper on PipelineConfig (see slice-2).", + "evidence": "Existing pattern: claude_code_version threads from kubernetes_spawner.py:747 through gateway_client.register_session at gateway_client.py:616 through the gateway session_create handler. The new fields follow the same shape." + }, + { + "name": "LiteLLM Deployment + Service + secret + NetworkPolicy", + "where": "k8s/base/litellm-deployment.yaml (new), k8s/base/litellm-service.yaml (new), k8s/base/kustomization.yaml (add resources), k8s/base/network-policies.yaml (add ingress allowing only gateway-labeled pods).", + "responsibility": "Deploy LiteLLM with a config that exposes ONE provider for the first cut (hosted Qwen per cq-6): a model_list entry whose `model_name` key is the natural model name (e.g. `qwen-2.5-coder`) that the gateway writes into body['model'] under Semantics B — NOT a Claude alias. Master key lives in a k8s Secret. Image is pinned to a digest (refine feedback Q3 — supply-chain mitigation). Replicas=1 for first cut; not on the critical path of Claude traffic. Service exposes port 4000 (LiteLLM default). The NetworkPolicy ingress permits only the gateway pod (labelled selector); egress to the actual hosted-Qwen provider is operator-supplied via a kustomize overlay (see component below) and NOT in k8s/base.", + "evidence": "k8s/base/gateway-deployment.yaml (existing pattern: Deployment + secret mounts + EGG_SECRETS_PATH=/secrets/secrets.env at lines 71-118); k8s/base/gateway-service.yaml (existing Service pattern); k8s/base/network-policies.yaml (existing namespace-scoped policies)." + }, + { + "name": "Gateway pod LITELLM_BASE_URL env var declaration (NEW, v2 — addresses reviewer_plan non-blocker on swap-out flexibility)", + "where": "k8s/base/gateway-deployment.yaml under the gateway container's env: section (e.g. after the existing GATEWAY_PORT/PROXY_PORT/HEALTH_PORT block at lines 77-82)", + "responsibility": "Declare LITELLM_BASE_URL as an explicit container env var (optional; default value http://litellm.egg-system.svc.cluster.local:4000). The UpstreamRegistry reads this env var so an operator can repoint at a different backend (kustomize overlay for production with a managed LiteLLM-equivalent, or for the swap-out seam refine feedback Q3 calls for) without an image rebuild.", + "evidence": "k8s/base/gateway-deployment.yaml:77-82 (existing env var declarations: GATEWAY_PORT, PROXY_PORT, HEALTH_PORT) — same pattern." + }, + { + "name": "NetworkPolicy egress overlay (NEW, v2 — addresses reviewer_plan non-blocker on egress vagueness)", + "where": "k8s/overlays//litellm-egress-policy.yaml (operator-supplied via kustomize overlay) — NOT in k8s/base.", + "responsibility": "For the cq-6 hosted-Qwen-first first cut, the operator must add a NetworkPolicy egress rule allowing the LiteLLM pod to reach the chosen backend (typical examples: api.together.xyz, api.fireworks.ai, openrouter.ai, api.deepinfra.com). The exact host depends on cq-6 follow-through and is operator-specific — k8s/base ships zero egress allowlist (defense-in-depth: LiteLLM pod cannot egress at all without an overlay). The slice-1 plan documents this in gateway/README.md and adds an example overlay scaffold under k8s/overlays/example/ (planner picks the exact filename) — example only, intentionally not active in k8s/base." + }, + { + "name": "secrets.env LITELLM_MASTER_KEY entry", + "where": "k8s/base/gateway-deployment.yaml — gateway already mounts /secrets/secrets.env; the new key is read by the new credential manager. Also: documentation in the gateway README on what secret to add.", + "responsibility": "Mirror today's ANTHROPIC_* injection pattern. Holding the key in the gateway pod, not the LiteLLM pod, is what cq-7 picked: 'Gateway holds a LiteLLM master key and injects it on every LiteLLM-bound request'." + }, + { + "name": "Slice-1 tests (gateway-side, in-cluster trivially testable)", + "where": "gateway/tests/test_upstream_router.py (new) + extensions to gateway/tests/test_session_manager.py", + "responsibility": "Unit tests with httpx.MockTransport: (a) Claude path: agent_upstream=None → request lands at api.anthropic.com client; verify base_url AND that the request body is forwarded byte-unchanged (i.e. body['model']=='opus' on the wire); verify Anthropic credential injection ran. (b) LiteLLM path: agent_upstream='litellm', agent_model_litellm='qwen-2.5-coder' → request lands at the litellm client; verify base_url; verify Authorization: Bearer ; **assert the body forwarded to LiteLLM has body['model']=='qwen-2.5-coder' (rewritten from 'opus')** — Semantics B's load-bearing assertion. (c) Same assertion for /v1/messages/count_tokens. (d) Upstream resolution failure (unknown upstream name, missing LITELLM_MASTER_KEY) → 502, no leak of Anthropic credential. (e) Tool-strip in private mode applies regardless of upstream (cq-9) — assert _filter_blocked_tools sees the rewritten body, so the tool-strip and the rewriter are wired in the documented order. (f) SSE accumulator and stream-resilience retry path runs against both upstreams with synthetic SSE bytes — proves the SSE plumbing is upstream-agnostic. (g) **Client disconnect mid-stream against the LiteLLM upstream** (v2 addition): close the downstream connection after partial SSE delivery; assert _SSEAccumulator flushes the partial response identically to the Anthropic path — exercises the gateway/gateway.py:9909 mid-stream synthetic-error-frame branch. (h) Session backward-compat: load a Session dict written by pre-slice-1 code (no agent_upstream / agent_model_litellm keys) and assert from_persistence yields None for both → resolver defaults to 'anthropic'. Tests are pure unit tests (httpx.MockTransport); no live LiteLLM needed." + } + ] + }, + { + "id": "slice-2", + "name": "Per-agent model config + consensus_wrapper plumbing", + "depends_on": ["slice-1"], + "summary": "Add the PipelineConfig.agent_models per-role field and the repositories.yaml default_model knob (cq-3). Thread the resolved per-role model into both the consensus wrapper's --model flag (so Claude Code sees the recognized alias) AND the gateway session registration (so the gateway routes to LiteLLM for non-Claude roles AND writes the natural model name into body['model'] on the wire). Both source-of-truth pointers derive from the SAME PipelineConfig lookup, eliminating drift (the cons of Option A in refine).", + "components": [ + { + "name": "PipelineConfig.agent_models", + "where": "orchestrator/models.py near overseer_decision_maker_model (current line 546) and overseer_advisor_model (current line 620)", + "responsibility": "New Pydantic field. **Key-validation is required** (planner picks the shape; the architect does not foreclose). Two acceptable shapes: (a) `agent_models: dict[AgentRole, str] = Field(default_factory=dict, ...)` with a Pydantic v2 `field_validator` that coerces string keys to AgentRole and rejects unknown role names (mirrors the AgentRole enum pattern already used elsewhere in orchestrator/models.py); (b) per-role explicit fields (e.g. `coder_model: str|None`, `refiner_model: str|None`, ...) following the overseer-model precedent at lines 546/620. Shape (a) is more compact for the operator (one field, one dict); shape (b) is more discoverable in IDEs and matches the existing overseer-model precedent. Either way, a free-string dict[str,str] with no validator is forbidden — typos like 'reviewer-refine' vs 'reviewer_refine' silently never resolve and the agent runs on Claude despite the operator's intent.", + "evidence": "Existing field pattern: orchestrator/models.py:546 (overseer_decision_maker_model: str), :620 (overseer_advisor_model: str). AgentRole enum already exists elsewhere in the codebase (consumed by concurrent_executor for role validation)." + }, + { + "name": "repositories.yaml default_model", + "where": "config/repositories.yaml (operator-supplied; only config/repositories.yaml.example exists in-repo); read via config/repo_config.py — add a get_default_model(repo: str) -> str | None helper", + "responsibility": "Operator can set a repo-level default_model (e.g. 'opus' or 'qwen-2.5-coder') that applies to every role not explicitly overridden. Resolution precedence (per cq-3): PipelineConfig.agent_models[role] > repositories.yaml default_model > built-in 'opus'. **Behavior when config/repositories.yaml is absent or has no default_model entry**: get_default_model(repo) returns None silently (mirrors the existing get_repo_setting at config/repo_config.py:248-267 — which returns the provided default of None on missing key). The resolver then falls back to built-in 'opus'. No exception, no error log — this is the dev / no-config path most contributors will hit.", + "evidence": "config/repo_config.py:248-267 (get_repo_setting): existing helper for per-repo settings — returns default on missing key without raising. The new get_default_model(repo) is a one-line wrapper on top: `return get_repo_setting(repo, 'default_model', None)`." + }, + { + "name": "Model resolver helper", + "where": "orchestrator/concurrent_executor.py (best fit, near get_agent_env at line 288) OR a new orchestrator/model_resolver.py (planner to pick — concurrent_executor.py is already importing PipelineConfig and AgentRole)", + "responsibility": "Single function resolve_model_for_role(pipeline_config, repo, role) -> tuple[str, str|None] returning (claude_code_model_alias, litellm_model_or_None). For Claude roles, returns ('opus', None) — Claude Code sees 'opus' and the gateway routes to anthropic. For non-Claude roles, returns ('opus', 'qwen-2.5-coder') — Claude Code STILL sees 'opus' (so its compaction math stays sane, per the issue's primary risk note), but the second tuple element drives the gateway's session.agent_model_litellm and session.agent_upstream='litellm'. Under Semantics B (pinned in v2), this second tuple element is what the gateway rewrites into body['model'] on the wire. This single function is the only source of truth; both the consensus_wrapper and the gateway session-create call read from it." + }, + { + "name": "Plumb resolved model into consensus_wrapper", + "where": "orchestrator/consensus_wrapper.py:620 (build_consensus_wrapped_command, model='opus' default) + call sites at orchestrator/concurrent_executor.py:454 (no model arg today) and orchestrator/routes/pipelines.py:2704 (same)", + "responsibility": "Both call sites already have access to role and PipelineConfig. Compute (claude_code_model_alias, _) = resolve_model_for_role(...) and pass claude_code_model_alias as the model arg. Functionally: Claude roles get 'opus' (today's behavior); non-Claude roles get 'opus' (the recognized alias the issue's mitigation requires) — they DO NOT pass the litellm model string here, so Claude Code's compaction math is unaffected. The litellm model string flows down the gateway-session path (next component)." + }, + { + "name": "Plumb resolved litellm_model + upstream into spawn-time session registration", + "where": "orchestrator/kubernetes_spawner.py:735 (the register_session call) — accept new optional agent_upstream/agent_model_litellm args, set from resolve_model_for_role's second return value. Default behavior preserved: for Claude roles the second return value is None, the session has agent_upstream=None, and the gateway defaults to 'anthropic'.", + "responsibility": "This is where the per-agent upstream metadata is recorded. The spawn site has all the inputs: pipeline_config (passed in via spawn_agent_job's call chain), repo (first repo in repos), agent_role.value." + }, + { + "name": "Slice-2 tests", + "where": "orchestrator/tests/test_model_resolver.py (new) + extensions to orchestrator/tests/test_concurrent_executor.py and orchestrator/tests/test_kubernetes_spawner.py", + "responsibility": "Unit tests for the resolver: (a) no override → ('opus', None); (b) PipelineConfig override → ('opus', ''); (c) repositories.yaml default → ('opus', ''); (d) precedence: pipeline config overrides repo default; (e) repo default 'opus' (Claude) → ('opus', None) — no LiteLLM routing; (f) typo'd role key in PipelineConfig.agent_models triggers a Pydantic validation error (proves the key-validator the planner picks is wired); (g) config/repositories.yaml absent → get_default_model returns None and resolver falls back to 'opus' (AC-10). Integration tests assert the spawn-time register_session call passes the right agent_upstream/agent_model_litellm values for representative configs. No live LiteLLM needed." + } + ] + } + ] + }, + "key_files_changed": { + "production_code_slice_1": [ + "gateway/gateway.py (router + credential-injection extension + body-rewriter + proxy call sites at lines 9320, 9355, 9752, 9767, 9778, 9789, 10019, 10028, 10032)", + "gateway/session_manager.py (Session dataclass + register_session signature, lines 288-405 + 548-617)", + "gateway/anthropic_credentials.py OR new gateway/litellm_credentials.py (sibling manager for LITELLM_MASTER_KEY, modeled on existing manager at lines 94-226 with mtime-based cache invalidation per #R6)", + "k8s/base/litellm-deployment.yaml (new)", + "k8s/base/litellm-service.yaml (new)", + "k8s/base/kustomization.yaml (register new resources)", + "k8s/base/network-policies.yaml (allow gateway→litellm ingress only; no default egress for litellm — operator overlay supplies egress)", + "k8s/base/gateway-deployment.yaml (NEW v2 addition: declare LITELLM_BASE_URL env var on the gateway container — see slice-1 component 'Gateway pod LITELLM_BASE_URL env var declaration')", + "orchestrator/gateway_client.py:602 (register_session signature)", + "orchestrator/kubernetes_spawner.py:735 (spawn-time session_info call)" + ], + "tests_slice_1": [ + "gateway/tests/test_upstream_router.py (new — httpx.MockTransport-driven, no live endpoint needed)", + "gateway/tests/test_session_manager.py (extend for new Session fields + backward-compat from_persistence test)" + ], + "production_code_slice_2": [ + "orchestrator/models.py:546-630 (PipelineConfig.agent_models — planner picks shape; key-validation required)", + "config/repo_config.py:248 (new get_default_model helper — wraps get_repo_setting, file is 836 lines and not in #2261 decomposition tables, safe to add here)", + "orchestrator/concurrent_executor.py:288-454 (get_agent_env + _spawn_agent + new model-resolver call site)", + "orchestrator/routes/pipelines.py:2702-2704 (restart call site)", + "orchestrator/consensus_wrapper.py:620 (build_consensus_wrapped_command — already takes model, just needs the new caller to pass it)", + "orchestrator/kubernetes_spawner.py:735 (plumb upstream + litellm_model into register_session)" + ], + "tests_slice_2": [ + "orchestrator/tests/test_model_resolver.py (new)", + "orchestrator/tests/test_concurrent_executor.py (extend with model-resolver assertions)", + "orchestrator/tests/test_kubernetes_spawner.py (extend with upstream-field plumbing assertions)" + ], + "docs_followup_not_in_scope_for_architect_files": [ + "docs/architecture/network-isolation.md (LiteLLM topology + traffic flow diagram)", + "docs/guides/agent-models.md (operator guide — how to set per-role models; Semantics B model_list pattern with natural keys; example egress overlay)", + "gateway/README.md (LITELLM_BASE_URL/LITELLM_MASTER_KEY config + body-rewrite documentation), orchestrator/README.md (agent_models)" + ] + }, + "runtime_primitives_with_evidence_2594": [ + { + "primitive": "_anthropic_client singleton + get_anthropic_client()", + "where": "gateway/gateway.py:9316-9329", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway in egg-system)", + "consumers": "proxy_anthropic_messages (gateway/gateway.py:9789), proxy_count_tokens (gateway/gateway.py:10032)", + "plan_action": "Replaced by UpstreamRegistry.get_client(upstream_name). The 'anthropic' entry preserves base_url, timeout, connection-pool limits exactly." + }, + { + "primitive": "proxy_anthropic_messages() Flask route at POST /v1/messages", + "where": "gateway/gateway.py:9752", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway), reached by in-sandbox-agent via ANTHROPIC_BASE_URL", + "consumers": "Every sandbox agent (Claude Code or egg_agent) — ANTHROPIC_BASE_URL=GATEWAY_K8S_URL is set unconditionally at orchestrator/kubernetes_spawner.py:807", + "plan_action": "Two new lines near 9775: resolve upstream from the session already fetched there; if upstream=='litellm' and session.agent_model_litellm is set, rewrite body['model'] (the v2 Semantics-B body rewriter). The streaming branch (9792+) uses the resolved client; credential injection at 9767 becomes upstream-aware." + }, + { + "primitive": "proxy_count_tokens() Flask route at POST /v1/messages/count_tokens", + "where": "gateway/gateway.py:10019", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway)", + "consumers": "Claude Code's compaction logic and the SDK token counters — invoked from any agent the gateway is fronting", + "plan_action": "Mirror of proxy_anthropic_messages: resolve upstream from the IP-keyed session (this route currently does NOT look up the session — slice-1 adds the same lookup it does in proxy_anthropic_messages; this is a small additive change), use that client, upstream-aware credential injection, AND v2 Semantics-B body rewrite. The rewrite is essential here precisely because Claude Code USES this response to drive compaction — if LiteLLM tokenized 'opus' (Claude tokenizer) instead of 'qwen-2.5-coder' (Qwen tokenizer), the count would mis-trigger compaction (R8)." + }, + { + "primitive": "_inject_anthropic_credentials(headers)", + "where": "gateway/gateway.py:9355", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway)", + "consumers": "proxy_anthropic_messages (9767), proxy_count_tokens (10028)", + "plan_action": "Extended signature to accept upstream parameter. Anthropic path unchanged. New 'litellm' branch reads LITELLM_MASTER_KEY via a new LiteLLMCredentialsManager (or extended AnthropicCredentialsManager) and sets Authorization: Bearer . Failure to find the key returns 502 (cq-8 fail-closed)." + }, + { + "primitive": "AnthropicCredential dataclass (header_name='x-api-key' or 'Authorization', header_value)", + "where": "gateway/anthropic_credentials.py:36-49", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway)", + "consumers": "_inject_anthropic_credentials at gateway/gateway.py:9371", + "plan_action": "Shape already accommodates Authorization: Bearer. A sibling LiteLLMCredential or a generalization is purely additive. Planner to choose the cleaner refactor — recommendation: sibling LiteLLMCredentialsManager so the supply-chain blast radius of LiteLLM key handling is in its own module and pytest target." + }, + { + "primitive": "_filter_blocked_tools(request_body, session_mode)", + "where": "gateway/gateway.py:9410", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway)", + "consumers": "proxy_anthropic_messages (9778)", + "plan_action": "UNCHANGED. cq-9 chose the conservative policy: tool-strip applies in private mode regardless of upstream. The filter doesn't need to know which upstream the request goes to. **Ordering note (v2)**: the slice-1 body-rewriter runs BEFORE _filter_blocked_tools so that the tool-strip sees the rewritten body; the rewriter only touches body['model'], so the downstream tool-strip's behavior is unaffected — but the canonical-body invariant is preserved." + }, + { + "primitive": "_SSEAccumulator (incremental SSE parser, MAX_CAPTURE_SIZE budget)", + "where": "gateway/gateway.py:9552", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway)", + "consumers": "proxy_anthropic_messages streaming branch (9881)", + "plan_action": "UNCHANGED. LiteLLM emits Anthropic-shaped SSE (it's an Anthropic /v1/messages translator). Verified upstream-agnostic in the slice-1 unit tests." + }, + { + "primitive": "Pre-stream retry + mid-stream synthetic SSE error frame (stream-resilience #1907)", + "where": "gateway/gateway.py:9811-9920", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway)", + "consumers": "proxy_anthropic_messages streaming branch", + "plan_action": "UNCHANGED. Both retry branches use the resolved client.send(http_req, stream=True). Works identically against LiteLLM. v2 slice-1 tests now explicitly include a client-disconnect-mid-stream case against the LiteLLM upstream (AC-2 part (g))." + }, + { + "primitive": "get_session_manager().get_session_by_ip(request.remote_addr)", + "where": "gateway/session_manager.py:741 (definition), gateway/gateway.py:9775 (call site in proxy_anthropic_messages)", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway)", + "consumers": "proxy_anthropic_messages — drives session_mode for the tool-strip; we extend it to drive the upstream resolution AND the body rewriter", + "plan_action": "Same lookup, same call site. proxy_count_tokens GAINS this lookup as part of slice-1 (it currently lacks it; adding it is a small additive change with no behavioral risk for Claude — the lookup is by IP and pure-read)." + }, + { + "primitive": "Session dataclass (mode, agent_role, phase, pipeline_id, etc.)", + "where": "gateway/session_manager.py:288-405", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway)", + "consumers": "All session-bound routes (gateway-wide)", + "plan_action": "Add two optional fields: agent_upstream: Literal['anthropic','litellm']|None, agent_model_litellm: str|None. Update to_dict_for_persistence/from_persistence at lines 338-405. Backward-compat: missing fields are None → default to 'anthropic'. Slice-1 unit test (AC-2 part (h)) asserts from_persistence on a pre-slice-1 dict yields None for both new fields." + }, + { + "primitive": "session_manager.register_session(...) + /api/v1/sessions/create payload", + "where": "gateway/session_manager.py:548 (Python helper), gateway/gateway.py:8560-8575 (HTTP payload acceptance), gateway/gateway.py:8846 (call site)", + "purpose": "production code (gateway pod)", + "execution_context": "deployed-pod (gateway) — invoked over HTTP by orchestrator's gateway_client", + "consumers": "Orchestrator at orchestrator/kubernetes_spawner.py:735 via orchestrator/gateway_client.py:602", + "plan_action": "Add new optional fields (agent_upstream, agent_model_litellm) to the Python signature, HTTP payload acceptance, and the gateway_client.register_session API. All optional → backward-compatible." + }, + { + "primitive": "build_consensus_wrapped_command(prompt_text, model='opus', ...)", + "where": "orchestrator/consensus_wrapper.py:620", + "purpose": "production code (trusted CI runner / orchestrator pod)", + "execution_context": "trusted-CI-runner / orchestrator pod — builds the bash script the agent Job runs", + "consumers": "orchestrator/concurrent_executor.py:454 (spawn), orchestrator/routes/pipelines.py:2704 (restart)", + "plan_action": "Signature already takes model. Both callers gain a one-line model resolution (resolve_model_for_role(pipeline_config, repo, role.value)[0]) before the call. The wrapper itself is unchanged." + }, + { + "primitive": "PipelineConfig.overseer_decision_maker_model, .overseer_advisor_model", + "where": "orchestrator/models.py:546, :620", + "purpose": "production code (orchestrator pod)", + "execution_context": "trusted-CI-runner / orchestrator pod", + "consumers": "Overseer spawn (orchestrator/kubernetes_spawner.py:1596), overseer monitor (orchestrator/overseer/monitor.py:259)", + "plan_action": "UNCHANGED. The new agent_models field sits alongside as a sibling. Overseer model selection is orthogonal." + }, + { + "primitive": "ANTHROPIC_BASE_URL=GATEWAY_K8S_URL injection", + "where": "orchestrator/kubernetes_spawner.py:807; consumed at sandbox/entrypoint.py:737-738 (setup_anthropic_api)", + "purpose": "production code (spawner writes, sandbox consumes)", + "execution_context": "trusted-CI-runner sets the env, in-sandbox-agent reads it — agent pods route all /v1/messages traffic to the gateway pod regardless of upstream", + "consumers": "Every spawned agent pod", + "plan_action": "UNCHANGED. The gateway is still the single point of egress for /v1/messages. The upstream is decided gateway-side, never sandbox-side. This is what preserves the zero-credential sandbox invariant." + }, + { + "primitive": "DEFAULT_MODEL = 'opus[1m]' (egg_agent client) and '--model' default 'opus[1m]' (egg_agent CLI)", + "where": "shared/egg_agent/client.py:62, shared/egg_agent/__main__.py:35", + "purpose": "production code (in-sandbox agent)", + "execution_context": "in-sandbox-agent", + "consumers": "python3 -m egg_agent invocations", + "plan_action": "UNCHANGED per cq-11. The Claude path keeps opus[1m]; non-Claude roles call build_consensus_wrapped_command with the recognized alias ('opus' without [1m]) per the recognized-alias mitigation." + }, + { + "primitive": "['--model', 'opus[1m]'] in legacy interactive sandbox/llm/runner.py:49", + "where": "sandbox/llm/runner.py:49", + "purpose": "production code (in-sandbox legacy CLI runner)", + "execution_context": "in-sandbox-agent (legacy path, may already be inert for the BRC flow)", + "consumers": "Legacy interactive runner only", + "plan_action": "UNCHANGED per cq-11. Planner to verify whether this path is exercised by any BRC role today; if not, leave it alone." + }, + { + "primitive": "_PROTECTED_ENV_KEYS in spawner", + "where": "orchestrator/kubernetes_spawner.py:138", + "purpose": "production code (orchestrator pod)", + "execution_context": "trusted-CI-runner", + "consumers": "spawn_agent_job at orchestrator/kubernetes_spawner.py:886-892", + "plan_action": "No change required for the recommended option (upstream is decided gateway-side from session metadata, not via env var). If the planner chooses to expose a debug/diagnostic env var to the agent pod for observability (e.g. EGG_AGENT_UPSTREAM=litellm — purely informational), that key should be added here so it cannot be tampered with by extra_env callers." + }, + { + "primitive": "Squid allowlist (gateway/allowed_domains.txt) — api.anthropic.com intentionally excluded", + "where": "gateway/allowed_domains.txt:9-15 (documented rationale), gateway/squid.conf", + "purpose": "production code (gateway pod's squid sidecar)", + "execution_context": "deployed-pod — outbound sandbox HTTPS", + "consumers": "Every outbound HTTP request from a sandbox agent", + "plan_action": "LiteLLM is reached only by the gateway pod, not by sandbox pods. The allowlist stays exactly as-is: LiteLLM host is NOT added to allowed_domains.txt (sandbox cannot reach LiteLLM directly, which is the same invariant as api.anthropic.com). The LiteLLM Service's network reachability is governed by a Kubernetes NetworkPolicy (gateway → litellm only), not Squid." + }, + { + "primitive": "k8s/base/gateway-deployment.yaml secrets.env mount at /secrets", + "where": "k8s/base/gateway-deployment.yaml:71-118", + "purpose": "production deployment", + "execution_context": "deployed-pod (gateway)", + "consumers": "Gateway credential managers (gateway/anthropic_credentials.py, gateway/jira_credentials.py, etc.)", + "plan_action": "Reuse the existing /secrets/secrets.env mount for the LITELLM_MASTER_KEY entry. No new mount; mirror the existing ANTHROPIC_API_KEY pattern. v2 addition: declare LITELLM_BASE_URL env var on the gateway container (default in-cluster Service DNS) so the swap-out seam is reachable without an image rebuild." + }, + { + "primitive": "GATEWAY_K8S_URL — http://gateway.egg-system.svc.cluster.local:9848", + "where": "orchestrator/kubernetes_spawner.py:124-132", + "purpose": "production code (orchestrator pod, also referenced sandbox-side)", + "execution_context": "trusted-CI-runner sets ANTHROPIC_BASE_URL to this; sandbox uses it", + "consumers": "Every agent pod", + "plan_action": "UNCHANGED. The gateway URL the sandbox sees is constant; the gateway internally fans out to anthropic-or-litellm. Decoupling the sandbox from upstream identity preserves the zero-credential and audit invariants." + }, + { + "primitive": "max_llm_cost_per_hour envelope (Anthropic-priced tokens)", + "where": "orchestrator/overseer/self_monitor.py:30-130; PipelineConfig comment refers to it at orchestrator/models.py:627", + "purpose": "production code (orchestrator overseer)", + "execution_context": "trusted-CI-runner", + "consumers": "Overseer self-monitor", + "plan_action": "OUT OF SCOPE per refine feedback Q4. The plan does NOT extend cost tracking to LiteLLM/Qwen pricing in this issue. Slice-1 + slice-2 land buildable but cost-tracking-unaware; a follow-up issue covers token-price extension. **Observable behavior (v2 addition, surfaced for operators)**: the existing cost tracker continues to apply Anthropic pricing to every token count it sees, including counts from LiteLLM-bound agents. The tracked dollar figure will be **incorrect** for non-Claude agents until the follow-up lands. This DOES NOT block the agent from running — max_llm_cost_per_hour will simply mis-estimate cost on the LiteLLM path." + }, + { + "primitive": "config/repo_config.py — get_repo_setting(repo, setting, default)", + "where": "config/repo_config.py:248", + "purpose": "production code (orchestrator + gateway)", + "execution_context": "trusted-CI-runner (orchestrator). Note: gateway also reads repositories.yaml via the EGG_REPO_CONFIG/EGG_SECRETS_PATH config mount.", + "consumers": "Various per-repo settings (restrict_to_configured_users, disable_auto_fix, etc.)", + "plan_action": "Add get_default_model(repo) wrapper that calls get_repo_setting(repo, 'default_model', None). Pure additive. **Decomposition cross-check (v2)**: config/repo_config.py is 836 lines (well under the 1500-line cap in scripts/file-size-allowlist.yaml); not listed in the #2261 decomposition tables in either orchestrator/CLAUDE.md or gateway/CLAUDE.md. Helper placement at line 248 is safe." + }, + { + "primitive": "config/repositories.yaml (operator-supplied) vs config/repositories.yaml.example (in-repo)", + "where": "config/repositories.yaml.example (in-repo); live file is operator-supplied via EGG_REPO_CONFIG path", + "purpose": "production config", + "execution_context": "trusted-CI-runner reads via config/repo_config.py:_load_config; gateway also reads via EGG_REPO_CONFIG mount", + "consumers": "config/repo_config.py:_load_config (config/repo_config.py:80)", + "plan_action": "v2 addition: the slice-2 get_default_model(repo) helper must handle the missing-file / missing-default_model case the same way get_repo_setting does today — returns None silently, resolver falls back to built-in 'opus' (AC-10). Dev environments that don't have a live repositories.yaml continue to work exactly as today." + } + ], + "key_constraints_carried_from_refine": { + "no_regression_on_claude_path": "Anthropic upstream is byte-for-byte identical: same base_url, same timeouts/limits, same credential injection, same SSE accumulator, same retry policy, same tool-strip. **Same body bytes** on the wire (the v2 body-rewriter is conditional on upstream=='litellm'; the Anthropic branch never touches body). The only structural change to the Claude path is that get_anthropic_client() becomes UpstreamRegistry.get_client('anthropic') under the hood — and the slice-1 tests assert this is a no-op.", + "zero_credential_sandbox": "Sandbox pods continue to set ANTHROPIC_BASE_URL=GATEWAY_K8S_URL with the placeholder OAuth token. Credentials (Anthropic and LiteLLM master key) live ONLY in the gateway pod's /secrets mount. Sandbox cannot reach LiteLLM directly (squid allowlist excludes it; only the gateway pod's NetworkPolicy lets it talk to LiteLLM).", + "gateway_mediated_visibility": "All non-Claude /v1/messages traffic continues to flow through the gateway pod, which keeps the existing audit-logging, transcript capture (_SSEAccumulator-based), and tool-strip uniform across upstreams.", + "compaction_math_mitigation": "Claude Code's --model flag stays a recognized Claude alias (e.g. 'opus' — without [1m]) even when the actual backend is Qwen. The gateway uses session.agent_model_litellm as the model string it writes into body['model'] on LiteLLM-bound requests. This is the cq-2-resolved approach + the v2 Semantics-B body-rewrite: the routing key is session metadata; the wire-level model identity is the rewritten body field, so LiteLLM dispatches on natural names and tokenizes against the correct backend. The body's model name as seen by Claude Code is informational only.", + "fail_closed_on_litellm_errors": "cq-8 chose fail-closed. UpstreamRegistry returns 502 if LITELLM_MASTER_KEY is missing or the LiteLLM upstream returns a connection error / timeout. No fallback to Claude. This matches today's Claude-side fail-closed behavior at gateway/gateway.py:9989-10016.", + "tool_strip_unchanged_in_private_mode": "cq-9 chose to keep the same tool-strip regardless of upstream. _filter_blocked_tools is upstream-agnostic and untouched. v2 ordering invariant: the LiteLLM body-rewriter runs BEFORE _filter_blocked_tools so the tool-strip operates on the canonical (rewritten) body. The rewriter only touches body['model'], so the tool-strip semantics are unchanged.", + "litellm_swap_out_seam": "Refine feedback Q3 said the design must keep LiteLLM swap-out-able if it hits a supply-chain incident (March 2026 PyPI). UpstreamRegistry is the seam: replacing LiteLLM with a different Anthropic-translator is one new registry entry + one secret rename + a k8s manifest swap + a kustomize-overlay change to LITELLM_BASE_URL. No call-site changes anywhere. The v2 LITELLM_BASE_URL env var declaration on the gateway pod is what makes the LiteLLM endpoint repointable from an overlay (vs requiring an image rebuild).", + "no_in_pipeline_validation": "cq-4 took empirical acceptance-test agent flipping OUT OF SCOPE for this pipeline. Slice-1 + slice-2 ship the seam and the config; the operator validates the end-to-end behavior in a separate pipeline run with a configured non-Claude backend. The plan/implement phases of THIS issue must not depend on a live LiteLLM endpoint being reachable.", + "build_now_validate_later": "Both slices must lint, type-check, and unit-test green on the CI runner WITHOUT a live LiteLLM endpoint. All slice-1 and slice-2 tests are pure unit tests (httpx.MockTransport in the gateway tests; standard mocks in orchestrator tests). No integration test in this scope requires a live non-Claude backend.", + "litellm_model_list_operator_pattern_semantics_b": "Under Semantics B (v2 pinned choice), operators configure LiteLLM's model_list with NATURAL model names: `model_list: [{model_name: 'qwen-2.5-coder', litellm_params: {model: 'openrouter/qwen/qwen3-coder', api_base: ..., api_key: ...}}, ...]`. The string `model_name` is what the gateway writes into body['model'] (from session.agent_model_litellm). The gateway → LiteLLM auth header is the master key (Authorization: Bearer ); LiteLLM uses its internal per-provider keys (e.g. OPENROUTER_API_KEY) to talk to the actual backend. There is NO requirement that LiteLLM's model_list keys equal Claude aliases — that constraint, which Semantics A would impose, was explicitly rejected." + }, + "open_risks_for_risk_analyst": [ + { + "id": "R1", + "title": "session_manager Session-field migration on existing persisted sessions", + "summary": "Adding agent_upstream/agent_model_litellm to the Session dataclass needs forward-compatible from_persistence/to_dict_for_persistence. If a gateway restart hits a pre-existing on-disk session file written by the old code, the new code must fall back to None for both. The existing pattern at gateway/session_manager.py:380-405 already handles this with .get(key) — extending it is safe, but the test plan must include a 'load a pre-existing session file' regression test (AC-2 part (h))." + }, + { + "id": "R2", + "title": "proxy_count_tokens currently has no session lookup", + "summary": "Slice-1 adds the IP-keyed session lookup to proxy_count_tokens (gateway/gateway.py:10019) so it can pick the right upstream AND apply the v2 Semantics-B body rewrite. The lookup is pure-read, but it's a behavioral change. Risk surface: ANY edge case (no session, expired session, stale-IP session) needs to default to 'anthropic' so the existing token-counting traffic is unaffected. The slice-1 tests must cover the 'no-session-for-this-IP' path explicitly and assert default-to-anthropic-no-body-rewrite." + }, + { + "id": "R3", + "title": "Drift between Claude Code's --model and gateway's session.agent_model_litellm", + "summary": "The whole compaction-math mitigation depends on Claude Code seeing a recognized alias while the gateway routes elsewhere. If the two source-of-truth points (consensus_wrapper --model arg and register_session payload) ever disagree, the agent gets a wrong-backend response or a request that fails token-counting. Mitigation: resolve_model_for_role is the SINGLE function that both callers consume. The slice-2 test plan must include an end-to-end test asserting the spawn-time tuple from resolve_model_for_role flows identically to both consensus_wrapper and register_session — no parallel computation." + }, + { + "id": "R4", + "title": "LiteLLM container supply-chain risk", + "summary": "Refine feedback Q3 raised the March 2026 PyPI incident. UpstreamRegistry provides the swap-out seam, but the LiteLLM image used in k8s/base/litellm-deployment.yaml must be pinned to a digest (not a floating tag) so a compromised tag push doesn't auto-roll to the gateway. The k8s manifest must specify image: ghcr.io/berriai/litellm@sha256:... — not :latest, not :stable. The v2 LITELLM_BASE_URL env-var declaration on the gateway pod makes the endpoint repointable from a kustomize overlay if an emergency rotation is needed." + }, + { + "id": "R5", + "title": "Stream-resilience retry assumption against LiteLLM", + "summary": "The pre-stream retry at gateway/gateway.py:9850 catches httpx.ReadError and httpx.RemoteProtocolError and retries once. Verified to behave identically against any backend exposing /v1/messages; risk is that LiteLLM's HTTP/2 stack or its provider-side connection-pool behavior triggers different exception classes. The slice-1 mock-transport tests assert the retry path with synthetic ReadError, but a real integration test (out-of-scope here per cq-4) is the only true validation. Documented as a known followup. v2 adds an explicit client-disconnect-mid-stream test against LiteLLM (AC-2 part (g)) so the _SSEAccumulator behavior is asserted under the new upstream too." + }, + { + "id": "R6", + "title": "LITELLM_MASTER_KEY rotation", + "summary": "Today the AnthropicCredentialsManager uses mtime-based cache invalidation (gateway/anthropic_credentials.py:126). The LiteLLM manager must follow the same pattern so the master key can be rotated by editing secrets.env without a gateway restart. Risk: if planner picks the 'extend AnthropicCredentialsManager' route instead of a sibling, the mtime invalidation already exists; if they pick a sibling, the new manager must replicate the mtime pattern." + }, + { + "id": "R7", + "title": "Backward compat with no-op default", + "summary": "Slice-1 must default agent_upstream to 'anthropic' for any session missing the field — including sessions registered by an older orchestrator against a new gateway (rollout ordering risk during k8s rolling update). The default-to-anthropic resolver covers this, but the rollout plan must spell out: deploy gateway first (additive; default-to-anthropic preserves old orchestrator behavior), then deploy orchestrator. Reverse order is also safe (new orchestrator + old gateway: new orchestrator omits the new fields, old gateway ignores them, no LiteLLM routing occurs)." + }, + { + "id": "R8", + "title": "Tokenizer mismatch on /v1/messages/count_tokens — Semantics A rejected (v2 addition)", + "summary": "If the gateway forwarded body['model']='opus' to LiteLLM unchanged (Semantics A), LiteLLM would compute tokens with the Anthropic tokenizer mapped to the model_list entry it dispatches on (its own internal mapping). If the operator's model_list maps `model_name: 'opus' → openrouter/qwen/qwen3-coder`, LiteLLM still uses the Anthropic tokenizer for counting because that's what `opus` resolves to in its tokenizer registry — but the actual generation backend uses the Qwen tokenizer, which has different token boundaries. The count_tokens response would be wrong; Claude Code's compaction math (which uses that count) would trigger at the wrong moment, producing over-length requests that hard-fail and wedge the agent. **v2 mitigation**: Semantics B's body-rewriter changes body['model'] to the natural model name BEFORE LiteLLM sees the request, so LiteLLM tokenizes against the correct backend. This is the load-bearing reason Semantics A was rejected." + } + ], + "task_planner_inputs": { + "slice_count": 2, + "slice_dependency": "slice-2 depends on slice-1 (the gateway-side router + body-rewriter + Session fields must land first so slice-2's per-agent model config has somewhere to be recorded and a wire-level mechanism to drive)", + "ordering_rationale": "cq-10 deferred slice decomposition to the planner. Of the four cq-10 options, the recommended shape is option 3 ([gateway router + LiteLLM topology + body-rewriter, no-op] → [per-agent model config + consensus_wrapper plumbing]) because: (a) slice-1 lands all the gateway-side work as a behavior-preserving no-op — easy to review, easy to unit-test without live infra, no functional change visible to any current pipeline; (b) slice-2 layers the operator-facing per-role config on top of a known-good seam — easy to reason about because the gateway side is already proven by slice-1's tests; (c) splitting them into two PRs keeps each PR within typical reviewer-attention bounds (gateway changes ~600-900 LOC; orchestrator changes ~400-700 LOC per reviewer_plan v1 estimate).", + "acceptance_criteria_seeds": [ + "AC-1: All existing Claude pipeline tests pass with no change (gateway-side and orchestrator-side suites).", + "AC-2: Slice-1 unit tests prove (a) Anthropic-upstream session → request lands at api.anthropic.com client; verify body['model'] is forwarded byte-unchanged (assert body['model']=='opus' on the wire under Anthropic upstream); (b) LiteLLM-upstream session with agent_model_litellm='qwen-2.5-coder' → request lands at LiteLLM client with Bearer-token credential AND body['model'] on the wire is 'qwen-2.5-coder' (the v2 Semantics-B rewrite); (c) same assertion applies to /v1/messages/count_tokens — body['model'] on the wire to LiteLLM is the rewritten natural name, not the Claude alias; (d) tool-strip in private mode applies to both upstreams and sees the rewritten body; (e) SSE accumulator + retry path runs against both upstreams; (f) backward-compat session from_persistence (R1); (g) client disconnect mid-stream against LiteLLM upstream — _SSEAccumulator flushes partial response identically to the Anthropic path; (h) JSON-parse failure on body in proxy_anthropic_messages when upstream='litellm' → body left unchanged, logger.warning emitted, request still proceeds to LiteLLM (defensive — should never happen on the production wire shape).", + "AC-3: Slice-1 unit tests prove session_manager from_persistence is backward-compatible: loading an on-disk session file written by the pre-slice-1 code yields agent_upstream=None and agent_model_litellm=None; the upstream resolver defaults to 'anthropic'; the body-rewriter is inert on that path.", + "AC-4: Slice-2 resolves PipelineConfig.agent_models[role] > repositories.yaml default_model > 'opus' precedence correctly; both consensus_wrapper --model and register_session agent_model_litellm/agent_upstream derive from the SAME resolve_model_for_role call.", + "AC-5: With no PipelineConfig.agent_models set and no repositories.yaml default_model, every existing pipeline spawns identically to today (Claude path, --model opus, no LiteLLM traffic, no body rewrite).", + "AC-6: Configuring PipelineConfig.agent_models={'refiner': 'qwen-2.5-coder'} for a test pipeline produces (a) consensus_wrapper --model 'opus' (recognized alias), and (b) gateway session with agent_upstream='litellm', agent_model_litellm='qwen-2.5-coder'. (c) The gateway's body-rewriter then changes body['model'] from 'opus' to 'qwen-2.5-coder' on every request from that agent's IP. The unit test asserts all three: the recorded session metadata, the resolved consensus_wrapper command, and the rewritten request body to LiteLLM. Does not require a live LiteLLM.", + "AC-7: LiteLLM Deployment manifest passes kubeval/kustomize-build; image is pinned to a digest; secret manifest references a Kubernetes Secret (not inline value); NetworkPolicy restricts inbound to the gateway pod's labels; k8s/base ships zero default egress for the LiteLLM pod (operator overlay supplies egress per cq-6 backend choice — example overlay scaffold lives under k8s/overlays/example/).", + "AC-8: Missing LITELLM_MASTER_KEY in secrets.env when a session declares agent_upstream='litellm' returns 502 to the agent (fail-closed per cq-8) and emits a logger.warning with no-key=True. No leak of the Anthropic credential into the LiteLLM-bound request.", + "AC-9: ANTHROPIC_BASE_URL sandbox env continues to point at the gateway service; sandbox pods cannot reach the LiteLLM Service directly (verified by a NetworkPolicy assertion or a documentation cross-link to allowed_domains.txt).", + "AC-10: When config/repositories.yaml is absent or contains no default_model for the queried repo, get_default_model(repo) returns None and the resolver falls back to the built-in 'opus' default. No exception, no error log. Dev environments without an opted-in repo config continue to work as today.", + "AC-11: PipelineConfig.agent_models with a typo'd role key (e.g. 'reviewer-refine' instead of 'reviewer_refine') fails Pydantic validation at PipelineConfig construction time — does NOT silently route the typo'd role to Claude. Whichever shape the planner picks for agent_models (dict[AgentRole,str] with field_validator vs explicit per-role fields), key-validation is in scope." + ], + "out_of_scope_explicitly": [ + "Empirical acceptance-test agent flip (cq-4 deferred to operator post-merge).", + "Cost tracking for LiteLLM/Qwen tokens (refine feedback Q4 deferred to follow-up). Observable behavior of the existing cost tracker on LiteLLM-bound agents: continues to apply Anthropic pricing to LiteLLM token counts; the tracked dollar figure will be incorrect for non-Claude agents until the follow-up issue lands. DOES NOT block the agent from running — max_llm_cost_per_hour will simply mis-estimate cost on the LiteLLM path.", + "Self-hosted vLLM/SGLang topology (refine feedback Q1 — hosted Qwen first; self-hosted is a later issue).", + "Multiple LiteLLM model_list entries beyond the first cut (single hosted-Qwen entry is enough to prove the seam).", + "Refactoring opus[1m] occurrences (cq-11 chose 'leave it' — the Claude path keeps opus[1m]; only non-Claude paths use a different model string).", + "Per-upstream tool-strip policy (cq-9 chose uniform; documented as a future option only)." + ] + }, + "external_research_done": [ + "LiteLLM Anthropic /v1/messages translator: https://docs.litellm.ai/docs/anthropic_unified — confirms LiteLLM exposes /v1/messages with full streaming and SSE-event shape compatibility.", + "Compaction risk surface: https://platform.claude.com/docs/en/build-with-claude/compaction — confirms compaction triggers on the model name Claude Code knows; the recognized-alias mitigation is sound.", + "Real-world setup writeup: https://dev.to/dcruver/running-claude-code-with-local-llms-via-vllm-and-litellm-599b — independent confirmation that LiteLLM + Claude Code + vLLM/Qwen works end-to-end and that 'route on metadata, not on body' is the standard mitigation.", + "claude-code-router rejection rationale: https://github.com/musistudio/claude-code-router/issues (historical write-ups about streaming-tool-call corruption on Qwen thinking-mode models)." + ], + "decisions_referenced": [ + "cq-1: Separate Deployment+Service in egg-system namespace", + "cq-2: Per-agent session metadata (IP-keyed lookup)", + "cq-3: Per-role PipelineConfig field + repositories.yaml default_model", + "cq-4: No empirical validation in this pipeline; operator validates post-merge", + "cq-5: Keep Claude Code harness for non-Claude models too", + "cq-6: Hosted Qwen-compatible provider for first cut", + "cq-7: Gateway holds LITELLM_MASTER_KEY in secrets.env", + "cq-8: Fail-closed on LiteLLM errors", + "cq-9: Uniform tool-strip in private mode regardless of upstream", + "cq-10: Slice decomposition deferred to planner (this analysis recommends option-3 shape)", + "cq-11: Leave opus[1m] in Claude defaults" + ] +} diff --git a/.egg-state/agent-outputs/2769-risk_analyst-output.json b/.egg-state/agent-outputs/2769-risk_analyst-output.json new file mode 100644 index 0000000000..83b2bfa12c --- /dev/null +++ b/.egg-state/agent-outputs/2769-risk_analyst-output.json @@ -0,0 +1,459 @@ +{ + "issue": 2769, + "phase": "plan", + "agent": "risk_analyst", + "title": "Risk Assessment: Support non-Claude models per agent via a LiteLLM proxy", + "summary": "Technical-risk assessment for adding a LiteLLM proxy in egg-system and turning the gateway into a per-agent upstream router. Overall risk is HIGH — the Claude path can stay unchanged (no regressions for current production), but the new path crosses two recently-burned third-party surfaces (LiteLLM was supply-chain compromised in March 2026 and has active SSE-translation tool-call bugs against non-Anthropic backends, both of which directly hit egg's primary use case of Claude Code running against Qwen) and rests on Claude Code's undocumented compaction-math behavior. The integration is buildable now and no-op by default, so production stability is preserved; but the empirical compatibility validation has been deferred to the operator (cq-4) and the plan must surface the residual risks the operator must own.", + + "overall_risk_level": "HIGH", + "recommendation": "PROCEED_WITH_MITIGATIONS", + "recommendation_rationale": "The recommended option (gateway-side upstream router + per-agent session metadata + LiteLLM Deployment in egg-system) keeps the Claude path structurally unchanged, which contains blast radius. The HIGH rating is driven by external risks that the team does not fully control (LiteLLM supply-chain, LiteLLM SSE-translation bugs, Claude Code compaction math) — not by the egg-side code change itself. With the listed mitigations (pinned + sha256-verified LiteLLM image, NetworkPolicy isolation, recognized-alias mitigation, fail-closed policy already chosen in cq-8), the change ships safely as a no-op-by-default integration; the real risk crystallises at operator-driven validation time (out of scope per cq-4), where the rollback plan is to leave the per-agent model config unset.", + + "risks": [ + { + "id": "R1", + "title": "LiteLLM supply-chain compromise (March 2026 PyPI incident + downstream CVEs)", + "category": "security", + "severity": "CRITICAL", + "likelihood": "MEDIUM", + "impact": "Direct credential harvesting, Kubernetes lateral movement, and persistent backdoor RCE in the egg-system namespace; spillover risk to gateway secrets (Anthropic OAuth, GitHub tokens, Atlassian credentials in /secrets/secrets.env) and to the cluster's K8s API surface. A compromised LiteLLM pod sharing the egg-system namespace with the gateway is one ServiceAccount RBAC misconfig away from cluster-wide impact.", + "description": "TeamPCP published backdoored LiteLLM 1.82.7 / 1.82.8 to PyPI on 24 March 2026 by compromising the maintainer's PyPI credentials (via a prior compromise of Trivy's CI pipeline). The malicious packages ran a three-stage payload: credential harvesting, Kubernetes lateral movement, and persistent RCE. Subsequent April–May 2026 audit-disclosed CVEs include CVE-2026-42208 (CVSS 9.3, unauthenticated SQL injection via the Authorization: Bearer header to /chat/completions — exploited in the wild within 36 hours of disclosure), CVE-2026-35029 (auth-bypass + RCE via /config/update), and CVE-2026-42271 (subprocess spawning via MCP preview endpoints). All fixed in v1.83.7+. The official LiteLLM-Cloud Docker image with strict version locking was NOT affected by the March incident; pip-installed unpinned versions WERE. Cited by the operator in feedback Q3 as recent and concerning.", + "affected_files": [ + "k8s/base/litellm-deployment.yaml (to be created)", + "k8s/base/litellm-service.yaml (to be created)", + "gateway/secrets.env layout (LITELLM_MASTER_KEY to be added)" + ], + "mitigation": { + "strategy": "Six layered controls: (1) Pin the LiteLLM image to a specific minor >= 1.83.7 with sha256 digest verification (image: ghcr.io/berriai/litellm@sha256:...) — not a floating tag. (2) Enable cosign signature verification at the imagePullPolicy layer or via a Kyverno/Cosign admission policy. (3) Run LiteLLM with a NetworkPolicy that allows ingress only from the gateway pod's ServiceAccount and egress only to the configured upstream provider domains — no general internet, no K8s API. (4) Give LiteLLM its own ServiceAccount with default (no-permission) RBAC; do NOT mount the gateway's kubeconfig or any cluster-scoped token. (5) Don't enable JWT auth (default off; CVE-2026-35030 only affects enabled deployments). (6) Document a kill-switch: if a future LiteLLM CVE drops, the operator can unset PipelineConfig.agent_models[*] cluster-wide and every agent reverts to the Claude path with no LiteLLM call sites. The cq-10 slice decomposition naturally supports this because the gateway router and the per-agent config land in dependent-but-separable slices.", + "effort": "MEDIUM", + "residual_risk": "MEDIUM — even a fully pinned image is only safe against the published CVEs; a future zero-day or repeat supply-chain incident is possible. The kill-switch and NetworkPolicy keep blast radius bounded if it happens." + }, + "requires_human_review": true, + "review_reason": "The operator already raised LiteLLM supply-chain concerns in feedback Q3 ('March 2026 PyPI incident is recent'). The plan must record image-pin + signature + NetworkPolicy + ServiceAccount-isolation as explicit, reviewer-verifiable acceptance criteria, not as 'best-effort'." + }, + { + "id": "R2", + "title": "LiteLLM /v1/messages streaming tool_use bug — drops input_json_delta for non-Anthropic backends", + "category": "compatibility", + "severity": "CRITICAL", + "likelihood": "HIGH", + "impact": "Claude Code (the chosen harness per cq-5) enters an infinite retry loop on every tool call: every tool_use block arrives with input: {}, fails Claude Code's required-field validation, and triggers an immediate retry. The agent never makes forward progress; consensus stalls; the operator sees what looks like an agent loop, not a backend bug.", + "description": "Multiple open GitHub issues against LiteLLM document this class of bug for the Anthropic /v1/messages streaming adapter when the upstream is non-Anthropic: issue #25561 (Gemini via vertex_ai), issue #25321 (regression in v1.82.x, works in v1.81.14), issue #24765 (GitHub Copilot path, content_block_start dropped → non-streaming fallback → output truncation → retry loop). Root cause is in LiteLLM's streaming_iterator.py: when the response transitions content blocks, processed chunks containing tool argument deltas are discarded rather than emitted as input_json_delta. Non-streaming requests for the same backend return the full arguments correctly. PR #13638 fixed an earlier instance of the same class. This is the egg-specific worst case: cq-5 chose Claude Code + LiteLLM (i.e. exactly the configuration these bugs were filed against), and the SDLC agents are tool-heavy (Bash, Read, Write, TodoWrite, MCP tools).", + "affected_files": [ + "gateway/gateway.py (proxy_anthropic_messages streaming path: 9792-9960)", + "gateway/gateway.py (_SSEAccumulator: 9552-9700)", + "k8s/base/litellm-deployment.yaml (image version pin)" + ], + "mitigation": { + "strategy": "Four layers: (1) Pin LiteLLM to a known-good minor (>= 1.83.7 satisfies the security CVE constraint; verify in the smoke test that v1.82.x regressions are no longer present) and document the pin rationale in the deployment YAML comment. (2) The acceptance-test validation step (deferred per cq-4) MUST include a long tool-heavy multi-turn run, not just a single round-trip — the bug only surfaces on content_block transitions across tool calls. (3) The plan must list this as a known-unresolved upstream issue that the operator owns at validation time; if a flip-validation reveals the bug, the rollback is to unset PipelineConfig.agent_models[role] and the agent reverts to Claude. (4) Optional defense: an in-gateway sentinel that detects empty tool input in a content_block_stop after a tool_use start and emits a structured error / OVERSEER_ALERT, rather than letting Claude Code loop silently. Belt-and-braces; not required for the no-op-by-default cut.", + "effort": "LOW (for the image pin + acceptance-test note); MEDIUM (if the empty-input sentinel is added)", + "residual_risk": "HIGH — the bug class is well-known and recurring; a regression on the next LiteLLM minor could re-introduce it. The fail-closed policy (cq-8) at least makes the failure visible, but only at validation time, not at PR-merge time." + }, + "requires_human_review": true, + "review_reason": "The chosen architecture (cq-5 + Claude Code + LiteLLM + non-Anthropic backend) is the exact configuration these bugs were filed against. The operator should know this is a real, recurring failure mode before approving the plan — not after discovering it at acceptance time." + }, + { + "id": "R3", + "title": "Claude Code compaction math drift — unrecognized / mis-sized model alias wedges long sessions", + "category": "compatibility", + "severity": "HIGH", + "likelihood": "MEDIUM", + "impact": "On long-running agents (coder, refiner — multi-thousand-turn sessions), Claude Code auto-compacts around 95% of the model's known context window. When the alias presented to Claude Code (e.g. 'opus') has a recognised window that is LARGER than the real LiteLLM-backed model's window (e.g. Qwen3 with 128K vs opus[1m] at 1M), auto-compaction fires too late, the in-progress request exceeds the real backend's limit, the request hard-fails with model_context_window_exceeded, the agent wedges, and consensus stalls. The reverse — alias smaller than backend — wastes context but does not stall.", + "description": "Per the issue body and corroborated by published write-ups (okhlopkov.com, dev.to vLLM+LiteLLM guide): Claude Code derives the context window from the model name and triggers auto-compaction at ~95% of that window. The compaction control's context_token_threshold default is 100,000 tokens for unrecognised models; for recognised Claude models it uses the model's real window. The cq-2/cq-5 mitigation is correct (route on session metadata, present a recognised alias), but it depends on a constraint the issue body does NOT enforce: the alias presented to Claude Code MUST have a context window ≤ the real backend's window. If the alias is 'opus[1m]' (1M tokens) and Qwen3 is 128K, auto-compaction never fires inside Qwen's window and long sessions hard-fail. The compaction-math is internal Claude Code behavior that is not documented as a public API; it can change between Claude Code versions without notice.", + "affected_files": [ + "shared/egg_agent/client.py:62 (DEFAULT_MODEL)", + "shared/egg_agent/__main__.py:35 (--model default)", + "orchestrator/consensus_wrapper.py:653-662 (model arg in agent command)", + ".egg-state/contracts/* (per-agent alias config landing here)" + ], + "mitigation": { + "strategy": "Three measures: (1) The plan must add an explicit invariant: the Claude-Code-facing alias's recognised window must be ≤ the real backend's window. Sample-coded examples in the docs: real Qwen3 128K → present 'sonnet' alias (200K) is WRONG; present a 100K-window alias (or set context_token_threshold explicitly) is RIGHT. (2) The plan must surface this as a per-(alias, real-model) pair recorded in PipelineConfig.agent_models or as a small lookup table next to it — not as 'whatever string the operator writes'. (3) Acceptance-test (deferred per cq-4) MUST include a session long enough to cross the auto-compaction boundary — otherwise the bug is undetected. The current build-now / validate-later split keeps Claude paths safe, but the validation step is non-optional and must be explicit on the plan.", + "effort": "LOW (invariant + doc note); MEDIUM (lookup table + validator on PipelineConfig)", + "residual_risk": "MEDIUM — Claude Code is closed-source and its compaction-math heuristics can change. The mitigation is good for the v4.x Claude Code branch but is not future-proof. Plan should note this as a Claude-Code-version-dependent assumption." + }, + "requires_human_review": false + }, + { + "id": "R4", + "title": "Qwen3 + vLLM streaming tool-call parser bugs — deferred but flagged for self-hosted target", + "category": "compatibility", + "severity": "MEDIUM", + "likelihood": "HIGH", + "impact": "If/when the operator moves from a hosted Qwen provider (cq-6) to self-hosted Qwen on vLLM (the long-term cost goal), several open vLLM bugs come back into scope: streaming + tool_choice with thinking disabled drops tool_calls (vllm issues #21565, #17655, #23992, #20611); Qwen3.5 / 3.6 emit XML tool_calls inside blocks that the qwen3_coder parser does not recover (#39056). Self-hosted Qwen with current vLLM produces silently truncated or empty tool calls under load.", + "description": "The cq-6 decision deferred self-hosted vLLM to a future iteration; the first acceptance-test backend is a hosted Qwen provider. That means R4 is OUT OF SCOPE for this issue's deliverables but IN SCOPE for the operator's eventual validation. Documenting it here because the plan's UpstreamRegistry abstraction will outlive the hosted-provider choice and the same gateway code will be used against vLLM eventually.", + "affected_files": [ + "k8s/base/litellm-config.yaml (vLLM endpoint config, when added)", + "future Qwen vLLM Deployment manifests" + ], + "mitigation": { + "strategy": "Two: (1) Document in the plan / follow-up issue that the eventual self-hosted vLLM path must validate with a streaming + tool-heavy + reasoning-enabled run; the hosted-provider validation does NOT cover this. (2) Pin vLLM to a version with #39056 / #21565 fixes landed when self-hosting (today: no single vLLM release is known-good; needs reassessment when the self-hosting work begins). Out-of-scope this issue but worth a follow-up tracker.", + "effort": "LOW (documentation only this issue)", + "residual_risk": "LOW for this issue's scope (hosted-provider only); HIGH for the eventual self-hosted-vLLM cut. The split is acceptable." + }, + "requires_human_review": false + }, + { + "id": "R5", + "title": "Hosted Qwen provider as a new credential + supply-chain surface", + "category": "security", + "severity": "MEDIUM", + "likelihood": "MEDIUM", + "impact": "A new third-party credential (e.g. Together / Fireworks / DeepInfra / OpenRouter) lands in /secrets/secrets.env or in the LiteLLM config. The provider becomes a new attack surface for source-code / transcript exfiltration (per feedback Q5: no compliance constraint, but per cq-9: the WebSearch/WebFetch strip is kept regardless of upstream — partial mitigation). Provider outage, pricing change, deprecation, or rate-limit change becomes a new operational dependency.", + "description": "The cq-6 selection ('hosted Qwen-compatible provider — fastest path; defers self-hosting but adds a third-party dependency and a new credential to hold') is explicit about the trade-off. The risk is operational, not architectural: hosted-provider credentials need rotation, monitoring, billing alerts, and an SLA. Feedback Q5 confirmed there is no data-residency constraint that rules this out. cq-9 confirmed the WebSearch/WebFetch tool-strip stays on regardless of upstream, which limits the additional exfiltration surface to whatever the provider does with API request bodies (i.e. prompt logging, data retention).", + "affected_files": [ + "gateway/secrets.env layout", + "k8s/base/litellm-config.yaml (provider config)" + ], + "mitigation": { + "strategy": "Three: (1) Choose a provider with a published data-retention policy and a 'no training, no logging' SLA option; document the chosen provider in the deployment YAML comment. (2) Per cq-7, the gateway holds a single LiteLLM master key and LiteLLM holds the real per-provider key — this contains credential blast radius (gateway never sees the provider key directly). (3) Standard secret rotation: track the provider key in secrets.env with a documented rotation cadence; surface key-not-set as a fail-closed error (cq-8) so an expired key surfaces immediately, not silently.", + "effort": "LOW", + "residual_risk": "LOW — bounded by the provider's published policies and the existing secret-management practice." + }, + "requires_human_review": false + }, + { + "id": "R6", + "title": "Gateway `register_session()` API surface missing upstream / model_alias fields (runtime-primitive gap per #2594)", + "category": "compatibility", + "severity": "MEDIUM", + "likelihood": "CERTAIN", + "impact": "The plan will assume the per-agent session metadata carries (upstream, model_alias). Today's Session schema (gateway/session_manager.py:548-563) and the orchestrator's gateway_client.register_session() (orchestrator/gateway_client.py:602-620) have no such fields. Adding them is mandatory for the recommended Option A. If the plan does not explicitly call this out as new schema, the implementer may revert to Option C (route on body model name) which conflicts with the cq-2 resolution.", + "description": "Per #2594 runtime-primitive flagging: the cq-2 resolution says 'orchestrator declares the model+upstream when it spawns the agent (session lookup by IP, same path used today for session_mode)'. But there is no such field today. Specifically: (a) Session dataclass (session_manager.py:287-328) has session_mode, container_id, agent_role, etc., but no upstream or model_alias. (b) register_session endpoint and the matching gateway_client method accept ~17 fields, none for upstream. (c) The gateway's proxy_anthropic_messages() looks up the session at gateway/gateway.py:9775 via get_session_by_ip() — adding session.upstream is the natural plug-in point. (d) Sessions persist to ~/.egg-state/sessions/sessions.json so the new fields must round-trip through the persistence layer.", + "affected_files": [ + "gateway/session_manager.py (Session dataclass + persistence)", + "gateway/gateway.py (register_session endpoint at line ~8846)", + "orchestrator/gateway_client.py (register_session() at line ~602)", + "orchestrator/kubernetes_spawner.py (call site that hands the new fields to gateway_client)" + ], + "mitigation": { + "strategy": "Plan must include an explicit task: 'extend the gateway session schema with optional upstream and model_alias fields, round-trip through Session.to_dict_for_persistence(), thread through register_session() on both sides, and default to None on existing sessions (backward compatible).' Make this its own contract task so the reviewer can verify the schema change is in the diff. The Session schema migration is purely additive (optional fields) so existing sessions are unaffected.", + "effort": "LOW", + "residual_risk": "LOW — clear, mechanical change." + }, + "requires_human_review": false + }, + { + "id": "R7", + "title": "build_consensus_wrapped_command() model arg untouched at both spawn sites (runtime-primitive gap per #2594)", + "category": "compatibility", + "severity": "MEDIUM", + "likelihood": "CERTAIN", + "impact": "PipelineConfig.agent_models[role] can be set in the contract / config, but if the spawn path doesn't thread it through, the agent still launches with --model opus and the gateway routes (correctly) to Anthropic — i.e. the per-agent config is silently inert. Worse: drift between the gateway's routing decision (Anthropic, because session.upstream is unset) and the agent's --model string (qwen-..., because PipelineConfig.agent_models[role] is set) → 'right' configuration produces 'wrong-backend' request.", + "description": "Per #2594 runtime-primitive flagging: build_consensus_wrapped_command() (orchestrator/consensus_wrapper.py:620) accepts a model arg with a hardcoded 'opus' default but BOTH callers (orchestrator/concurrent_executor.py:454 and orchestrator/routes/pipelines.py:2704) call it with prompt_text only. The default-to-opus path is invisible at the call site. The plan must thread the per-agent model into both call sites AND into the kubernetes_spawner's register_session() call so both ends agree.", + "affected_files": [ + "orchestrator/concurrent_executor.py:454 (initial spawn)", + "orchestrator/routes/pipelines.py:2704 (restart spawn)", + "orchestrator/consensus_wrapper.py:620-662 (wrapper)", + "orchestrator/kubernetes_spawner.py (session-create call)" + ], + "mitigation": { + "strategy": "Plan task: 'compute the resolved model once (from PipelineConfig.agent_models[role] with default = opus) at spawn time, pass it as both (a) the --model arg to build_consensus_wrapped_command and (b) the model_alias field on register_session, so the agent CLI and the gateway routing decision are derived from the same source.' Restart path (pipelines.py:2704) re-derives from the same PipelineConfig (which is persisted in the contract / pipeline state), not from a separate code path. Add a defensive assert in the consensus wrapper that errors if model and the session's model_alias disagree.", + "effort": "LOW-MEDIUM", + "residual_risk": "LOW — single source of truth (PipelineConfig.agent_models) is the right pattern." + }, + "requires_human_review": false + }, + { + "id": "R8", + "title": "Claude Code harness assumption — recognised-alias mitigation rests on an undocumented internal heuristic", + "category": "compatibility", + "severity": "MEDIUM", + "likelihood": "MEDIUM", + "impact": "The cq-5 mitigation ('keep Claude Code harness, use a recognised alias') assumes Claude Code's compaction trigger, extended-thinking feature gating, and any other model-name-keyed feature gates will behave correctly when the alias is presented in the request body but the real backend is non-Claude. If Claude Code adds a feature that probes the model server in a way that fails on LiteLLM (e.g. capability discovery, model-specific tool descriptions), the mitigation silently breaks on a Claude Code version bump.", + "description": "Claude Code is closed-source. Its model-name-keyed behavior is observed empirically and documented externally (community write-ups), not by Anthropic. A Claude Code minor version bump could change the heuristic. The mitigation is correct today but not durable across Claude Code versions, and we have no way to detect a regression except by running the validation again.", + "affected_files": [ + "sandbox/CLAUDE.md (records Claude Code version)", + "shared/egg_agent/client.py (Claude Agent SDK version)" + ], + "mitigation": { + "strategy": "Two: (1) Plan must record the Claude Code version that the recognised-alias mitigation was validated against; pin or bound the Claude Code version in the sandbox image (the session schema already carries claude_code_version). (2) Add a CI smoke test or scheduled job that exercises a tool-heavy multi-turn loop against the non-Claude backend (when an endpoint is available) so a Claude Code version bump that regresses the heuristic is caught. Out of scope for this issue per cq-4; record as a follow-up.", + "effort": "LOW (this issue: doc + version pin); MEDIUM (follow-up: CI smoke test)", + "residual_risk": "MEDIUM — bounded by the Claude Code release cadence and our willingness to keep the pin current." + }, + "requires_human_review": false + }, + { + "id": "R9", + "title": "Cost-tracking envelope (max_llm_cost_per_hour) breaks silently on the LiteLLM path", + "category": "operational", + "severity": "MEDIUM", + "likelihood": "CERTAIN", + "impact": "Per feedback Q4, cost-tracking is explicitly deferred to a follow-up. The plan must not silently regress today's behavior. Today the overseer monitor reads cost_usd from SDK results (orchestrator/overseer/self_monitor.py:123-130, ~137 _get_hourly_llm_cost). The SDK reports Anthropic-priced costs. For LiteLLM-bound agents, the SDK still reports a number, but it is either zero (if LiteLLM omits the field) or wrong (if LiteLLM passes the field through with the alias's Anthropic pricing applied to Qwen tokens). The $5/hr envelope fires on phantom data; the operator can't know real spend.", + "description": "Cited explicitly in the operator's feedback Q4 reply: 'Defer to a follow-up — do not extend max_llm_cost_per_hour cost tracking to LiteLLM/Qwen token pricing in this issue.' This is acknowledged technical debt, not a design defect, but the plan must document it as a known regression on the LiteLLM path so the follow-up issue is visible.", + "affected_files": [ + "orchestrator/overseer/self_monitor.py:123-130", + "shared/egg_contracts/usage.py:171", + "sandbox/agent-config/rules/overseer.md:94, 314" + ], + "mitigation": { + "strategy": "Two: (1) The plan creates a follow-up issue 'Extend max_llm_cost_per_hour to LiteLLM-priced tokens' with the per-provider pricing table as the unit of work. (2) For this issue, the plan should document in the LiteLLM deployment README / config that cost_usd is unreliable on the LiteLLM path until the follow-up lands, so operators don't make decisions on it. No code change in this issue.", + "effort": "LOW (doc + follow-up issue)", + "residual_risk": "LOW for this issue (operator is informed); MEDIUM for the follow-up (real $$ exposure if an agent loops on the LiteLLM path without cost visibility)." + }, + "requires_human_review": false + }, + { + "id": "R10", + "title": "SSE accumulator hardcoded to Anthropic event names — silent transcript corruption if LiteLLM emits foreign events", + "category": "compatibility", + "severity": "MEDIUM", + "likelihood": "LOW", + "impact": "_SSEAccumulator (gateway/gateway.py:9552-9700) hardcodes Anthropic event names (message_start, content_block_start, content_block_delta, message_delta, error). If LiteLLM's Anthropic-translation layer mis-translates a backend response (e.g. emits an OpenAI-style 'data: {choices: ...}' line), the accumulator silently drops the event: no error log, transcript captures empty content, usage stays None, stop_reason None. Audit logs and checkpoints are quietly wrong.", + "description": "Per Q5 of the gateway-research probe: the accumulator parses specific Anthropic event names. LiteLLM SHOULD emit Anthropic-shaped SSE for /v1/messages — that is the whole point of the endpoint — but the documented bugs (R2 above) and the GitHub Copilot path (#24765 — content_block_start dropped) show that the translation is imperfect in real-world conditions. The failure mode is asymmetric: streaming bugs in LiteLLM tend to drop events rather than emit foreign ones, so the most likely manifestation is empty content blocks (already covered by R2), not foreign event names. Listed here as a separate risk because the fallout (silent transcript corruption, wrong audit logs) is different.", + "affected_files": [ + "gateway/gateway.py (_SSEAccumulator: 9552-9700)", + "gateway/gateway.py (transcript capture: _capture_streaming_response, ~9942-9953)" + ], + "mitigation": { + "strategy": "Two: (1) Plan should add a low-cost defensive log/warning in _SSEAccumulator when an unknown event_type is seen — current code branches on known names and silently drops unknown events. A warning-level log preserves the transcript shape but surfaces the misdetection. (2) Acceptance-test (deferred) MUST diff a LiteLLM-bound transcript against a Claude-bound transcript for the same prompt to catch silent drops.", + "effort": "LOW", + "residual_risk": "LOW — the bug surface is well-bounded; defensive logging is cheap." + }, + "requires_human_review": false + }, + { + "id": "R11", + "title": "ANTHROPIC_BASE_URL not in `_PROTECTED_ENV_KEYS` — defense-in-depth gap", + "category": "security", + "severity": "LOW", + "likelihood": "LOW", + "impact": "Per-agent extra_env (passed via the spawner) can mutate ANTHROPIC_BASE_URL today, redirecting agent traffic away from the gateway. Pre-LiteLLM this is theoretical (no caller sets it); post-LiteLLM, the plan may introduce a per-agent ANTHROPIC_BASE_URL or X-Egg-Upstream injection — the surface widens.", + "description": "_PROTECTED_ENV_KEYS in orchestrator/kubernetes_spawner.py:138 includes EGG_SESSION_TOKEN, GATEWAY_URL, http_proxy, etc., but NOT ANTHROPIC_BASE_URL. Per cq-2 the routing signal is per-agent session metadata (set by orchestrator at session-create time), not a per-agent env var, so this is not the chosen routing mechanism — but defense-in-depth says ANTHROPIC_BASE_URL should be protected regardless.", + "affected_files": [ + "orchestrator/kubernetes_spawner.py:138 (_PROTECTED_ENV_KEYS)" + ], + "mitigation": { + "strategy": "Plan adds ANTHROPIC_BASE_URL to _PROTECTED_ENV_KEYS as a one-line defense-in-depth change alongside the upstream-router work. Optional but cheap.", + "effort": "LOW", + "residual_risk": "NEGLIGIBLE." + }, + "requires_human_review": false + }, + { + "id": "R12", + "title": "max_turns hardcoded to 1000 in consensus wrapper — may not fit non-Claude models", + "category": "compatibility", + "severity": "LOW", + "likelihood": "MEDIUM", + "impact": "consensus_wrapper.py:623 hardcodes max_turns=1000. Non-Claude models may have very different per-turn budgets (smaller context → more turns to compensate, OR longer per-turn latency → fewer turns affordable). A non-Claude agent may hit the wall earlier or waste budget vs Claude.", + "description": "The cq-* decisions did not surface this knob. Listed here because per-model max-turns is a likely follow-up tuning need, and the plan should not lock in 1000 as if it were universal.", + "affected_files": [ + "orchestrator/consensus_wrapper.py:623" + ], + "mitigation": { + "strategy": "Plan should make max_turns either a PipelineConfig.agent_models entry (alongside model) or note as a TODO for the follow-up validation phase. No code change required this issue if it stays a follow-up.", + "effort": "LOW", + "residual_risk": "LOW — operator-controllable knob, not a correctness issue." + }, + "requires_human_review": false + }, + { + "id": "R13", + "title": "Empirical validation deferred per cq-4 — residual integration risk lands on the operator at flip-time", + "category": "operational", + "severity": "MEDIUM", + "likelihood": "CERTAIN", + "impact": "Per the cq-4 resolution, the empirical compatibility validation cannot occur inside this pipeline (no configured access to non-Claude models). The plan ships buildable, no-op-by-default integration. ALL the compatibility risks above (R2 LiteLLM streaming bug, R3 compaction math, R8 Claude Code heuristic drift, R10 SSE accumulator drift) remain unverified until the operator flips an agent at validation time. If a validation flip exposes a bug, the rollback (unset PipelineConfig.agent_models[role]) is safe but it means the LiteLLM path was never green at PR-merge time.", + "description": "This is an explicit decision (cq-4) and not a defect. Recording it here because the risk profile of a 'shipped but unvalidated' integration is different from a 'shipped and validated' one — reviewers, operators, and follow-up issues need to know which class they are in.", + "affected_files": [ + "(plan-phase deliverables: contract acceptance criteria)" + ], + "mitigation": { + "strategy": "Plan must (1) explicitly enumerate the deferred validations as known-unresolved acceptance criteria the operator will sign off on at flip-time, and (2) ensure the no-op-by-default rollback path is structurally guaranteed (i.e. with PipelineConfig.agent_models = {}, no LiteLLM bytes are sent — every code path falls back to the Claude singleton). Reviewer should verify this by inspection: searching for 'litellm' in the diff should yield only conditional / opt-in code paths.", + "effort": "LOW (documentation + structural guarantee)", + "residual_risk": "MEDIUM — bounded by the operator's diligence at flip-time." + }, + "requires_human_review": true, + "review_reason": "The deferred validation is on the operator at flip-time; the plan reviewer must confirm the structural guarantee of the no-op-by-default rollback is real, not nominal." + }, + { + "id": "R14", + "title": "UpstreamRouter abstraction needed per feedback Q3 — design choice not yet made", + "category": "architecture", + "severity": "LOW", + "likelihood": "CERTAIN", + "impact": "Feedback Q3 explicitly requested 'a thin UpstreamRouter/Registry abstraction so LiteLLM can be replaced if it hits a maintenance/supply-chain problem'. If the plan hard-wires LiteLLM into proxy_anthropic_messages() / get_anthropic_client() instead of introducing a tiny registry, the next time we want to swap LiteLLM out (or add a parallel non-LiteLLM upstream) we pay double — once for the change and once for the abstraction. Given R1 (already-burned supply-chain risk on LiteLLM), the swap-out point is not theoretical.", + "description": "No UpstreamRouter / UpstreamRegistry / ModelRegistry abstraction exists in the repo today. The cleanest cut is a 1-2-class file (gateway/_upstreams.py?) that exposes (a) an enum / string set of upstream names, (b) a registry mapping name → httpx client + base_url + credential type, (c) resolve(session) → upstream. This is consistent with the in-flight gateway.py decomposition (CLAUDE.md slice-14 already pre-allocates new sub-packages).", + "affected_files": [ + "gateway/_upstreams.py (to be created)", + "gateway/gateway.py (get_anthropic_client refactor: 9316-9329)" + ], + "mitigation": { + "strategy": "Plan should explicitly require an UpstreamRegistry-shaped abstraction (not a pair of if/else branches inline in proxy_anthropic_messages). Keyed by upstream name from the session metadata; Anthropic is the default entry; LiteLLM is registered when the deployment is present. Replaces today's _anthropic_client singleton with registry.get('anthropic') and adds registry.get('litellm') alongside.", + "effort": "LOW-MEDIUM", + "residual_risk": "LOW — a small abstraction, easy to verify, matches the operator's explicit request." + }, + "requires_human_review": false + }, + { + "id": "R15", + "title": "Failure-policy visibility — fail-closed (cq-8) requires operator monitoring to catch stalls", + "category": "operational", + "severity": "LOW", + "likelihood": "MEDIUM", + "impact": "Per cq-8 the failure policy is fail-closed (502 to agent, no fallback). This is correct, but it means a misconfigured LiteLLM upstream produces a hung agent (502 → consensus stall → eventual timeout) rather than an alert. The operator must monitor pipeline state. The recommended-option Option C variant (fail-closed + auto-HITL escalation on non-Claude spawn / stall) was explicitly NOT chosen.", + "description": "The cq-8 resolution chose 'fail closed (no fallback)' over 'fail closed + auto-HITL escalation'. The plan must surface that this is the chosen behavior and that operator monitoring is the detection path — not a structural alarm.", + "affected_files": [ + "gateway/gateway.py (proxy_anthropic_messages error branch)" + ], + "mitigation": { + "strategy": "Plan should record this as a known-and-accepted trade-off referencing cq-8, not as a discoverable behavior. The follow-up Option-C-style auto-HITL escalation can be a separate issue if pipeline stalls due to LiteLLM 5xx become common.", + "effort": "LOW (documentation)", + "residual_risk": "LOW." + }, + "requires_human_review": false + } + ], + + "runtime_primitive_audit_per_2594": { + "description": "Per issue #2594 (and the recurring #2474 failure pattern), the plan must explicitly declare which runtime primitives exist today vs which are net-new. Below: primitives the plan WILL rely on, and their existence status.", + "primitives": [ + { + "name": "get_session_by_ip().upstream / .model_alias", + "status": "MISSING — not in today's Session dataclass (gateway/session_manager.py:287-328). Plan must add them as optional fields with persistence.", + "addressed_in_risk": "R6" + }, + { + "name": "register_session(upstream=..., model_alias=...)", + "status": "MISSING — neither the orchestrator client (gateway_client.py:602-620) nor the gateway endpoint accept these fields today. Plan must extend both signatures.", + "addressed_in_risk": "R6" + }, + { + "name": "PipelineConfig.agent_models[role]", + "status": "MISSING — only overseer_decision_maker_model (models.py:546) and overseer_advisor_model (models.py:620) exist today. New optional dict field. Plan must add with a role-key validator.", + "addressed_in_risk": "R7" + }, + { + "name": "build_consensus_wrapped_command(model=...)", + "status": "EXISTS but UNUSED — accepts a model arg with default 'opus' (consensus_wrapper.py:622) but neither caller passes it (concurrent_executor.py:454, pipelines.py:2704). Plan must thread the resolved model into both call sites.", + "addressed_in_risk": "R7" + }, + { + "name": "LiteLLM Deployment + Service + LITELLM_MASTER_KEY in /secrets/secrets.env", + "status": "MISSING — no k8s/base/litellm-*.yaml exists; secrets.env layout (gateway-deployment.yaml:71-118) lists Anthropic / GitHub / Atlassian only. Plan must add Deployment, Service, secret entry, and document the rotation contract.", + "addressed_in_risk": "R1, R5" + }, + { + "name": "Internal cluster DNS allowlisting (litellm.egg-system.svc.cluster.local)", + "status": "NOT NEEDED if the gateway routes to LiteLLM through its httpx client (the recommended path — same pattern as today's get_anthropic_client). gateway/allowed_domains.txt is the Squid allowlist for direct sandbox traffic; internal cluster DNS is implicit cluster-network traffic and NOT routed through Squid. Plan should NOT add svc.cluster.local entries to allowed_domains.txt — it should keep the gateway-as-proxy pattern uniform with Anthropic.", + "addressed_in_risk": "(no explicit risk — note in plan)" + }, + { + "name": "AnthropicCredential extension for LiteLLM master key", + "status": "EXTENDABLE — gateway/anthropic_credentials.py defines a concrete AnthropicCredential dataclass with header_name + header_value, no abstract base. Plan must either (a) extend the dataclass to a third credential type with a distinct header name, or (b) introduce a sibling LiteLLMCredential type, plus a manager-level discriminator. The injection site (gateway.py:9368-9371) already does headers[cred.header_name] = cred.header_value, so it is LiteLLM-friendly today.", + "addressed_in_risk": "(implicit in R6 / plan structure)" + }, + { + "name": "Claude Code's recognised-alias compaction math", + "status": "DOCUMENTED EXTERNALLY (community write-ups, undocumented by Anthropic) — assumed to work today against Claude 4.5+ but not version-stable. Plan must pin Claude Code version and document this as a heuristic dependency.", + "addressed_in_risk": "R3, R8" + }, + { + "name": "_SSEAccumulator (Anthropic event names)", + "status": "EXISTS, EVENT-NAME-HARDCODED — gateway.py:9552-9700 hardcodes message_start / content_block_* / message_delta / error. LiteLLM /v1/messages SHOULD emit the same shape, but real-world LiteLLM bugs do drop / mis-translate events on non-Anthropic backends. Plan should add a defensive log on unknown event_type.", + "addressed_in_risk": "R10" + }, + { + "name": "max_llm_cost_per_hour enforcement", + "status": "EXISTS but UNINSTRUMENTED FOR LITELLM — orchestrator/overseer/self_monitor.py:123-130 reads cost from .cost attribute (likely SDK-reported cost_usd, Anthropic-priced). Plan documents this as out-of-scope-but-known-broken per feedback Q4.", + "addressed_in_risk": "R9" + }, + { + "name": "_PROTECTED_ENV_KEYS protection for ANTHROPIC_BASE_URL", + "status": "MISSING — orchestrator/kubernetes_spawner.py:138 doesn't list ANTHROPIC_BASE_URL today. Defense-in-depth gap noted.", + "addressed_in_risk": "R11" + }, + { + "name": "max_turns per model", + "status": "HARDCODED — consensus_wrapper.py:623 max_turns=1000 for all roles / models. Out of scope for this issue but flagged.", + "addressed_in_risk": "R12" + } + ] + }, + + "trust_boundary_audit": { + "description": "Trust-boundary surfaces the plan crosses, per #2594.", + "boundaries": [ + { + "name": "Sandbox → Gateway", + "status": "UNCHANGED — sandbox stays zero-credential, still reaches gateway via ANTHROPIC_BASE_URL=GATEWAY_K8S_URL. Plan adds no new credential to the sandbox." + }, + { + "name": "Gateway → Anthropic", + "status": "UNCHANGED — same singleton client, same credential injection, same path." + }, + { + "name": "Gateway → LiteLLM (NEW)", + "status": "NEW — gateway holds LITELLM_MASTER_KEY in /secrets/secrets.env; injects on every LiteLLM-bound request (cq-7). Trust boundary: gateway trusts LiteLLM with the full request body (system prompt + transcript) and tool definitions; LiteLLM holds the real per-backend keys. NetworkPolicy must restrict LiteLLM's inbound to gateway-pod only and its outbound to the configured upstream domains." + }, + { + "name": "LiteLLM → Hosted Qwen provider (NEW)", + "status": "NEW — the chosen hosted-provider sees request bodies (including source code, secrets-like strings in agent transcripts, tool definitions). Per feedback Q5, no compliance constraint rules this out, BUT the cq-9 tool-strip stays on regardless of upstream (WebSearch/WebFetch) which limits the additional exfiltration surface. Plan should document this boundary explicitly so future operators don't accidentally regress cq-9." + }, + { + "name": "Orchestrator → Gateway register_session() (NEW fields)", + "status": "EXTENDED — same authenticated channel as today, but now carries (upstream, model_alias) per agent. No new credential, just new fields on an existing endpoint." + } + ] + }, + + "rollback_plan": { + "description": "Rollback paths if the LiteLLM path goes wrong at any time (validation, post-deploy, or in response to a future LiteLLM CVE).", + "paths": [ + { + "scenario": "Per-agent flip exposes a Claude-Code-on-Qwen tool_use drop (R2) or compaction wedge (R3)", + "action": "Unset PipelineConfig.agent_models[role] (or remove the role's entry from the per-pipeline config) → the spawn path falls back to 'opus' default → the gateway routes to Anthropic singleton → the agent runs Claude as before. No code rollback needed; config-only.", + "guarantee": "Structural — requires no_op_by_default to be true (R13). Reviewer must verify." + }, + { + "scenario": "Future LiteLLM CVE / supply-chain incident (R1)", + "action": "(a) Scale LiteLLM Deployment to 0 replicas (or delete the Deployment) — every LiteLLM-bound request fails closed per cq-8 → fail-fast surfaces the misconfig. (b) Cluster-wide config: unset PipelineConfig.agent_models[*] across all pipelines → all agents revert to Claude path. (c) Roll the LITELLM_MASTER_KEY in secrets.env (rotate to defang the compromised key).", + "guarantee": "Operational — requires monitoring to detect, then a few minutes of config / kubectl actions." + }, + { + "scenario": "LiteLLM pinned image (>= 1.83.7) ships a regression", + "action": "Re-pin to a known-good earlier image (>= 1.83.7); the deployment YAML records the previous-known-good sha256.", + "guarantee": "Operational — assumes the deployment YAML history is in git." + }, + { + "scenario": "Hosted Qwen provider outage / pricing change / data-policy change", + "action": "Same as 'per-agent flip exposes a bug' — unset agent_models, revert to Claude. The new provider key in secrets.env can be revoked separately.", + "guarantee": "Structural + operational." + } + ] + }, + + "areas_requiring_human_review": [ + "R1 (LiteLLM supply-chain controls — pinning + signature + NetworkPolicy + ServiceAccount-isolation must be explicit, reviewer-verifiable acceptance criteria)", + "R2 (LiteLLM SSE tool_use drop bug — operator must know this is the worst-case failure mode for the chosen architecture (cq-5 + Claude Code + LiteLLM) before approving the plan)", + "R13 (deferred empirical validation — the plan reviewer must confirm the no-op-by-default structural guarantee is real, not nominal; the operator owns the validation residual)" + ], + + "open_questions_for_plan_phase": [ + "Slice decomposition (cq-10 deferred to plan phase): the likely two-slice DAG is [gateway upstream router + LiteLLM Deployment, no-op by default] → [per-agent model config + consensus_wrapper plumbing]. The risk-analyst notes this matches the rollback structure: slice 1 is harmless if slice 2 never lands.", + "UpstreamRegistry interface shape (feedback Q3) — flagged as R14. The risk-analyst recommends an explicit small registry class rather than inline if/else.", + "ANTHROPIC_BASE_URL inclusion in _PROTECTED_ENV_KEYS (R11) — cheap defense-in-depth, plan should include or explicitly decline.", + "Claude Code version pin (R8) — the recognised-alias mitigation is heuristic-dependent; plan should specify the Claude Code minor that the validation will run against and bound updates.", + "Defensive logging in _SSEAccumulator on unknown event_type (R10) — cheap, plan should include.", + "max_turns per model (R12) — out of scope or in scope? Recommend follow-up if out of scope." + ], + + "complexity_assessment": "HIGH — the change is no-op by default and the rollback is structurally clean (R13 + rollback_plan), but the risk surface (LiteLLM supply chain R1 + LiteLLM streaming bugs R2 + Claude Code compaction R3 + Claude Code heuristic drift R8) is dominated by external dependencies that egg cannot patch. The plan ships safely; the integration's correctness must be re-validated by the operator at every flip and at every Claude Code / LiteLLM version bump.", + + "external_research_sources": [ + "https://docs.litellm.ai/blog/security-update-march-2026 — March 2026 LiteLLM supply-chain incident", + "https://docs.litellm.ai/blog/security-hardening-april-2026 — Subsequent CVE fix release (1.83.7+)", + "https://www.trendmicro.com/en_us/research/26/c/inside-litellm-supply-chain-compromise.html — TeamPCP attack chain analysis", + "https://snyk.io/blog/poisoned-security-scanner-backdooring-litellm/ — Trivy → LiteLLM credential pivot", + "https://osintteam.blog/litellm-ai-gateway-exposed-to-sql-injection-flaw-cve-2026-42208-57657211230c — CVE-2026-42208 SQL injection (CVSS 9.3)", + "https://github.com/BerriAI/litellm/issues/25561 — /v1/messages streaming tool_use drop on vertex_ai/gemini", + "https://github.com/BerriAI/litellm/issues/25321 — Same bug in v1.82.x non-Anthropic models", + "https://github.com/BerriAI/litellm/issues/24765 — /v1/messages → GitHub Copilot content_block_start drop", + "https://github.com/vllm-project/vllm/issues/21565 — Qwen3 streaming + tool_choice + thinking-disabled drops tool calls", + "https://github.com/vllm-project/vllm/issues/17655 — Qwen3 reasoning parser bug in non-think mode", + "https://github.com/vllm-project/vllm/issues/39056 — Qwen3.5 / 3.6 XML tool_calls inside blocks lost", + "https://platform.claude.com/docs/en/build-with-claude/compaction — Claude Code compaction-control SDK reference", + "https://okhlopkov.com/claude-code-compaction-explained/ — Claude Code auto-compaction heuristic write-up", + "https://dev.to/dcruver/running-claude-code-with-local-llms-via-vllm-and-litellm-599b — Recognised-alias mitigation pattern" + ] +} diff --git a/Makefile b/Makefile index eaf4e106ae..07ddd3d295 100644 --- a/Makefile +++ b/Makefile @@ -504,10 +504,20 @@ k3s-secrets: ## Create gateway secrets from ~/.config/egg/ fi @echo "==> Creating gateway-secrets in egg-system namespace..." @echo " (all files under ~/.config/egg/ become keys in the secret)" + @# LiteLLM master key (issue #2769): the in-cluster LiteLLM + @# Deployment expects ``gateway-secrets.litellm-master-key`` so the + @# gateway's injected x-api-key matches LiteLLM's master_key. The + @# value lives in ``secrets.env`` as ``LITELLM_MASTER_KEY=...``; + @# extract it and surface it as a discrete literal key so both + @# sides of the wire share one source of truth. Empty value is the + @# no-op default (the manifest reads the Secret with + @# ``optional: true``). + @LITELLM_KEY="$$(grep -E '^[[:space:]]*LITELLM_MASTER_KEY[[:space:]]*=' "$$HOME/.config/egg/secrets.env" 2>/dev/null | tail -n1 | cut -d= -f2- | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$$//' -e 's/^"//' -e 's/"$$//' -e "s/^'//" -e "s/'$$//")"; \ export KUBECONFIG=$${KUBECONFIG:-/etc/rancher/k3s/k3s.yaml} && \ kubectl apply -f k8s/base/namespaces.yaml && \ kubectl -n egg-system create secret generic gateway-secrets \ --from-file=$$HOME/.config/egg/ \ + --from-literal=litellm-master-key="$$LITELLM_KEY" \ --dry-run=client -o yaml | kubectl apply -f - deploy: k3s-secrets ## Deploy egg to k3s diff --git a/config/repo_config.py b/config/repo_config.py index 13e1c68e74..986f5afc31 100644 --- a/config/repo_config.py +++ b/config/repo_config.py @@ -307,6 +307,56 @@ def should_disable_auto_fix(repo: str) -> bool: return cast(bool, get_repo_setting(repo, "disable_auto_fix", False)) +def get_default_agent_model(repo: str) -> str | None: + """Return the repository-level default agent model, or ``None`` when unset. + + This is the second tier of the per-agent model resolution precedence + (see ``orchestrator/agent_model_resolution.py``): + + 1. ``PipelineConfig.agent_models[role]`` (per-pipeline override) + 2. ``repositories.yaml`` ``default_agent_model`` (this helper) + 3. Built-in ``"opus"`` default + + The value follows the same classifier as ``agent_models``: a recognised + Claude alias (``opus``, ``opus[1m]``, ``sonnet``, ``sonnet[1m]``, + ``haiku``, ``claude-*``) routes through the Anthropic upstream, anything + else routes through the in-cluster LiteLLM proxy with the alias + ``"opus"`` presented to Claude Code (cq-5 mitigation). + + Args: + repo: Repository in "owner/repo" format + + Returns: + The configured model string, or ``None`` when the repo has no + per-repo entry, the entry omits ``default_agent_model``, or the + ``repositories.yaml`` file is absent (a missing config file is the + same observable as a missing entry — preserves the + no-op-by-default invariant for callers like ``resolve_agent_model`` + that run inside spawn paths where the config file may not be + present, e.g. unit tests and ephemeral CI environments). + + Raises: + ValueError: When ``default_agent_model`` is set to a non-string + YAML value (e.g. ``default_agent_model: 4``). Surfacing the + misconfiguration loudly here keeps it out of ``classify_model``, + where a non-string would otherwise raise an opaque ``TypeError`` + from the regex internals. + """ + try: + value = get_repo_setting(repo, "default_agent_model", None) + except FileNotFoundError: + return None + if value is None: + return None + if not isinstance(value, str): + raise ValueError( + f"default_agent_model for {repo!r} must be a string, got " + f"{type(value).__name__}: {value!r}. Set it to a recognised " + f"Claude alias (opus, sonnet, …) or a LiteLLM model name." + ) + return value + + try: from egg_config.validators import validate_checks except ImportError: diff --git a/config/repositories.yaml.example b/config/repositories.yaml.example index 3567620b38..b606c5d8a4 100644 --- a/config/repositories.yaml.example +++ b/config/repositories.yaml.example @@ -109,6 +109,17 @@ readable_repos: # **/*_test.go in tests_globs). Security-relevant blocklists # (.egg-state/contracts/, .github/) are hard-coded and cannot be # relaxed. See docs/guides/sdlc-pipeline.md#per-repository-role-patterns. +# - default_agent_model: Repository-level default for the per-agent +# model knob added in #2769. Used by every agent role unless the +# pipeline submission overrides it via ``agent_models``. A recognised +# Claude alias (opus, opus[1m], sonnet, sonnet[1m], haiku, claude-*) +# routes through the Anthropic upstream; anything else routes through the +# in-cluster LiteLLM proxy with the recognised alias "opus" presented +# to Claude Code (cq-5 mitigation). Precedence: +# PipelineConfig.agent_models[role] +# > this default_agent_model +# > built-in "opus" default +# Default: not set (every role runs on built-in "opus"). repo_settings: # Example: # YOUR_USERNAME/egg: @@ -190,6 +201,13 @@ repo_settings: # tests_globs: ["**/__tests__/**", "**/*.test.ts", "**/*.test.tsx", "**/*.spec.ts"] # code_globs: ["**/*.ts", "**/*.tsx", "**/*.js"] # docs_globs: ["**/*.md", "docs/"] + # + # # Per-agent model example (#2769): route every role on this repo through + # # the LiteLLM proxy by default, picking up the hosted-Qwen model entry + # # populated in the LiteLLM ConfigMap. The pipeline-level ``agent_models`` + # # field still wins when set. + # YOUR_USERNAME/qwen-pilot: + # default_agent_model: qwen3-coder-30b # User mode configuration (optional) # When auth_mode is set to "user" for a repo, operations will be diff --git a/config/secrets.template.env b/config/secrets.template.env index f53bf03623..2298830918 100644 --- a/config/secrets.template.env +++ b/config/secrets.template.env @@ -24,6 +24,24 @@ ANTHROPIC_API_KEY="" ANTHROPIC_OAUTH_TOKEN="" +# ============================================================================= +# LiteLLM Proxy (Optional — non-Claude model routing, issue #2769) +# ============================================================================= +# Master key the gateway injects as ``x-api-key`` on every request routed +# through the in-cluster LiteLLM proxy. The same value is consumed by the +# LiteLLM Deployment via the ``gateway-secrets`` Secret (key +# ``litellm-master-key``) so both sides of the wire agree. +# +# Per-backend provider credentials (e.g. TOGETHER_API_KEY for a hosted +# Qwen provider — see cq-6) are NOT stored here; they go in LiteLLM's +# own env-var slots so the gateway never sees the raw provider key. +# +# Leave empty to disable LiteLLM routing — no agent will be routed to +# LiteLLM with this unset, regardless of any per-pipeline +# ``agent_models`` override. The Claude path is byte-identical. + +LITELLM_MASTER_KEY="" + # ============================================================================= # Slack Integration # ============================================================================= diff --git a/docs/architecture/README.md b/docs/architecture/README.md index cc82808dc2..6f179ca11b 100644 --- a/docs/architecture/README.md +++ b/docs/architecture/README.md @@ -187,6 +187,7 @@ The SDLC pipeline orchestrates agent-based development with structurally enforce - [Gateway Auto-Filter](gateway-auto-filter.md) - Restricted-path rejection on push and the commit-authorship registry that backs attribution - [Credential Injection](credential-injection.md) - Zero-credential sandbox with API key proxy - [Network Isolation](network-isolation.md) - Public/private network modes +- [Upstream Routing](upstream-routing.md) - `UpstreamRegistry` seam, LiteLLM topology, per-session routing decision, and no-op-by-default invariant for non-Claude agent backends ([#2769](https://github.com/jwbron/egg/issues/2769)) - [SDLC Pipeline](sdlc-pipeline.md) - Structurally enforced agent checkpoints - [Declarative Setup](declarative-setup.md) - Python-based setup - [Logging](logging.md) - Structured JSON logging diff --git a/docs/architecture/orchestrator.md b/docs/architecture/orchestrator.md index 02911b5f04..3fa05940b8 100644 --- a/docs/architecture/orchestrator.md +++ b/docs/architecture/orchestrator.md @@ -244,6 +244,25 @@ See `plugins/refine-plan/skills/refine-plan/agents/applier.md`'s "Out of scope: - Orchestrator-side drain hook: `orchestrator/routes/pipelines.py::_drain_wontdo_batch_after_apply` — invoked from both the auto-advance and HITL-resolution apply-phase exit paths; writes per-Task `jira_action_status` back via the `on_entry_result` callback. - Issue-level decision record: [#1557 decision-15](https://github.com/jwbron/egg/issues/1557) (trust-boundary for Jira transitions). +## Upstream Routing (Per-Agent Model Backends, [#2769](https://github.com/jwbron/egg/issues/2769)) + +Per-agent `/v1/messages` traffic routes through an `UpstreamRegistry` +in the gateway that resolves the upstream (Anthropic vs LiteLLM) from +per-session metadata declared by the orchestrator at session-create +time — the same IP-keyed session lookup that already drives +`session_mode`. Slice 1 of #2769 lands the router and the LiteLLM +Deployment + Service in `egg-system`, no-op by default; slice 2 adds +`PipelineConfig.agent_models` and a repository-level +`default_agent_model` for the orchestrator-side resolution. Until an +operator opts in, every existing pipeline keeps running on Claude +with byte-identical gateway behavior. + +See [Upstream Routing](upstream-routing.md) for the gateway-side +seam (registry, credential layout, request lifecycle, failure +policy, and the no-op-by-default invariant). The operator-facing +setup ships in slice 2 as `docs/guides/per-agent-models.md` (that +file does not exist until slice 2 lands). + ## Network Mode Pipelines can specify an explicit network mode that controls internet access for spawned containers: diff --git a/docs/architecture/upstream-routing.md b/docs/architecture/upstream-routing.md new file mode 100644 index 0000000000..f7f8ac887b --- /dev/null +++ b/docs/architecture/upstream-routing.md @@ -0,0 +1,431 @@ +# Upstream Routing — LiteLLM Proxy Seam + +This document describes the gateway-side seam that lets per-agent +`/v1/messages` traffic route to either `api.anthropic.com` (default) +or a LiteLLM proxy that fronts non-Claude backends (the first target +is a hosted Qwen provider). It covers the `UpstreamRegistry` +abstraction, the per-session routing decision, the credential layout, +the LiteLLM topology in Kubernetes, and the **no-op-by-default** +invariant that keeps every Claude-bound agent on byte-identical paths +when LiteLLM is not configured. + +Status: this seam lands in two stacked changes for [#2769](https://github.com/jwbron/egg/issues/2769). Slice 1 — described here — is the +**gateway router + LiteLLM Deployment**, no-op by default. Slice 2 +adds the orchestrator-side resolution (`PipelineConfig.agent_models`, +`default_agent_model`, `resolve_agent_model`) and the gateway-side +body rewrite (`_rewrite_upstream_model`) that together connect an +agent role to a concrete upstream model. Operators looking to actually +flip a role to a non-Claude backend should start at the +[Per-Agent Models guide](../guides/per-agent-models.md) — it walks +through the two configuration knobs, the precedence chain, the cq-5 +recognized-alias mitigation, and the cq-4 operator smoke test +end-to-end. + +## Why a router, not a hard-wired second client + +Today's gateway hard-wires the Anthropic upstream: a singleton +`httpx.Client` is opened against `https://api.anthropic.com` +(`gateway/gateway.py:9320` `get_anthropic_client`) and every +`/v1/messages` and `/v1/messages/count_tokens` request injects an +Anthropic credential before forwarding +(`gateway/gateway.py:9355` `_inject_anthropic_credentials`). + +We want any agent role to be independently switchable to a non-Claude +model — first for cost (the primary driver behind #2769), eventually +to mix self-hosted weights with Claude on per-role boundaries — while +every Claude-bound agent stays byte-identically on the existing path +until the operator opts in. We also want a clean swap-out point in +case LiteLLM is unsuitable in the future (refine-phase feedback Q3: +the March 2026 PyPI incident is recent enough that hard-wiring a +specific proxy framework is a known risk). + +The shape that satisfies both constraints is a small per-request +registry of `(httpx.Client, credential_resolver)` pairs keyed by +upstream name, with the proxy routes resolving the upstream from the +per-session metadata that already drives `session_mode`. The router +adds no new tool-side surface and no new agent-visible API — agents +keep talking to `http://egg-gateway:9848/v1/messages` as before. The +gateway is the only component that knows there is more than one +upstream. + +## Topology — where LiteLLM runs in the cluster + +The HITL on `cq-1` selected **a separate Deployment + Service in the +`egg-system` namespace** (1 LiteLLM pod, gateway calls it over the +cluster-internal Service DNS). The alternative shapes — a sidecar in +the gateway pod, or a fully separate namespace with its own +NetworkPolicy — were declined: sidecar couples lifecycle to the +gateway (one restart kills both), and a separate namespace doubles +the ops surface for a defense-in-depth gain we can add later. + +``` +┌───────────────────────────────────────────────────────────────────┐ +│ UPSTREAM ROUTING │ +├───────────────────────────────────────────────────────────────────┤ +│ │ +│ egg-agents/ (untrusted) egg-system/ (trusted) │ +│ ┌─────────────────┐ ┌─────────────────────┐ │ +│ │ sandbox pod │ │ egg-gateway │ │ +│ │ Claude Code │ ANTHROPIC_ │ │ │ +│ │ no creds │ BASE_URL │ UpstreamRegistry │ │ +│ │ │ ──────────▶ │ ├── anthropic ───┼──▶ api.anthropic.com +│ │ /v1/messages │ :9848 │ └── litellm ───┼──▶ litellm.egg-system │ +│ └─────────────────┘ │ │ │ .svc.cluster.local +│ │ Session lookup │ │ :4000 +│ │ by remote_addr → │ │ │ +│ │ session.upstream │ │ ▼ +│ └─────────────────────┘ │ ┌────────────────┐ +│ │ │ litellm pod │ +│ │ │ (egg-system) │ +│ │ │ Anthropic-shape│ +│ │ │ translator → │ +│ │ │ hosted Qwen │ +│ │ └────────────────┘ +└───────────────────────────────────────────────────────────────────┘ +``` + +Key topology properties: + +- **Sandbox sees one URL.** Agents still use + `ANTHROPIC_BASE_URL=http://egg-gateway:9848`. The router is + invisible to the sandbox; the gateway picks the upstream based on + the per-agent session. +- **LiteLLM is gateway-only.** No NetworkPolicy change on the + `egg-agents` namespace — agent pods continue to be denied + egress to `litellm.egg-system.svc.cluster.local`. Only the + gateway pod calls LiteLLM. +- **Operator-owned isolation for the LiteLLM pod itself.** LiteLLM + in turn talks out to the configured provider (hosted Qwen for + the first cut, per `cq-6`); per-provider credentials live in the + LiteLLM ConfigMap, not the gateway. +- **No-op until configured.** The Service comes up healthy with an + empty `model_list`. Until an operator populates the ConfigMap + and sets the LiteLLM master key, every `/v1/messages` request + still routes to `api.anthropic.com` via the default + `session.upstream = "anthropic"`. + +The manifests land at: + +| File | Resource | +|------|----------| +| `k8s/base/litellm-deployment.yaml` | Deployment running a pinned LiteLLM image in `egg-system`, mounting the ConfigMap at `/app/config.yaml` | +| `k8s/base/litellm-service.yaml` | ClusterIP Service named `litellm` exposing port `4000` (LiteLLM's default) | +| `k8s/base/litellm-configmap.yaml` | Empty `model_list` so the pod is healthy but serves nothing until operators populate it | +| `k8s/base/kustomization.yaml` | Registers the three new resources | + +## The router — `UpstreamRegistry` + +The router is a new module, `gateway/upstream_registry.py`, that owns +the (singleton `httpx.Client`, credential resolver) pair for each +known upstream: + +| Upstream key | Client base URL | Credential resolver | +|--------------|-----------------|---------------------| +| `"anthropic"` | `https://api.anthropic.com` (preserves the `# noqa: EGG200` annotation at `gateway/gateway.py:9325`) | `AnthropicCredentialsManager` — unchanged | +| `"litellm"` | `LITELLM_BASE_URL` env var, default `http://litellm.egg-system.svc.cluster.local:4000` | LiteLLM resolver in `gateway/anthropic_credentials.py` (see below) | + +Public surface: + +- `class UpstreamRegistry` — keyed lookup of upstream clients + + credential resolvers +- `get(upstream: str) -> (httpx.Client, credential_resolver)` — + returns the pair for a known upstream +- `class UnknownUpstreamError` — raised by `get()` on miss; the + proxy routes translate this to a 400 with a descriptive body +- `get_upstream_registry()` — module-level accessor mirroring the + `get_anthropic_client()` lifetime semantics so the registry is + built once and shared across requests + +Both clients share the same `httpx.Timeout(120.0, connect=10.0)` +and `httpx.Limits(max_connections=100, max_keepalive_connections=20)` +shape as today's `_anthropic_client` so latency / pooling behavior +does not silently regress for the Claude path. + +## Credentials + +### Layout + +Anthropic credentials are unchanged: the +`AnthropicCredentialsManager` loads `ANTHROPIC_API_KEY` / +`ANTHROPIC_OAUTH_TOKEN` from `~/.config/egg/secrets.env` with an +mtime-invalidated cache (`gateway/anthropic_credentials.py`). + +LiteLLM gets a parallel resolver that loads `LITELLM_MASTER_KEY` +from the same `secrets.env` file via the existing `parse_env_file` +helper at `gateway/anthropic_credentials.py:52` and caches with the +same mtime invalidation. When the key is set, the resolver returns +a credential shaped `header_name="x-api-key"`, +`header_value=`. When the key is **unset** — +the default for every existing deployment — the resolver returns +`None` and emits no startup warning. This matches the existing +Anthropic-resolver behavior when the API key is absent (the +"no-credential" path is well-trodden) and is the source of the +**no-op-by-default** invariant: without `LITELLM_MASTER_KEY`, +nothing about the Claude flow changes. + +The key lands as a documented entry in `config/secrets.template.env` +with an explicit "leave empty to disable LiteLLM routing — no agent +will be routed to LiteLLM with this unset" comment, one block below +the existing `ANTHROPIC_API_KEY` block. This is the choice from +`cq-7`: the gateway holds the LiteLLM master key and injects it on +every LiteLLM-bound request, mirroring today's Anthropic credential +injection. LiteLLM itself holds the real per-backend keys in its +ConfigMap (e.g. the hosted Qwen provider's API key). The sandbox +sees neither. + +### Injection + +`_inject_anthropic_credentials` (`gateway/gateway.py:9355`) is +generalized to `_inject_upstream_credentials(headers, upstream)`. +The old symbol stays as a back-compat alias that calls through with +`upstream="anthropic"`. Dispatch picks the resolver from the +registry and produces a `(headers, error_response_tuple)` pair +identical in shape to today's: a missing credential returns the same +401 / `authentication_error` JSON body the Anthropic path returns, +just sourced from the LiteLLM resolver when the upstream is LiteLLM. + +The OAuth-passthrough fallback (when no gateway credential is +configured but the client sent `Authorization` / `x-api-key` +headers itself, used by Claude Code's own OAuth flow) is preserved +verbatim for the Anthropic path. The LiteLLM path does not exercise +that fallback because the sandbox never holds LiteLLM credentials. + +## Per-session routing + +The routing decision is **per session**, not per request body. The +HITL on `cq-2` settled this: the orchestrator declares the upstream +when it spawns the agent (the same IP-keyed session lookup that +already drives `session_mode`), and the gateway treats the model +name in the request body as informational only. The alternatives — +a custom HTTP header from the sandbox, or model-name sniffing — +were declined: the header shape duplicates session state, and +model-name sniffing conflicts with the `cq-5` compaction mitigation +that intentionally presents `opus` to Claude Code even when the +upstream is non-Claude. + +### Session fields + +Two new optional fields land on the `Session` dataclass at +`gateway/session_manager.py:288`: + +- `upstream: str = "anthropic"` — the registered upstream name to + route this session's `/v1/messages` traffic to +- `upstream_model: str | None = None` — the on-the-wire model name + to forward upstream (only meaningful for LiteLLM; see slice 2's + body rewrite) + +Both default to the Claude path. Sessions persisted **before** this +change rehydrate cleanly: `Session.from_persistence` tolerates +dicts that omit the two keys, falling back to the defaults. This is +exercised by the back-compat round-trip test in +`gateway/tests/test_session_manager.py`. + +### Registration plumbing + +`SessionManager.register_session` (`gateway/session_manager.py:548`) +gains two optional parameters (`upstream`, `upstream_model`) and +stores them on the returned `Session`. The +`/api/v1/sessions/create` handler at `gateway/gateway.py:8507` +parses both from the request body (defaulting to `"anthropic"` / +`None`), validates `upstream` against the names registered in +`UpstreamRegistry` (unknown name → HTTP 400 with a descriptive +error), and passes both through to `register_session`. The existing +`audit_log("session_created", ...)` call includes the upstream and +upstream_model so the per-session routing decision is auditable. + +The orchestrator-side wire shape catches up in +`GatewayClient.register_session` (`orchestrator/gateway_client.py:602`): +two optional kwargs, included in the POST body only when set +(matches the existing optional-field pattern in the surrounding +calls). No slice-1 caller passes them — that's purely the wire +contract. Slice 2 is where the orchestrator's spawner actually +decides per-agent upstreams and calls +`register_session(upstream=…, upstream_model=…)`. + +### Proxy routes + +`proxy_anthropic_messages` (`gateway/gateway.py:9753`) and +`proxy_count_tokens` (`gateway/gateway.py:10020`) each replace +`client = get_anthropic_client()` with a registry lookup keyed by +`session.upstream`. When no session is found (e.g. an +unauthenticated probe), the upstream defaults to `"anthropic"` — +this preserves today's behavior verbatim. The `headers, error = +_inject_anthropic_credentials(headers)` call becomes +`_inject_upstream_credentials(headers, session.upstream)`. + +Everything **inside** the request loop stays unchanged: + +- `_filter_blocked_tools(...)` (the private-mode WebSearch / + WebFetch strip from `cq-9` — kept identical regardless of + upstream; the conservative read is that an attacker on a + compromised sandbox can still exfiltrate through search queries + even when the upstream is self-hosted, so the strip stays + defense-in-depth on all routes) +- `_SSEAccumulator` parse and the per-session transcript capture +- The pre-stream retry + mid-stream synthetic-error resilience + loop from [#1907](https://github.com/jwbron/egg/issues/1907) (see [Credential Injection + → Upstream Stream Resilience](credential-injection.md#upstream-stream-resilience)) +- Connection-pool / timeout / limits — both clients share the same + shape (see "The router" above) + +The router is intentionally a **per-request lookup at the top of +the handler**, not new branches inside any of these inner loops. +That keeps the security-critical inner code byte-identical and the +review surface small. + +## Request lifecycle — both upstreams + +### Anthropic (default — every agent, until configured otherwise) + +``` +sandbox ──POST /v1/messages──▶ gateway + │ + ├─ session_manager.get_session_by_ip(...) + │ → session.upstream = "anthropic" (default) + │ → session.upstream_model = None (default) + │ + ├─ _inject_upstream_credentials(headers, "anthropic") + │ → registry.get("anthropic")[1].get_credential() + │ → AnthropicCredentialsManager (Anthropic key from secrets.env) + │ + ├─ _filter_blocked_tools(body, session.mode) + │ + ├─ client = registry.get("anthropic")[0] + │ → existing _anthropic_client, base_url=https://api.anthropic.com + │ + └─ stream upstream → accumulate → return SSE downstream +``` + +Wire shape: byte-identical to today. This is the regression guard: +with `agent_models == {}` everywhere (the slice-2 default) and no +`LITELLM_MASTER_KEY` configured, every gateway request follows the +exact same path it does today. + +### LiteLLM (per-agent, once configured) + +``` +sandbox ──POST /v1/messages──▶ gateway + │ + ├─ session_manager.get_session_by_ip(...) + │ → session.upstream = "litellm" + │ → session.upstream_model = "qwen3-coder-30b" (set by orchestrator at session create) + │ + ├─ _inject_upstream_credentials(headers, "litellm") + │ → registry.get("litellm")[1].get_credential() + │ → LiteLLM resolver (LITELLM_MASTER_KEY from secrets.env) + │ + ├─ _filter_blocked_tools(body, session.mode) + │ (same strip regardless of upstream — cq-9 conservative choice) + │ + ├─ [slice 2] _rewrite_upstream_model(body, session.upstream_model) + │ → body["model"] = "qwen3-coder-30b" + │ (Claude-Code-facing alias stays "opus" so compaction math + │ remains sane — cq-5 mitigation) + │ + ├─ client = registry.get("litellm")[0] + │ → LiteLLM client, base_url from LITELLM_BASE_URL + │ + └─ stream upstream → accumulate → return SSE downstream +``` + +The body rewrite is **slice 2**'s addition; in slice 1 the router +is in place but no caller sets `session.upstream = "litellm"`, so +the LiteLLM branch is exercised only by unit tests until slice 2 +lands the spawn-side plumbing. + +## Failure policy + +The HITL on `cq-8` settled this: when the LiteLLM proxy is +unreachable or errors for a non-Claude agent, the gateway **fails +closed** — a 502 surfaces to the agent, no fallback to Claude. The +alternatives (transparent Claude fallback, HITL escalation) were +declined: fallback produces a quietly-mixed transcript that erodes +the cost goal motivating the work, and HITL escalation is a +follow-up that can be layered on top of the fail-closed default if +operators demand it. + +This is the same policy the Claude path uses today on upstream +errors. The existing `except httpx.ConnectError / TimeoutException / +Exception` handlers in `proxy_anthropic_messages` and +`proxy_count_tokens` cover the LiteLLM case verbatim — both +upstreams produce the same 502 / 504 error contracts. + +## No-op-by-default invariant + +The invariant that makes this safe to ship before a live non-Claude +endpoint exists has three independent guards: + +1. **No LiteLLM master key configured.** The LiteLLM resolver + returns `None`, so any request that *did* somehow route to + LiteLLM would fail credential injection with the standard 401 + — but no request routes there because of guard 2. +2. **Session upstream defaults to `"anthropic"`.** Both + `Session.upstream` and the `/api/v1/sessions/create` handler's + default value are `"anthropic"`. Slice 1 has no caller that + passes `upstream="litellm"`. Slice 2 only sets it when + `PipelineConfig.agent_models` or repository-level + `default_agent_model` names a non-Claude model. +3. **`agent_models` default is empty.** Slice 2's + `PipelineConfig.agent_models` field defaults to `{}`, and the + repository-level `default_agent_model` defaults to `None`. + Without operator action, every spawn resolves to the built-in + `"opus"` Claude path with `upstream="anthropic"`. + +Any one of the three suffices to keep the LiteLLM client cold on a +given deployment. All three are independent: a misconfiguration on +one does not silently activate the LiteLLM path. + +## HITL decisions that shape this design + +The `#2769` refine phase resolved eleven `cq-*` decisions. The six +that directly shape the slice-1 architecture: + +| Decision | Resolution | Effect on this seam | +|----------|------------|---------------------| +| `cq-1` | Separate Deployment + Service in `egg-system` | Topology section above — LiteLLM is a sibling of the gateway, not a sidecar | +| `cq-2` | Per-agent session metadata (IP-keyed lookup, like `session_mode`) | Per-session routing decision; model name in request body is informational only | +| `cq-5` | Keep Claude Code harness for non-Claude models; present recognized alias | Compaction-math mitigation; body rewrite lives in slice 2 | +| `cq-7` | Gateway holds `LITELLM_MASTER_KEY` in `secrets.env`; injects per-request | Credentials section above — mirrors today's Anthropic injection | +| `cq-8` | Fail closed on LiteLLM errors (502, no fallback) | Failure policy section above | +| `cq-9` | Keep the private-mode WebSearch / WebFetch strip on every upstream | Proxy-routes section — `_filter_blocked_tools` runs identically regardless of upstream | + +The full set is at [`.egg-state/contracts/issue-2769.json`](../../.egg-state/contracts/issue-2769.json). + +## Files + +| File | Role in this seam | +|------|-------------------| +| `gateway/upstream_registry.py` *(new)* | `UpstreamRegistry`, `UnknownUpstreamError`, `get_upstream_registry()` — keyed (client, credential_resolver) lookup | +| `gateway/anthropic_credentials.py` | LiteLLM credential resolver alongside `AnthropicCredentialsManager`; shared `parse_env_file` (`anthropic_credentials.py:52`) and mtime cache | +| `gateway/gateway.py:9320` `get_anthropic_client` | Unchanged; registry holds this as the `"anthropic"` entry | +| `gateway/gateway.py:9355` `_inject_anthropic_credentials` | Generalized to `_inject_upstream_credentials(headers, upstream)`; old symbol kept as back-compat alias | +| `gateway/gateway.py:9753` `proxy_anthropic_messages` | Resolves upstream from `session.upstream`; SSE / tool-filter / retry inner loops unchanged | +| `gateway/gateway.py:10020` `proxy_count_tokens` | Same routing change as `proxy_anthropic_messages` | +| `gateway/gateway.py:8507` `/api/v1/sessions/create` | Accepts `upstream` and `upstream_model`; validates `upstream` against `UpstreamRegistry`; audit log includes both | +| `gateway/session_manager.py:288` `Session` | New `upstream: str = "anthropic"` and `upstream_model: str \| None = None` fields with back-compat `from_persistence` | +| `gateway/session_manager.py:548` `register_session` | New optional `upstream` / `upstream_model` kwargs | +| `orchestrator/gateway_client.py:602` `register_session` | New optional `upstream` / `upstream_model` kwargs; included in POST body only when set | +| `k8s/base/litellm-deployment.yaml` *(new)* | LiteLLM pod in `egg-system` | +| `k8s/base/litellm-service.yaml` *(new)* | ClusterIP `litellm:4000` (default `LITELLM_BASE_URL`) | +| `k8s/base/litellm-configmap.yaml` *(new)* | Empty `model_list` — gateway-only callable; operators populate post-deploy | +| `k8s/base/kustomization.yaml` | Registers the three new LiteLLM resources | +| `config/secrets.template.env` | New `LITELLM_MASTER_KEY=""` block with the disable-when-empty note | + +## Related Documentation + +- [Credential Injection](credential-injection.md) — Anthropic + credential resolver shape and the upstream-reset resilience + pattern that both router branches inherit unchanged +- [Orchestrator Architecture](orchestrator.md) — Spawner and + session-creation context for slice 2's per-agent model + resolution +- [Per-Agent Models Guide](../guides/per-agent-models.md) — Operator + walkthrough for the slice-2 configuration plumbing + (`PipelineConfig.agent_models`, repo-level `default_agent_model`, + the `resolve_agent_model` precedence + classifier, the gateway-side + `_rewrite_upstream_model` helper, the cq-5 recognized-alias + mitigation, and the cq-4 hosted-Qwen smoke test) +- [Network Isolation](network-isolation.md) — Cluster network + posture; LiteLLM is gateway-only and NetworkPolicy is unchanged +- Issue [#2769](https://github.com/jwbron/egg/issues/2769) — Original + motivation, refine-phase decisions, plan-phase risk analysis diff --git a/docs/guides/per-agent-models.md b/docs/guides/per-agent-models.md new file mode 100644 index 0000000000..81e80ef7ae --- /dev/null +++ b/docs/guides/per-agent-models.md @@ -0,0 +1,469 @@ +# Per-Agent Models — Running a Single Agent on a Non-Claude Backend + +This guide walks an operator through the **slice-2** plumbing of +[#2769](https://github.com/jwbron/egg/issues/2769): flipping any single +SDLC phase agent role (refiner, coder, tester, a reviewer, …) to a +non-Claude model via the gateway's LiteLLM proxy, without touching +agent prompts, the sandbox, or source code. By the end you will have +configured a pipeline that runs (e.g.) the refiner on hosted Qwen while +every other role continues to run on Claude. + +The seam itself — the gateway-side `UpstreamRegistry`, the LiteLLM +Deployment, the per-session routing decision — is described in +[Upstream Routing](../architecture/upstream-routing.md) and was the +slice-1 deliverable. This guide is the operator-facing complement: the +two configuration fields slice 2 introduces, how they compose, the +**recognized-alias** mitigation that keeps Claude Code's auto-compaction +math sane on non-Claude routes, and the end-to-end smoke test that +validates a live LiteLLM endpoint. + +## Mental model in one sentence + +> Each agent role independently resolves to a `(claude_code_alias, +> upstream, upstream_model)` triple — the `--model` flag handed to +> Claude Code stays a recognized Claude alias on every route, while +> the gateway rewrites the request body's `model` field to the +> upstream-side name (`"qwen3-coder-30b"`, …) on LiteLLM-bound +> requests so the proxy targets the right backend. + +The default path (every agent on Claude, every request to +`api.anthropic.com`) is **byte-identical to today** — no config means +no behavioral change. + +## The two configuration knobs + +### Per-pipeline override — `PipelineConfig.agent_models` + +`agent_models: dict[str, str]` on +[`PipelineConfig`](../../orchestrator/models.py) (the field lives next +to the existing overseer-tier model fields around +`orchestrator/models.py:757`). Keys are +[`AgentRole`](../../shared/egg_contracts/agent_roles.py) values, but +**only the SDLC phase producer and reviewer roles** (`"coder"`, +`"refiner"`, `"tester"`, the `"reviewer_*"` roles, …) — the roles +spawned through the paths that consult the resolver. A Pydantic +validator rejects both typos *and* otherwise-valid roles the resolver +never honors — the utility roles `autofixer` / `conflict_resolver` and +the interface roles `overseer` / `inspector` spawn through dedicated +paths that bypass `resolve_agent_model`, so an override naming one of +them would silently no-op at spawn. The honored set is the constant +`agent_roles.MODEL_OVERRIDE_ROLES`; the validator surfaces a +misconfigured key immediately at construction time. Values are +free-form model strings — interpreted by the resolver (below). + +```python +PipelineConfig( + agent_models={ + "refiner": "qwen3-coder-30b", # → LiteLLM, "qwen3-coder-30b" + "coder": "sonnet", # → Anthropic, "sonnet" + }, + # other fields unchanged … +) +``` + +Default-constructed `PipelineConfig.agent_models` is `{}` — every +existing pipeline continues to spawn every role on the built-in +`"opus"` Claude default with `upstream="anthropic"`. + +### Repository-level default — `default_agent_model` + +`default_agent_model: str | None` on `repositories.yaml`, read by +[`get_default_agent_model(repo)`](../../config/repo_config.py) in +`config/repo_config.py` (mirroring the existing +`get_repo_setting(repo, key, default)` pattern at +`config/repo_config.py:248`). Documented in +[`config/repositories.yaml.example`](../../config/repositories.yaml.example) +with an inline precedence note. + +```yaml +# ~/.config/egg/repositories.yaml +repo_settings: + acme-corp/widgets: + # Every role defaults to Sonnet on this repo, unless a pipeline + # passes its own agent_models entry. + default_agent_model: sonnet +``` + +When set, the field applies to **every role** that the per-pipeline +`agent_models` doesn't already pin. + +### Precedence + +The resolver +[`resolve_agent_model(role, pipeline_config, repo)`](../../orchestrator/agent_model_resolution.py) +in `orchestrator/agent_model_resolution.py` walks the chain: + +1. `pipeline_config.agent_models.get(role.value)` — per-pipeline, + per-role override +2. `get_default_agent_model(repo)` — repository-level default +3. Built-in `"opus"` — the historical default; preserves today's + behavior unchanged + +The result is an `AgentModelDecision` dataclass with fields +`(claude_code_alias: str, upstream: str, upstream_model: str | None)`. + +### Classifier — Claude alias vs. LiteLLM upstream + +The resolver's classifier divides model strings into two camps: + +| Model string pattern | `upstream` | `claude_code_alias` | `upstream_model` | +|----------------------|------------|---------------------|------------------| +| `opus`, `opus[1m]`, `sonnet`, `sonnet[1m]`, `haiku` | `"anthropic"` | the model string verbatim | `None` | +| `claude-*` (e.g. `claude-3-5-sonnet-20241022`) | `"anthropic"` | the model string verbatim | `None` | +| Everything else (e.g. `qwen3-coder-30b`, `mistral-large`) | `"litellm"` | **`"opus"`** (the cq-5 mitigation — see below) | the model string verbatim | + +Two consequences: + +- The Anthropic path produces no wire change at all when the resolved + model is a Claude alias — `upstream_model` is `None`, so the gateway + forwards the body verbatim. +- The LiteLLM path always presents Claude Code the alias `"opus"` — + see *Recognized-alias mitigation* below for why this matters. + +## How the resolved decision threads through spawn + +Two sites resolve and thread the decision into the consensus-wrapper +command + gateway session-create call: + +| Spawn site | Threads `--model` to | Threads `upstream` / `upstream_model` to | +|------------|----------------------|-------------------------------------------| +| Initial spawn at `orchestrator/concurrent_executor.py:454` | `build_consensus_wrapped_command(model=decision.claude_code_alias, …)` | `GatewayClient.register_session(upstream=…, upstream_model=…)` at `orchestrator/kubernetes_spawner.py:735` (the new kwargs land on the slice-1 wire contract) | +| Restart path at `orchestrator/routes/pipelines.py:2704` | Same `--model` resolution | `restart_agent_job` → `spawn_agent_job` → `register_session(upstream=…, upstream_model=…)` — a restart respawns the Job and registers a **new** gateway session carrying the resolved decision | + +When the resolved decision is the default Claude case +(`upstream="anthropic"`, `upstream_model=None`), +`ConcurrentPhaseExecutor._spawn_agent` omits the `upstream` / +`upstream_model` kwargs from its `spawn_fn` call entirely +(`concurrent_executor.py:501-503`), so test mocks and legacy spawn +paths see the pre-#2769 call signature. One layer down, +`spawn_agent_job` still passes both kwargs to +`GatewayClient.register_session` (`kubernetes_spawner.py:770-771`) — as +`None` on the default path — and `register_session` drops `None` +values from the session-create request body +(`gateway_client.py:707-712`), so the wire shape stays byte-identical +to today. This is the slice-2 regression guard exercised by the +existing concurrent-executor tests. + +## The gateway-side body rewrite + +For LiteLLM-bound requests, the gateway has to make the upstream model +name reach LiteLLM — Claude Code is still sending `"model": "opus"` in +the body it constructs. The rewrite lives in a new helper +`_rewrite_upstream_model(request_body, upstream_model)` colocated with +`_filter_blocked_tools` in `gateway/gateway.py` (slice-1 added +`_filter_blocked_tools` at `gateway/gateway.py:9496`; the new helper +sits adjacent to it). + +Called from `proxy_anthropic_messages` +(`gateway/gateway.py:9839`) and `proxy_count_tokens` +(`gateway/gateway.py:10135`) **after** `_filter_blocked_tools` and +**before** building the upstream request: + +``` +sandbox ──POST /v1/messages──▶ gateway + │ + ├─ session_manager.get_session_by_ip(...) + │ → session.upstream = "litellm" + │ → session.upstream_model = "qwen3-coder-30b" + │ + ├─ _inject_upstream_credentials(headers, "litellm") + ├─ _filter_blocked_tools(body, session.mode) + ├─ _rewrite_upstream_model(body, session.upstream_model) + │ → body["model"]: "opus" → "qwen3-coder-30b" + │ + ├─ client = registry.get("litellm")[0] + └─ stream upstream → accumulate → return SSE downstream +``` + +Behavior at the edges: + +- **`upstream="anthropic"` (or `upstream_model is None`)** — the body + is returned unchanged. The Anthropic regression guard is enforced + with a test that sends a non-default incoming `"model"` (e.g. + `"opus"`) and asserts byte-identical forwarding. +- **Invalid JSON in the request body** — the rewrite is a no-op (the + original bytes are returned). Parsing failures never crash the + proxy; they fall through to the existing upstream and let the + upstream produce whatever error response it would have produced. + +## Recognized-alias mitigation (cq-5) — *plausible, not empirically proven* + +> The behavior described in this section is the cq-5 *mitigation* for +> a harness compatibility risk — not a proven guarantee. Confirming +> compaction math actually stays sane on long non-Claude sessions is +> the cq-4-deferred operator smoke test below. Read the invariants +> as "what the resolver enforces structurally" rather than "what +> Claude Code is guaranteed to do." + +The harness Claude Code runs the agent inside makes **decisions keyed +on the model name** that we cannot ask it to override: + +- It derives a model's **context window** — and therefore + **auto-compaction timing** — from the `--model` flag. +- Some features (e.g. extended thinking) are gated on recognized + model names. + +An unrecognized model name gets a fallback context window that does +not match the real backend. On long sessions this produces +over-length requests that hard-fail and wedge the agent. + +The cq-5 mitigation, baked into the resolver's classifier above: + +> For every LiteLLM-routed agent, Claude Code is told `--model opus`. +> The gateway separately rewrites the on-the-wire `"model"` field to +> the LiteLLM-side name (`"qwen3-coder-30b"` etc.). + +Two structural invariants follow (these are enforced by the resolver +and tested explicitly; *whether* they are sufficient to keep Claude +Code's compaction math sane on a given backend remains the smoke +test's job): + +- The Claude-Code-facing alias is **always** a recognized Claude name + (`opus` by default). +- The model name actually requested upstream is **always** the + configured `upstream_model`. LiteLLM dispatches to the right + `model_list` entry. + +Tests in `orchestrator/tests/test_agent_model_resolution.py` and +`tests/gateway/test_anthropic_proxy.py` assert both invariants +explicitly — in particular that the Claude-Code-facing alias for a +LiteLLM-routed agent is `"opus"`, never the upstream model name. + +## Operator walkthrough — Qwen for the refiner role + +This walks through enabling hosted Qwen for a single role end-to-end. +None of these steps modify egg source code. + +### 1. Provision LiteLLM credentials + +Add the LiteLLM gateway master key (a random secret you generate; the +gateway forwards it on every LiteLLM-bound request as +`x-api-key`): + +```bash +# ~/.config/egg/secrets.env +LITELLM_MASTER_KEY= +``` + +This is the **only** gateway-held LiteLLM credential. Provider-side +keys (the hosted-Qwen API key, etc.) go in **LiteLLM's** environment, +not the gateway's `secrets.env` — see the next step. The gateway pod +rereads `secrets.env` on mtime change. + +### 2. Configure LiteLLM `model_list` + +Populate the LiteLLM ConfigMap with a `model_list` entry naming the +upstream-side model and provider: + +```yaml +# k8s/base/litellm-configmap.yaml (partial) +data: + config.yaml: | + model_list: + - model_name: qwen3-coder-30b + litellm_params: + model: together_ai/Qwen/Qwen2.5-Coder-32B-Instruct + api_key: os.environ/TOGETHER_API_KEY + general_settings: + master_key: os.environ/LITELLM_MASTER_KEY +``` + +Then surface `TOGETHER_API_KEY` (or your provider's equivalent) to the +LiteLLM Deployment — typically as a Secret env-var binding on the +LiteLLM pod, **not** through the gateway. Apply the change and roll +the LiteLLM Deployment. + +> **Hosted-provider choice.** Hosted Qwen is the cq-6 first target; +> any LiteLLM-supported backend works. Self-hosted vLLM / SGLang is +> deferred until the no-op-by-default path is validated against a +> hosted provider first. + +### 3a. Per-repository default (recommended for stable rollouts) + +If every role on this repo should default to the same model, edit +`~/.config/egg/repositories.yaml`: + +```yaml +repo_settings: + acme-corp/widgets: + # Every role on this repo defaults to qwen3-coder-30b, unless a + # pipeline passes its own agent_models entry. + default_agent_model: qwen3-coder-30b +``` + +`default_agent_model` is a *repo-level default for every role* — to +flip exactly one role to a non-Claude backend (the common case for +the cq-4 smoke test), prefer 3b. + +### 3b. Per-pipeline override (recommended for the smoke test) + +For one-off pipelines (or the smoke test below), pass `agent_models` +in the `submit_task` MCP-tool arguments +([`orchestrator/mcp_tools.py:74`](../../orchestrator/mcp_tools.py) — +`required: ["description", "repo"]`; `issue_number`, `branch`, and +`config` are optional): + +```json +{ + "description": "Smoke-test the refiner on hosted Qwen", + "repo": "acme-corp/widgets", + "issue_number": 1234, + "config": { + "agent_models": { + "refiner": "qwen3-coder-30b" + } + } +} +``` + +The equivalent `POST /pipelines` HTTP body (see +`orchestrator/routes/pipelines.py:1336` where the handler reads +`data.get("issue_number")`) uses the same field names plus an +optional `branch` override: + +```json +{ + "issue_number": 1234, + "repo": "acme-corp/widgets", + "branch": "egg/issue-1234/work", + "config": { + "agent_models": { + "refiner": "qwen3-coder-30b" + } + } +} +``` + +> The orchestrator silently ignores unrecognized top-level keys +> (`data.get("issue_number")` reads `issue_number` specifically), so +> a misspelled `"issue": 1234` would submit a pipeline with **no +> issue binding** without surfacing an error. Use `issue_number`. + +Per-pipeline `agent_models` entries **override** the repo-level +`default_agent_model`. Both can be unset — the resolver falls back to +the built-in `"opus"`. + +### 4. Run the pipeline and observe routing + +Submit a pipeline. The gateway audit log records the per-session +routing decision **once per session** (slice-1's +`audit_log("session_created", …)` extension at +`gateway/gateway.py:8920` includes the resolved `upstream` and +`upstream_model`); every subsequent `/v1/messages` request from that +session inherits the decision implicitly via the session-keyed +lookup, with no per-request routing log line: + +- **Refiner session-created line**: `upstream=litellm`, + `upstream_model=qwen3-coder-30b`. Subsequent refiner requests have + their body forwarded to + `litellm.egg-system.svc.cluster.local:4000` with the + `_rewrite_upstream_model` helper substituting the `"model"` field; + LiteLLM routes to the hosted Qwen backend. +- **Every other session-created line**: `upstream=anthropic`, + `upstream_model=null`. Subsequent requests have their body + forwarded byte-identically to `api.anthropic.com`. + +If anything is misconfigured (LiteLLM master key absent, LiteLLM pod +unreachable, etc.), the failure policy is **fail closed**: a 502 +surfaces to the agent. No silent fallback to Claude — this is cq-8 +and intentionally matches today's Anthropic-side failure shape so a +mixed transcript can't quietly erode the cost goal motivating the +work. See [Upstream Routing → Failure +policy](../architecture/upstream-routing.md#failure-policy). + +### 5. Exercise the cq-4 smoke test + +The two compatibility properties only the live path can prove: + +1. **Tool-heavy multi-turn loop.** Pick an issue whose refine phase + exercises the agent's tool surface (file reads, web fetches, MCP + tools). Watch the transcript: each tool call should round-trip + cleanly with no stream corruption (the + `claude-code-router`-style failure mode for Qwen thinking-mode + models was the explicit reason LiteLLM was chosen — confirm we are + not seeing it in practice). +2. **Auto-compaction boundary.** Run a session long enough that + Claude Code triggers its auto-compaction step. With the + recognized-alias mitigation, the compaction should fire on + schedule and the agent should resume cleanly. If the agent wedges + on an over-length request, the alias mitigation has failed and the + model needs a custom context-window override (an out-of-scope + follow-on for #2769). + +Capture the transcript via `egg-checkpoint show ` and the +gateway audit log via the structured-logging stream +([architecture/logging](../architecture/logging.md)). + +This validation step is **operator-driven and out of scope for +slice-2 merge** (per the cq-4 resolution): merging slice 2 ships only +the buildable seam. The smoke test runs once an operator has a live +LiteLLM endpoint to point at. + +## No-op-by-default invariant — three independent guards + +The combination of slices 1 and 2 keeps the LiteLLM client cold on a +deployment that has not been configured. The guards (described in +detail at [Upstream Routing → No-op-by-default +invariant](../architecture/upstream-routing.md#no-op-by-default-invariant)): + +1. **No LiteLLM master key configured.** The LiteLLM credential + resolver returns `None`; any request that *did* somehow route to + LiteLLM would fail credential injection with the standard 401. +2. **Session upstream defaults to `"anthropic"`.** Both the + `Session.upstream` default and the `/api/v1/sessions/create` + handler default are `"anthropic"`. Slice 2 only sets it when the + resolver returns a LiteLLM decision. +3. **`agent_models` default is empty.** Both + `PipelineConfig.agent_models` and repo-level + `default_agent_model` default to nothing — the resolver returns + the built-in `"opus"` Anthropic path. + +Any single guard suffices. All three are independent; a +misconfiguration on one does not silently activate the LiteLLM path. + +## Slice-2 primitives at a glance + +| Primitive | Location | Purpose | +|-----------|----------|---------| +| `PipelineConfig.agent_models: dict[str, str]` | `orchestrator/models.py:757` (alongside the existing `PipelineConfig` fields) | Per-pipeline, per-role model override; Pydantic validator rejects keys outside the phase producer / reviewer set (`agent_roles.MODEL_OVERRIDE_ROLES`) | +| `default_agent_model: str \| None` | `config/repositories.yaml.example` (documented schema) | Repository-level default applied when `agent_models` does not pin the role | +| `get_default_agent_model(repo)` | `config/repo_config.py` (mirrors `get_repo_setting` at `config/repo_config.py:248`) | Reader for the repo-level default; returns `None` when absent | +| `resolve_agent_model(role, pipeline_config, repo)` + `AgentModelDecision` | `orchestrator/agent_model_resolution.py` *(new module)* | Walks precedence + classifies into `(claude_code_alias, upstream, upstream_model)` | +| Spawn-side plumbing | `orchestrator/concurrent_executor.py:454`, `orchestrator/routes/pipelines.py:2704`, `orchestrator/kubernetes_spawner.py:735` | Threads `--model` to the consensus wrapper and `upstream` / `upstream_model` to `GatewayClient.register_session` | +| `_rewrite_upstream_model(request_body, upstream_model)` | `gateway/gateway.py` (adjacent to `_filter_blocked_tools` at `gateway/gateway.py:9496`); called from `proxy_anthropic_messages` (`gateway/gateway.py:9839`) and `proxy_count_tokens` (`gateway/gateway.py:10135`) | Rewrites the body's `"model"` field on LiteLLM-bound requests; no-op for `"anthropic"` and for invalid JSON | + +The slice-1 primitives this guide builds on (`UpstreamRegistry`, +`Session.upstream` / `upstream_model`, the LiteLLM credential +resolver, the LiteLLM k8s manifests) are catalogued in [Upstream +Routing → Files](../architecture/upstream-routing.md#files). + +## HITL decisions that shape this guide + +| Decision | Resolution | Where it shows up | +|----------|------------|-------------------| +| `cq-3` | Per-role field on `PipelineConfig` **and** `repositories.yaml` default | Two-knob config above; precedence chain in the resolver | +| `cq-4` | No agent-flip in this pipeline; operator smoke-test deferred | Smoke-test section is operator-driven, not gating merge | +| `cq-5` | Keep Claude Code; recognized-alias mitigation | `claude_code_alias` is always a Claude name; gateway rewrites body's `model` | +| `cq-6` | First validation backend is hosted Qwen | Step 2 example uses hosted Qwen; self-hosted deferred | +| `cq-8` | Fail closed on LiteLLM errors (502, no fallback) | Step 4 error-path note; same behavior as today's Anthropic upstream errors | +| `cq-11` | Leave `[1m]` for Claude | Non-Claude model strings simply do not carry `[1m]`; the resolver routes them via the LiteLLM path | + +The full set is at +[`.egg-state/contracts/issue-2769.json`](../../.egg-state/contracts/issue-2769.json). + +## Related Documentation + +- [Upstream Routing](../architecture/upstream-routing.md) — slice-1 + architecture: gateway router, `UpstreamRegistry`, per-session + routing, LiteLLM topology, credential layout, failure policy, and + the no-op-by-default invariant. +- [Credential Injection](../architecture/credential-injection.md) — + Anthropic credential resolver shape; the LiteLLM resolver follows + the same mtime-invalidated cache pattern. +- [Orchestrator Architecture](../architecture/orchestrator.md) — + Spawner / session-creation context that the slice-2 plumbing + hooks into. +- [Agent Roles](../reference/agent-roles.md) — Canonical role names + accepted as keys in `agent_models`. +- Issue [#2769](https://github.com/jwbron/egg/issues/2769) — Original + motivation, refine-phase decisions, plan-phase risk analysis. diff --git a/docs/index.md b/docs/index.md index e821e0ff78..bf89d2c692 100644 --- a/docs/index.md +++ b/docs/index.md @@ -22,6 +22,7 @@ This index helps both humans and LLMs navigate the documentation efficiently. | [Gateway Auto-Filter](architecture/gateway-auto-filter.md) | Restricted-path rejection on push (`403 restricted_path_modified`) and the commit-authorship registry that backs attribution | | [Credential Injection](architecture/credential-injection.md) | Zero-credential sandbox with API key proxy via gateway | | [Network Isolation](architecture/network-isolation.md) | Public/private network modes and domain allowlist | +| [Upstream Routing](architecture/upstream-routing.md) | `UpstreamRegistry` gateway seam for per-agent non-Claude backends, LiteLLM topology in `egg-system`, per-session routing decision, credential layout, and the no-op-by-default invariant ([#2769](https://github.com/jwbron/egg/issues/2769)) | | [Kubernetes Migration](architecture/kubernetes-migration.md) | Docker to k8s (k3s) migration: architecture, network isolation, developer workflow | | [SDLC Pipeline](architecture/sdlc-pipeline.md) | Structurally enforced agent checkpoints and verification gates | | [Slice-DAG Implement Phase](architecture/slice-dag.md) | `Phase` → `Slice` schema rename, forest validation, slice scheduler (waves, two-tier `max_cycles`, failure cascade), stacked-PR reconciler, per-slice branches and BRC trackers | @@ -58,6 +59,7 @@ This index helps both humans and LLMs navigate the documentation efficiently. | [Anchor Recovery](guides/anchor-recovery.md) | Agent post-compaction state recovery via persistent anchors | | [Deployment Diagnostics](guides/deployment-diagnostics.md) | When to use `/deployment-diagnose` vs `/agent-diagnose`, evidence boundaries, and redaction guarantees | | [File Decomposition Pattern](guides/decomposition-pattern.md) | Canonical sub-package + explicit re-export barrel pattern for decomposing oversize Python files under the `scripts/file-size-allowlist.yaml` cap; covers conversion mechanics, method-modules-on-class shape, audit recipe, allowlist rebase, and routes-handling convention | +| [Per-Agent Models](guides/per-agent-models.md) | Operator guide for flipping a single agent role to a non-Claude model via the LiteLLM proxy: `PipelineConfig.agent_models` + `default_agent_model` precedence, the `resolve_agent_model` classifier, the gateway-side `_rewrite_upstream_model` helper, the recognized-alias compaction mitigation (cq-5), and the end-to-end Qwen smoke test ([#2769](https://github.com/jwbron/egg/issues/2769)) | ### Deploy @@ -148,6 +150,7 @@ Each major component has detailed documentation: | **Health check framework** | [Health Checks README](../orchestrator/health_checks/README.md) | [Orchestrator Architecture](architecture/orchestrator.md), [Orchestrator README](../orchestrator/README.md) | | **Pipeline health monitoring** | [Pipeline Health Monitoring](guides/pipeline-health-monitoring.md) | [Health Checks README](../orchestrator/health_checks/README.md), [Agent Roles](reference/agent-roles.md), [Orchestrator Architecture](architecture/orchestrator.md) | | **Generating repository documentation** | [GitHub Automation: Documentation Onboarding](guides/github-automation.md#documentation-onboarding) | [Onboarding prompt](../shared/prompts/onboarding-docs-prompt.md), `egg-onboarding-docs` CLI | +| **Running a single agent on a non-Claude model** | [Per-Agent Models](guides/per-agent-models.md) | [Upstream Routing](architecture/upstream-routing.md), [Agent Roles](reference/agent-roles.md), `orchestrator/agent_model_resolution.py` | ## Quick Navigation diff --git a/gateway/CLAUDE.md b/gateway/CLAUDE.md index d40708df87..9d16b96ad2 100644 --- a/gateway/CLAUDE.md +++ b/gateway/CLAUDE.md @@ -4,6 +4,7 @@ Policy-enforcement sidecar that sits between agents and GitHub. Validates git/gh - **[README.md](README.md)** — architecture, policy rules, configuration - **[../docs/index.md](../docs/index.md)** — full documentation index +- **[../docs/architecture/upstream-routing.md](../docs/architecture/upstream-routing.md)** — `UpstreamRegistry` seam, LiteLLM topology, per-session routing decision, and the no-op-by-default invariant for non-Claude agent backends ([#2769](https://github.com/jwbron/egg/issues/2769)) ## Testing diff --git a/gateway/anthropic_credentials.py b/gateway/anthropic_credentials.py index 4cc20ba2c3..c43ff06529 100644 --- a/gateway/anthropic_credentials.py +++ b/gateway/anthropic_credentials.py @@ -227,3 +227,102 @@ def reset_credentials_manager() -> None: """Reset the global credentials manager (for testing).""" global _credentials_manager _credentials_manager = None + + +# ============================================================================= +# LiteLLM Credentials +# ============================================================================= +# LiteLLM is an optional non-Anthropic upstream (see issue #2769 cq-7). The +# gateway holds a single LiteLLM "master key" in the same secrets.env file as +# ANTHROPIC_API_KEY and injects it as `x-api-key` on every LiteLLM-bound +# request. LiteLLM itself holds the real per-backend credentials (hosted +# provider API keys) via its own env-var slots, so a missing +# ``LITELLM_MASTER_KEY`` here disables routing-to-LiteLLM rather than +# silently downgrading the security boundary. + + +class LiteLLMCredentialsManager: + """Manages the LiteLLM master key for gateway proxy injection. + + Mirrors ``AnthropicCredentialsManager`` — reads + ``LITELLM_MASTER_KEY`` out of ``secrets.env`` with the same + mtime-invalidated cache, and returns an ``AnthropicCredential`` shaped + ``(header_name='x-api-key', header_value=)``. Returns ``None`` when + the key is absent so the proxy route falls into the same "no credential" + error branch as today's Anthropic path. + """ + + def __init__(self, secrets_path: Path | None = None) -> None: + self._secrets_path = secrets_path or SECRETS_PATH + self._credential: AnthropicCredential | None = None + self._cached_mtime: float = 0 + self._lock = threading.Lock() + + def get_credential(self) -> AnthropicCredential | None: + """Return the cached LiteLLM credential, reloading on mtime change.""" + try: + current_mtime = self._secrets_path.stat().st_mtime + except OSError: + with self._lock: + self._credential = None + self._cached_mtime = 0 + return None + + with self._lock: + if current_mtime != self._cached_mtime: + self._load_credential() + self._cached_mtime = current_mtime + return self._credential + + def _load_credential(self) -> None: + if not self._secrets_path.exists(): + # Not warning here: a missing secrets file is the default for + # operators who have not opted into LiteLLM routing. + self._credential = None + return + + secrets = parse_env_file(self._secrets_path) + if not secrets: + # Empty/unreadable secrets.env was already warned about by the + # Anthropic manager; don't double-warn for the LiteLLM resolver. + self._credential = None + return + + master_key = secrets.get("LITELLM_MASTER_KEY", "").strip() + if not master_key: + # Absent key is the no-op default — no warning (see docstring). + self._credential = None + return + + self._credential = AnthropicCredential( + header_name="x-api-key", + header_value=master_key, + ) + logger.info( + "LiteLLM master key loaded from secrets", + key_prefix=master_key[:8] + "..." if len(master_key) >= 8 else "", + ) + + def reload(self) -> None: + """Force reload of the LiteLLM credential (for testing / config updates).""" + with self._lock: + self._cached_mtime = 0 + self._credential = None + + +# Global LiteLLM credentials manager instance. +_litellm_credentials_manager: LiteLLMCredentialsManager | None = None + + +def get_litellm_credentials_manager() -> LiteLLMCredentialsManager: + """Get or create the global LiteLLM credentials manager.""" + global _litellm_credentials_manager + if _litellm_credentials_manager is None: + _litellm_credentials_manager = LiteLLMCredentialsManager() + return _litellm_credentials_manager + + +def reset_litellm_credentials_manager() -> None: + """Reset the global LiteLLM credentials manager (for testing).""" + global _litellm_credentials_manager + _litellm_credentials_manager = None diff --git a/gateway/gateway.py b/gateway/gateway.py index 1ccd9c6860..34051296bc 100644 --- a/gateway/gateway.py +++ b/gateway/gateway.py @@ -84,7 +84,10 @@ check_agent_gh_operation, get_agent_pattern, # noqa: F401 — re-exported for test patching ) - from .anthropic_credentials import get_credentials_manager + from .anthropic_credentials import ( + get_credentials_manager, + get_litellm_credentials_manager, + ) from .checkpoint_handler import ( _get_checkpoint_repo_for_path, capture_and_store_checkpoint, @@ -201,6 +204,10 @@ validate_session_for_request, ) from .transcript_buffer import get_transcript_buffer + from .upstream_registry import ( + UnknownUpstreamError, + get_upstream_registry, + ) from .worktree_manager import ( REPOS_BASE_DIR, WORKTREE_BASE_DIR, @@ -214,7 +221,10 @@ check_agent_gh_operation, get_agent_pattern, # noqa: F401 — re-exported for test patching ) - from anthropic_credentials import get_credentials_manager # type: ignore[no-redef] + from anthropic_credentials import ( # type: ignore[no-redef] + get_credentials_manager, + get_litellm_credentials_manager, + ) from checkpoint_handler import ( # type: ignore[no-redef, import-untyped] _get_checkpoint_repo_for_path, capture_and_store_checkpoint, @@ -350,6 +360,10 @@ validate_session_for_request, ) from transcript_buffer import get_transcript_buffer # type: ignore[no-redef, import-untyped] + from upstream_registry import ( # type: ignore[no-redef, import-untyped] + UnknownUpstreamError, + get_upstream_registry, + ) from worktree_manager import ( # type: ignore[no-redef, import-untyped] REPOS_BASE_DIR, WORKTREE_BASE_DIR, @@ -8573,6 +8587,10 @@ def session_create() -> tuple[Response, int] | Response: branch = data.get("branch") # Optional git branch for non-pushing sessions jira_ticket = data.get("jira_ticket") # Optional Atlassian ticket key — advisory only synthetic = data.get("synthetic", False) # Orchestrator-internal temp session + # Per-session upstream routing (issue #2769). Default to "anthropic" so + # pre-#2769 callers keep byte-identical session shape. + upstream = data.get("upstream", "anthropic") + upstream_model = data.get("upstream_model") # Validate required fields if not container_id: @@ -8653,6 +8671,22 @@ def session_create() -> tuple[Response, int] | Response: if not isinstance(synthetic, bool): return make_error("Invalid synthetic: must be a boolean") + # Validate upstream / upstream_model (issue #2769). + # ``upstream`` must be a name the UpstreamRegistry will serve — refuse + # silently routing unknown upstreams to Anthropic. + if not isinstance(upstream, str): + return make_error("Invalid upstream: must be a string") + if not get_upstream_registry().is_known(upstream): + known = ", ".join(sorted(get_upstream_registry().known_upstreams())) + return make_error(f"Invalid upstream: '{upstream}'. Must be one of: {known}") + if upstream_model is not None: + if not isinstance(upstream_model, str): + return make_error("Invalid upstream_model: must be a string") + if not upstream_model: + return make_error("Invalid upstream_model: must be non-empty if provided") + if len(upstream_model) > 256: + return make_error("Invalid upstream_model: must be 256 characters or fewer") + # Validate worktree_container_id if provided if worktree_container_id is not None: if not isinstance(worktree_container_id, str): @@ -8857,6 +8891,8 @@ def session_create() -> tuple[Response, int] | Response: branch=branch, jira_ticket=jira_ticket if isinstance(jira_ticket, str) and jira_ticket else None, synthetic=synthetic, + upstream=upstream, + upstream_model=upstream_model, ) # Pre-populate checkpoint context so non-pushing sessions (reviewers, @@ -8896,6 +8932,8 @@ def session_create() -> tuple[Response, int] | Response: "filtered_repos": filtered_repos, "worktree_count": len(worktrees), "worktree_errors": worktree_errors if worktree_errors else None, + "upstream": upstream, + "upstream_model": upstream_model, }, ) @@ -9352,16 +9390,78 @@ def _filter_response_headers(headers: Any) -> dict[str, str]: return {k: v for k, v in headers.items() if k.lower() not in skip} -def _inject_anthropic_credentials( +def _inject_upstream_credentials( headers: dict[str, str], + upstream: str = "anthropic", ) -> tuple[dict[str, str], tuple[Any, int] | None]: """ - Inject Anthropic credentials into headers. + Inject upstream credentials into headers. + + Dispatches per-upstream so the gateway can carry both Anthropic and + LiteLLM credentials side-by-side (issue #2769 cq-7). For the Anthropic + upstream this is byte-identical to the legacy + ``_inject_anthropic_credentials`` helper — same OAuth/API-key precedence, + same 401 error shape on missing credentials, same client-supplied auth + fall-through. The LiteLLM upstream uses ``x-api-key`` only (no OAuth + path) and has no client-supplied-auth fall-through because Claude Code + never carries a LiteLLM master key. + + An upstream the registry does not serve is rejected with a 502 — it + is never silently treated as Anthropic. + + Args: + headers: Mutable header dict — credential is appended in place. + upstream: ``"anthropic"`` (default — back-compat) or ``"litellm"``. Returns: (headers, None) on success (headers, error_response_tuple) on failure - caller should return this """ + # Refuse to silently treat an unknown upstream as Anthropic. Falling + # through to the Anthropic branch produced an observable error-code + # inconsistency — 401 vs 502 for the same invalid input depending on + # unrelated Anthropic-credential state. An unregistered upstream now + # fails closed with a 502, matching the proxy routes' own + # UnknownUpstreamError handling (issue #2769 review). + if not get_upstream_registry().is_known(upstream): + logger.warning( + "Unknown upstream for credential injection, refusing request", + upstream=upstream, + ) + return headers, ( + jsonify( + { + "error": { + "type": "api_error", + "message": f"Unknown upstream '{upstream}'", + } + } + ), + 502, + ) + + if upstream == "litellm": + cred = get_litellm_credentials_manager().get_credential() + if cred: + headers[cred.header_name] = cred.header_value + return headers, None + logger.warning( + "No LiteLLM master key available for proxy request", + upstream=upstream, + ) + return headers, ( + jsonify( + { + "error": { + "type": "authentication_error", + "message": "No LiteLLM credentials available", + } + } + ), + 401, + ) + + # Default: anthropic upstream — preserves the legacy behavior verbatim. credentials_manager = get_credentials_manager() cred = credentials_manager.get_credential() @@ -9397,6 +9497,18 @@ def _inject_anthropic_credentials( ) +def _inject_anthropic_credentials( + headers: dict[str, str], +) -> tuple[dict[str, str], tuple[Any, int] | None]: + """Back-compat alias delegating to the upstream-aware injector. + + Kept so external test mocks targeting ``_inject_anthropic_credentials`` + continue to work. New code paths should call + ``_inject_upstream_credentials(headers, upstream)`` directly. + """ + return _inject_upstream_credentials(headers, upstream="anthropic") + + # Tools blocked in private mode to prevent data exfiltration # These tools route through Anthropic's infrastructure, bypassing container network controls # See PR #686 security findings and PR #702 analysis @@ -9459,6 +9571,63 @@ def _filter_blocked_tools(request_body: bytes, session_mode: str | None) -> byte return request_body +def _rewrite_upstream_model(request_body: bytes, upstream_model: str | None) -> bytes: + """ + Rewrite the request body's top-level ``"model"`` field to *upstream_model*. + + This is the gateway-side half of the cq-5 mitigation for per-agent + non-Claude routing (#2769 task-2-6). Claude Code is handed a recognised + Claude alias (``"opus"``) so its compaction heuristics stay calibrated + against a known family, but LiteLLM has no ``"opus"`` entry in its + ``model_list`` — it routes on whatever model name appears in the request + body. The orchestrator threads the upstream-side model name onto the + session at spawn time (``Session.upstream_model``); this helper swaps it + into the body before the gateway forwards to the non-Anthropic upstream. + + Mirrors ``_filter_blocked_tools``'s bytes-in / bytes-out contract: parse + failures and a ``None`` *upstream_model* both return the original bytes + unchanged so the regression guard on the Anthropic path is byte-identical + and a malformed body never escapes as a 500. Callers MUST guard with + ``if upstream_name != "anthropic"`` so the helper is only invoked on the + LiteLLM-routed path; on Anthropic, the body must be forwarded verbatim. + + Args: + request_body: Raw JSON request body. + upstream_model: Upstream-side model name to substitute, or ``None`` + (no-op — body returned unchanged). + + Returns: + Modified request body with ``"model"`` set to *upstream_model* (if + rewriting was applied) or the original body unchanged. + """ + if upstream_model is None: + return request_body + + try: + body = json.loads(request_body) + except (json.JSONDecodeError, TypeError) as e: + logger.warning( + "Failed to parse request body for upstream-model rewrite", + error=str(e), + ) + return request_body + + if not isinstance(body, dict): + return request_body + + original = body.get("model") + if original == upstream_model: + return request_body + + body["model"] = upstream_model + logger.info( + "Rewrote upstream model in request body", + original_model=original, + upstream_model=upstream_model, + ) + return json.dumps(body).encode() + + def _is_streaming_request(request_body: bytes) -> bool: """ Check if request body indicates streaming mode. @@ -9762,22 +9931,37 @@ def proxy_anthropic_messages() -> tuple[Response, int] | Response: """ start_time = time.time() - # Build headers with injected auth - headers = _get_forwarded_headers(request.headers) - headers, error = _inject_anthropic_credentials(headers) - if error: - return error - - request_body = request.get_data() - # Look up session by IP to determine mode (Claude Code doesn't send session tokens) session_manager = get_session_manager() session = session_manager.get_session_by_ip(request.remote_addr or "") session_mode = session.mode if session else None container_id = session.container_id if session else None + # Resolve per-session upstream (issue #2769). With no session, default to + # "anthropic" so today's Claude path is byte-identical when an unrelated + # client probes /v1/messages without first registering a session. + upstream_name = session.upstream if session else "anthropic" + + # Build headers with injected auth — dispatches per-upstream so the + # LiteLLM master key is never paired with an Anthropic request and vice + # versa. + headers = _get_forwarded_headers(request.headers) + headers, error = _inject_upstream_credentials(headers, upstream=upstream_name) + if error: + return error + + request_body = request.get_data() request_body = _filter_blocked_tools( request_body, session_mode ) # Remove web tools in private mode + # cq-5 mitigation (#2769 task-2-6): on the LiteLLM-routed path, swap + # the cq-5 alias ("opus") that Claude Code was handed for the + # upstream-side model name LiteLLM's model_list keys on. The + # Anthropic path keeps the original body byte-identical so today's + # wire shape is preserved. + if upstream_name != "anthropic": + request_body = _rewrite_upstream_model( + request_body, session.upstream_model if session else None + ) is_streaming = _is_streaming_request(request_body) # Parse request body for transcript capture @@ -9786,7 +9970,30 @@ def proxy_anthropic_messages() -> tuple[Response, int] | Response: except json.JSONDecodeError, TypeError: request_json = {} - client = get_anthropic_client() + # Resolve the upstream httpx client per request. The Anthropic path + # keeps calling ``get_anthropic_client()`` so behavior — including + # existing test mocks patching that symbol — is byte-identical to + # today. Non-Anthropic upstreams (currently only ``"litellm"``) resolve + # through the registry, which is the swap-out seam (issue #2769 + # feedback Q3). + if upstream_name == "anthropic": + client = get_anthropic_client() + else: + try: + client, _ = get_upstream_registry().get(upstream_name) + except UnknownUpstreamError: + logger.warning( + "Unknown upstream on session, refusing request", + upstream=upstream_name, + ) + return jsonify( + { + "error": { + "type": "api_error", + "message": f"Unknown upstream '{upstream_name}'", + } + } + ), 502 try: if is_streaming: @@ -9983,34 +10190,34 @@ def generate() -> Any: ) except httpx.ConnectError as e: - logger.error("Anthropic API connection failed", error=str(e)) + logger.error("Upstream connection failed", upstream=upstream_name, error=str(e)) return jsonify( { "error": { "type": "api_error", - "message": f"Failed to connect to Anthropic API: {e}", + "message": f"Failed to connect to {upstream_name} upstream: {e}", } } ), 502 except httpx.TimeoutException as e: - logger.error("Anthropic API request timed out", error=str(e)) + logger.error("Upstream request timed out", upstream=upstream_name, error=str(e)) return jsonify( { "error": { "type": "api_error", - "message": f"Anthropic API request timed out: {e}", + "message": f"{upstream_name} upstream request timed out: {e}", } } ), 504 except Exception as e: - logger.exception("Anthropic API proxy error") + logger.exception("Upstream proxy error", upstream=upstream_name) return jsonify( { "error": { "type": "api_error", - "message": f"Anthropic API proxy error: {e}", + "message": f"{upstream_name} upstream proxy error: {e}", } } ), 502 @@ -10024,18 +10231,52 @@ def proxy_count_tokens() -> tuple[Response, int] | Response: This endpoint allows Claude Code to use ANTHROPIC_BASE_URL to route token counting requests through the gateway. """ + # Mirror the per-session upstream lookup used by proxy_anthropic_messages + # so count_tokens and messages always agree on which backend serves a + # given agent (issue #2769). + session_manager = get_session_manager() + session = session_manager.get_session_by_ip(request.remote_addr or "") + upstream_name = session.upstream if session else "anthropic" + headers = _get_forwarded_headers(request.headers) - headers, error = _inject_anthropic_credentials(headers) + headers, error = _inject_upstream_credentials(headers, upstream=upstream_name) if error: return error - client = get_anthropic_client() + if upstream_name == "anthropic": + client = get_anthropic_client() + else: + try: + client, _ = get_upstream_registry().get(upstream_name) + except UnknownUpstreamError: + logger.warning( + "Unknown upstream on session, refusing request", + upstream=upstream_name, + ) + return jsonify( + { + "error": { + "type": "api_error", + "message": f"Unknown upstream '{upstream_name}'", + } + } + ), 502 + + # cq-5 mitigation (#2769 task-2-6): mirror proxy_anthropic_messages — + # the LiteLLM-routed path must have the cq-5 alias swapped for the + # upstream-side model name LiteLLM keys on, while the Anthropic path + # forwards bytes verbatim so the regression guard stays intact. + count_tokens_body = request.get_data() + if upstream_name != "anthropic": + count_tokens_body = _rewrite_upstream_model( + count_tokens_body, session.upstream_model if session else None + ) try: response = client.post( "/v1/messages/count_tokens", headers=headers, - content=request.get_data(), + content=count_tokens_body, ) return Response( response.content, @@ -10044,34 +10285,34 @@ def proxy_count_tokens() -> tuple[Response, int] | Response: ) except httpx.ConnectError as e: - logger.error("Anthropic API connection failed", error=str(e)) + logger.error("Upstream connection failed", upstream=upstream_name, error=str(e)) return jsonify( { "error": { "type": "api_error", - "message": f"Failed to connect to Anthropic API: {e}", + "message": f"Failed to connect to {upstream_name} upstream: {e}", } } ), 502 except httpx.TimeoutException as e: - logger.error("Anthropic API request timed out", error=str(e)) + logger.error("Upstream request timed out", upstream=upstream_name, error=str(e)) return jsonify( { "error": { "type": "api_error", - "message": f"Anthropic API request timed out: {e}", + "message": f"{upstream_name} upstream request timed out: {e}", } } ), 504 except Exception as e: - logger.exception("Anthropic API proxy error") + logger.exception("Upstream proxy error", upstream=upstream_name) return jsonify( { "error": { "type": "api_error", - "message": f"Anthropic API proxy error: {e}", + "message": f"{upstream_name} upstream proxy error: {e}", } } ), 502 diff --git a/gateway/session_manager.py b/gateway/session_manager.py index 7a66776f56..026fab496e 100644 --- a/gateway/session_manager.py +++ b/gateway/session_manager.py @@ -325,6 +325,14 @@ class Session: auto_commit_sha: str | None = None # SHA from post-agent auto-commit jira_ticket: str | None = None # Advisory Jira ticket key (issue #1556) synthetic: bool = False # Orchestrator-internal temp session — skip checkpoint capture + # Per-session upstream routing (issue #2769). ``upstream`` chooses which + # registered UpstreamRegistry entry serves /v1/messages traffic for this + # session ("anthropic" — default — or "litellm"). ``upstream_model`` is the + # upstream-side model name LiteLLM expects when the gateway rewrites the + # request body's ``model`` field (slice 2). Default-Anthropic sessions + # leave both at their defaults and are byte-identical to pre-#2769 sessions. + upstream: str = "anthropic" + upstream_model: str | None = None def is_expired(self) -> bool: """Check if session has expired.""" @@ -374,6 +382,12 @@ def to_dict_for_persistence(self) -> dict[str, Any]: result["jira_ticket"] = self.jira_ticket if self.synthetic: result["synthetic"] = True + # Only persist the upstream fields when they differ from the default — + # keeps disk layout byte-identical for the Claude-only path (issue #2769). + if self.upstream != "anthropic": + result["upstream"] = self.upstream + if self.upstream_model is not None: + result["upstream_model"] = self.upstream_model return result @classmethod @@ -402,6 +416,8 @@ def from_persistence(cls, data: dict[str, Any]) -> Session: auto_commit_sha=data.get("auto_commit_sha"), jira_ticket=data.get("jira_ticket"), synthetic=bool(data.get("synthetic", False)), + upstream=data.get("upstream", "anthropic"), + upstream_model=data.get("upstream_model"), ) @@ -560,6 +576,8 @@ def register_session( branch: str | None = None, jira_ticket: str | None = None, synthetic: bool = False, + upstream: str = "anthropic", + upstream_model: str | None = None, ) -> tuple[str, Session]: """ Register a new session for a container. @@ -576,10 +594,46 @@ def register_session( agent_anchor_id: Optional agent anchor ID for scoped anchor file writes claude_code_version: Optional Claude Code version string branch: Optional git branch for non-pushing pipeline sessions + upstream: Upstream registry name driving ``/v1/messages`` + traffic for this session (issue #2769). Defaults to + ``"anthropic"`` so all pre-#2769 callers keep byte-identical + behavior; pass ``"litellm"`` to route through the LiteLLM + proxy in egg-system. Validated against ``UpstreamRegistry``. + upstream_model: Optional upstream-side model name used by the + slice-2 body-rewrite path. ``None`` (default) leaves the + incoming request body's ``model`` field unchanged. Returns: Tuple of (session_token, Session) + + Raises: + ValueError: If ``upstream`` is not a name the gateway's + ``UpstreamRegistry`` serves, or if ``upstream_model`` is + given but is not a non-empty string of ≤256 characters. """ + # Validate the upstream against the registry (issue #2769). The + # /api/v1/sessions/create route validates too, but guard here so a + # direct caller (the slice-2 spawner, tests) cannot register a + # session whose upstream the gateway will not serve — that would + # land in the proxy routes' inconsistent unknown-upstream path. + from upstream_registry import get_upstream_registry # type: ignore[import-untyped] + + registry = get_upstream_registry() + if not registry.is_known(upstream): + known = ", ".join(sorted(registry.known_upstreams())) + raise ValueError(f"Unknown upstream '{upstream}'. Must be one of: {known}") + + # Mirror the /api/v1/sessions/create route's upstream_model checks + # (issue #2769) so a direct caller cannot store a malformed model + # name that would only surface in the slice-2 body-rewrite path. + if upstream_model is not None: + if not isinstance(upstream_model, str): + raise ValueError("upstream_model must be a string") + if not upstream_model: + raise ValueError("upstream_model must be non-empty if provided") + if len(upstream_model) > 256: + raise ValueError("upstream_model must be 256 characters or fewer") + # Generate cryptographically secure token token = secrets.token_urlsafe(SESSION_TOKEN_BYTES) token_hash = _hash_token(token) @@ -603,6 +657,8 @@ def register_session( claude_code_version=claude_code_version, jira_ticket=jira_ticket, synthetic=synthetic, + upstream=upstream, + upstream_model=upstream_model, ) if branch: diff --git a/gateway/tests/test_session_manager.py b/gateway/tests/test_session_manager.py index 9e8031f140..9ef2552956 100644 --- a/gateway/tests/test_session_manager.py +++ b/gateway/tests/test_session_manager.py @@ -2148,3 +2148,191 @@ def test_stop_joins_thread(self, manager): manager.stop_background_pruner(timeout=2.0) assert manager._prune_thread is None assert manager._prune_shutdown is None + + +# =========================================================================== +# Upstream fields — slice-1 of issue #2769 (TASK-1-4) +# =========================================================================== +# +# Slice 1 adds two optional fields to the Session dataclass: +# - ``upstream: str = "anthropic"`` +# - ``upstream_model: str | None = None`` +# +# The fields must round-trip through ``to_dict_for_persistence`` / +# ``from_persistence`` and existing persisted dicts WITHOUT the fields +# must rehydrate cleanly (back-compat with persisted sessions on disk). +# ``SessionManager.register_session`` must accept the two optional kwargs +# and store them on the returned Session. +# =========================================================================== + + +class TestSessionUpstreamFields: + """Slice-1 ``upstream`` / ``upstream_model`` Session fields.""" + + def _make_session(self, **overrides): + now = datetime.now(UTC) + defaults = { + "session_token": "test-token", + "session_token_hash": _hash_token("test-token"), + "container_id": "test-container", + "container_ip": "172.18.0.5", + "mode": "private", + "created_at": now, + "last_seen": now, + "expires_at": now + timedelta(hours=24), + } + defaults.update(overrides) + return Session(**defaults) + + def test_default_upstream_is_anthropic(self): + """A Session built without upstream= keeps the no-op default.""" + session = self._make_session() + assert session.upstream == "anthropic" + assert session.upstream_model is None + + def test_explicit_upstream_litellm(self): + session = self._make_session(upstream="litellm", upstream_model="qwen3-coder-30b") + assert session.upstream == "litellm" + assert session.upstream_model == "qwen3-coder-30b" + + def test_roundtrip_default_upstream(self): + """Default values round-trip through persistence.""" + session = self._make_session() + d = session.to_dict_for_persistence() + restored = Session.from_persistence(d) + assert restored.upstream == "anthropic" + assert restored.upstream_model is None + + def test_roundtrip_explicit_litellm(self): + session = self._make_session(upstream="litellm", upstream_model="qwen3-coder-30b") + d = session.to_dict_for_persistence() + restored = Session.from_persistence(d) + assert restored.upstream == "litellm" + assert restored.upstream_model == "qwen3-coder-30b" + + def test_persisted_dict_without_upstream_fields_rehydrates_cleanly(self): + """Back-compat guard: existing on-disk sessions written before + slice-1 land MUST still load with default ``upstream='anthropic'`` + and ``upstream_model=None``. This is the most important slice-1 + invariant for session_manager — a missed default breaks every + already-running gateway on upgrade. + """ + now = datetime.now(UTC) + legacy_dict = { + "session_token_hash": _hash_token("legacy-token"), + "container_id": "legacy-container", + "container_ip": "172.18.0.5", + "mode": "private", + "created_at": now.isoformat(), + "last_seen": now.isoformat(), + "expires_at": (now + timedelta(hours=24)).isoformat(), + # NOTE: no "upstream" or "upstream_model" keys — the legacy + # shape from before #2769 slice-1. + } + restored = Session.from_persistence(legacy_dict) + assert restored.upstream == "anthropic" + assert restored.upstream_model is None + + def test_anthropic_default_omitted_from_persistence(self): + """Mirroring the existing pattern (synthetic etc.), the default + ``upstream='anthropic'`` should not bloat the persisted dict. + + This is a lenient guard: the conditional asserts below accept + either an omitted field (preferred) or one present at its + default value — an always-emitted default is acceptable as long + as the back-compat read still works. The test fails only if a + non-default value is persisted unexpectedly. + """ + session = self._make_session() + d = session.to_dict_for_persistence() + # Either omitted entirely (preferred) OR present with the + # default value — both are acceptable. The test fails ONLY + # if a non-default value is present unexpectedly. + if "upstream" in d: + assert d["upstream"] == "anthropic" + if "upstream_model" in d: + assert d["upstream_model"] is None + + +class TestSessionManagerRegisterUpstream: + """``SessionManager.register_session`` accepts the new upstream kwargs.""" + + @pytest.fixture + def manager(self, tmp_path): + return SessionManager(persistence_file=tmp_path / "sessions.json") + + def test_register_without_upstream_kwargs_defaults_to_anthropic(self, manager): + """Back-compat: existing callers that don't pass the new kwargs + get the no-op Anthropic default.""" + _token, session = manager.register_session( + container_id="legacy-caller", + container_ip="172.18.0.5", + mode="private", + ) + assert session.upstream == "anthropic" + assert session.upstream_model is None + + def test_register_with_litellm_upstream_stores_both_fields(self, manager): + """Explicit LiteLLM registration stores both fields on the Session.""" + _token, session = manager.register_session( + container_id="qwen-agent", + container_ip="172.18.0.7", + mode="private", + upstream="litellm", + upstream_model="qwen3-coder-30b", + ) + assert session.upstream == "litellm" + assert session.upstream_model == "qwen3-coder-30b" + + def test_register_with_anthropic_upstream_explicit(self, manager): + """Explicit ``upstream='anthropic'`` is a valid no-op.""" + _token, session = manager.register_session( + container_id="explicit-anthropic-agent", + container_ip="172.18.0.8", + mode="private", + upstream="anthropic", + ) + assert session.upstream == "anthropic" + assert session.upstream_model is None + + def test_register_with_unknown_upstream_raises(self, manager): + """An upstream the gateway cannot serve is rejected at registration + (issue #2769 review) — a direct caller must not be able to create + a bogus-upstream session that bypasses the session-create route's + validation. + """ + with pytest.raises(ValueError, match="Unknown upstream"): + manager.register_session( + container_id="bogus-agent", + container_ip="172.18.0.9", + mode="private", + upstream="bogus_upstream_name", + ) + + def test_register_with_empty_upstream_model_raises(self, manager): + """An empty ``upstream_model`` is rejected at registration, mirroring + the session-create route's validation (issue #2769 review) — a + direct caller must not be able to store a malformed model name. + """ + with pytest.raises(ValueError, match="upstream_model must be non-empty"): + manager.register_session( + container_id="empty-model-agent", + container_ip="172.18.0.10", + mode="private", + upstream="litellm", + upstream_model="", + ) + + def test_register_with_oversized_upstream_model_raises(self, manager): + """An ``upstream_model`` over 256 characters is rejected at + registration, mirroring the session-create route (issue #2769 + review). + """ + with pytest.raises(ValueError, match="256 characters or fewer"): + manager.register_session( + container_id="long-model-agent", + container_ip="172.18.0.11", + mode="private", + upstream="litellm", + upstream_model="x" * 257, + ) diff --git a/gateway/upstream_registry.py b/gateway/upstream_registry.py new file mode 100644 index 0000000000..c45979882b --- /dev/null +++ b/gateway/upstream_registry.py @@ -0,0 +1,201 @@ +""" +Upstream Registry for Gateway LLM Proxy. + +Provides a per-upstream registry pairing an httpx.Client (base_url, timeout, +connection limits) with a credential resolver. Used by the +``/v1/messages`` and ``/v1/messages/count_tokens`` proxy routes to select +between today's Anthropic upstream and a future LiteLLM-translation proxy +in ``egg-system``. + +The registry is the gateway-side trust boundary for upstream selection: +the orchestrator declares the per-agent upstream at session-create time +(``Session.upstream``), and the proxy route resolves the registry entry +per request before injecting credentials and forwarding bytes upstream. + +With no agent configured for LiteLLM, no LiteLLM-bound request ever fires +— ``UpstreamRegistry`` is the seam, not a behavior change. See cq-1 and +cq-7 on issue #2769. +""" + +from __future__ import annotations + +import os +import sys +import threading +from collections.abc import Callable +from pathlib import Path + +import httpx + +# Add shared directory to path for egg_logging +_shared_path = Path(__file__).parent.parent / "shared" +if _shared_path.exists(): + sys.path.insert(0, str(_shared_path)) +from egg_logging import get_logger + +# Import gateway modules - try relative import first (module mode), +# fall back to absolute import (standalone / pytest load path). Mirrors the +# pattern used by gateway/gateway.py so this module is importable both as +# ``gateway.upstream_registry`` and as a top-level ``upstream_registry``. +try: + from .anthropic_credentials import ( + AnthropicCredential, + LiteLLMCredentialsManager, + get_credentials_manager, + get_litellm_credentials_manager, + ) +except ImportError: # pragma: no cover - exercised by standalone import paths + _gateway_dir = str(Path(__file__).parent) + if _gateway_dir not in sys.path: + sys.path.insert(0, _gateway_dir) + from anthropic_credentials import ( # type: ignore[no-redef] + AnthropicCredential, + LiteLLMCredentialsManager, + get_credentials_manager, + get_litellm_credentials_manager, + ) + +logger = get_logger("gateway.upstream-registry") + + +# Anthropic upstream base URL — matches today's hard-wired client. +ANTHROPIC_BASE_URL = "https://api.anthropic.com" # noqa: EGG200 - proxy target URL, not a direct LLM call + +# LiteLLM proxy Service DNS — overridable via env var so operators can +# point at a different proxy without rebuilding the gateway image. +LITELLM_BASE_URL_DEFAULT = "http://litellm.egg-system.svc.cluster.local:4000" + +# The upstream names this registry serves. Single source of truth for +# ``get`` / ``is_known`` / ``known_upstreams``; adding a fourth upstream +# means adding a name here and an ``_ensure_`` constructor in ``get``. +KNOWN_UPSTREAMS: tuple[str, ...] = ("anthropic", "litellm") + + +# Type alias for the credential resolver shape — both anthropic and litellm +# resolvers return ``AnthropicCredential | None``. Reusing the dataclass keeps +# the credential-injection code path uniform (header_name / header_value). +CredentialResolver = Callable[[], "AnthropicCredential | None"] + + +class UnknownUpstreamError(KeyError): + """Raised when ``UpstreamRegistry.get`` is called with an unregistered name.""" + + +class UpstreamRegistry: + """Per-upstream registry of (httpx.Client, credential_resolver) pairs. + + Each entry is created lazily on first ``get(name)``. The clients share the + same timeout / connection-pool characteristics as today's + ``_anthropic_client`` so behavior is byte-identical on the Anthropic path. + """ + + def __init__( + self, + litellm_base_url: str | None = None, + ) -> None: + self._litellm_base_url = litellm_base_url or os.environ.get( + "LITELLM_BASE_URL", LITELLM_BASE_URL_DEFAULT + ) + self._clients: dict[str, httpx.Client] = {} + self._resolvers: dict[str, CredentialResolver] = {} + self._lock = threading.Lock() + + def _make_client(self, base_url: str) -> httpx.Client: + """Create an httpx.Client with the same shape as today's singleton.""" + return httpx.Client( + base_url=base_url, # noqa: EGG200 - gateway proxy client, not direct LLM call + timeout=httpx.Timeout(120.0, connect=10.0), + limits=httpx.Limits(max_connections=100, max_keepalive_connections=20), + ) + + def _ensure_anthropic(self) -> None: + if "anthropic" in self._clients: + return + self._clients["anthropic"] = self._make_client(ANTHROPIC_BASE_URL) + # The anthropic resolver wraps the existing global + # AnthropicCredentialsManager so its mtime-invalidated cache is shared + # with any direct callers of get_credentials_manager(). + self._resolvers["anthropic"] = lambda: get_credentials_manager().get_credential() + + def _ensure_litellm(self) -> None: + if "litellm" in self._clients: + return + self._clients["litellm"] = self._make_client(self._litellm_base_url) + self._resolvers["litellm"] = lambda: get_litellm_credentials_manager().get_credential() + + def get(self, upstream: str) -> tuple[httpx.Client, CredentialResolver]: + """Return ``(client, credential_resolver)`` for ``upstream``. + + Raises ``UnknownUpstreamError`` if ``upstream`` is not a registered + name. Registration is implicit on first call for the canonical + upstreams (``anthropic`` and ``litellm``) — both share construction + semantics with today's ``get_anthropic_client()``. + """ + if upstream not in KNOWN_UPSTREAMS: + raise UnknownUpstreamError(upstream) + with self._lock: + if upstream == "anthropic": + self._ensure_anthropic() + elif upstream == "litellm": + self._ensure_litellm() + + return self._clients[upstream], self._resolvers[upstream] + + def is_known(self, upstream: str) -> bool: + """Return True if ``upstream`` is a name the registry will serve.""" + return upstream in KNOWN_UPSTREAMS + + def known_upstreams(self) -> tuple[str, ...]: + """Return the canonical upstream names the registry will serve.""" + return KNOWN_UPSTREAMS + + def close(self) -> None: + """Close all open httpx clients. For tests / teardown.""" + with self._lock: + for client in self._clients.values(): + try: + client.close() + except Exception: + logger.debug("Failed to close httpx client") + self._clients.clear() + self._resolvers.clear() + + +# Module-level singleton mirrors today's ``_anthropic_client`` lifetime. +_upstream_registry: UpstreamRegistry | None = None +_registry_lock = threading.Lock() + + +def get_upstream_registry() -> UpstreamRegistry: + """Return the module-level ``UpstreamRegistry`` singleton.""" + global _upstream_registry + if _upstream_registry is None: + with _registry_lock: + if _upstream_registry is None: + _upstream_registry = UpstreamRegistry() + return _upstream_registry + + +def reset_upstream_registry() -> None: + """Reset the module-level registry. For tests only.""" + global _upstream_registry + with _registry_lock: + if _upstream_registry is not None: + try: + _upstream_registry.close() + except Exception: + pass + _upstream_registry = None + + +__all__ = [ + "ANTHROPIC_BASE_URL", + "KNOWN_UPSTREAMS", + "LITELLM_BASE_URL_DEFAULT", + "CredentialResolver", + "LiteLLMCredentialsManager", + "UnknownUpstreamError", + "UpstreamRegistry", + "get_upstream_registry", + "reset_upstream_registry", +] diff --git a/k8s/base/kustomization.yaml b/k8s/base/kustomization.yaml index 45ebae656f..d3012c44da 100644 --- a/k8s/base/kustomization.yaml +++ b/k8s/base/kustomization.yaml @@ -8,4 +8,9 @@ resources: - orchestrator-service.yaml - gateway-deployment.yaml - gateway-service.yaml + # LiteLLM proxy (issue #2769) — no-op until operators populate the + # ConfigMap's ``model_list``; only the gateway pod calls into it. + - litellm-configmap.yaml + - litellm-deployment.yaml + - litellm-service.yaml - network-policies.yaml diff --git a/k8s/base/litellm-configmap.yaml b/k8s/base/litellm-configmap.yaml new file mode 100644 index 0000000000..01475ac750 --- /dev/null +++ b/k8s/base/litellm-configmap.yaml @@ -0,0 +1,41 @@ +apiVersion: v1 +kind: ConfigMap +metadata: + name: litellm-config + namespace: egg-system + labels: + app.kubernetes.io/name: litellm + app.kubernetes.io/component: litellm + app.kubernetes.io/part-of: egg +data: + config.yaml: | + # LiteLLM Proxy configuration (issue #2769). + # + # This ConfigMap ships with an EMPTY ``model_list`` so the Deployment + # comes up healthy but serves NO models until an operator populates + # entries here and overrides this ConfigMap via a kustomize overlay or + # ``kubectl apply``. Without entries, no agent request can route to + # LiteLLM (the gateway only forwards when a session's ``upstream`` is + # explicitly set to ``"litellm"``). + # + # Operators add backends per LiteLLM's documented format, e.g. for a + # hosted Qwen provider (cq-6): + # + # model_list: + # - model_name: qwen3-coder-30b + # litellm_params: + # model: together_ai/Qwen/Qwen2.5-Coder-32B-Instruct + # api_key: os.environ/TOGETHER_API_KEY + # + # The matching provider-side API key (e.g. ``TOGETHER_API_KEY``) is + # set via LiteLLM's standard env-var convention. It is NOT stored + # in gateway secrets.env — only the LiteLLM master key crosses that + # boundary (cq-7). + model_list: [] + + # Health-check config — used by the kubelet readiness/liveness + # probes on the Deployment. The Deployment probes /health/readiness + # and /health/liveliness; both pass with an empty model_list because + # LiteLLM treats "no models registered" as a healthy idle state. + general_settings: + master_key: os.environ/LITELLM_MASTER_KEY diff --git a/k8s/base/litellm-deployment.yaml b/k8s/base/litellm-deployment.yaml new file mode 100644 index 0000000000..91bd457920 --- /dev/null +++ b/k8s/base/litellm-deployment.yaml @@ -0,0 +1,109 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: litellm + namespace: egg-system + labels: + app.kubernetes.io/name: litellm + app.kubernetes.io/component: litellm + app.kubernetes.io/part-of: egg +spec: + # Single replica is sufficient — only the gateway pod is allowed to call + # this Service (no agent pods talk to it directly), so request volume is + # bounded by gateway concurrency. Issue #2769 cq-1. + replicas: 1 + selector: + matchLabels: + app.kubernetes.io/name: litellm + app.kubernetes.io/component: litellm + template: + metadata: + labels: + app.kubernetes.io/name: litellm + app.kubernetes.io/component: litellm + app.kubernetes.io/part-of: egg + spec: + enableServiceLinks: false + securityContext: + runAsNonRoot: true + runAsUser: 1000 + runAsGroup: 1000 + fsGroup: 1000 + containers: + - name: litellm + # Pinned image — bump via overlay when intentionally updating. + # LiteLLM's Anthropic-translation interface is the contract the + # gateway depends on; do not float ``latest`` (cq-1, feedback Q3). + image: ghcr.io/berriai/litellm:main-v1.55.10 + imagePullPolicy: IfNotPresent + args: + - --config + - /app/config.yaml + - --port + - "4000" + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + readOnlyRootFilesystem: true + ports: + - name: api + containerPort: 4000 + protocol: TCP + env: + # readOnlyRootFilesystem is true and only /tmp is writable. + # Point HOME at the writable tmpfs so any home-relative path + # LiteLLM touches at startup (cache / config dirs, and the + # XDG_* defaults that derive from $HOME) lands somewhere + # writable instead of failing against the read-only root. + - name: HOME + value: /tmp + # LiteLLM master key — clients (the gateway) inject this as + # ``x-api-key`` on every request. Issue #2769 cq-7. Operators + # populate the matching value in + # ``~/.config/egg/secrets.env`` as ``LITELLM_MASTER_KEY``. + # Sourced here from the existing gateway-secrets Secret so + # both sides of the wire share one value of truth. + - name: LITELLM_MASTER_KEY + valueFrom: + secretKeyRef: + name: gateway-secrets + key: litellm-master-key + optional: true + livenessProbe: + httpGet: + path: /health/liveliness + port: api + initialDelaySeconds: 15 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 6 + readinessProbe: + httpGet: + path: /health/readiness + port: api + initialDelaySeconds: 5 + periodSeconds: 5 + timeoutSeconds: 3 + failureThreshold: 6 + resources: + requests: + cpu: 100m + memory: 256Mi + limits: + cpu: "1" + memory: 1Gi + volumeMounts: + - name: config + mountPath: /app/config.yaml + subPath: config.yaml + readOnly: true + - name: tmp + mountPath: /tmp + volumes: + - name: config + configMap: + name: litellm-config + - name: tmp + emptyDir: {} diff --git a/k8s/base/litellm-service.yaml b/k8s/base/litellm-service.yaml new file mode 100644 index 0000000000..521988e31f --- /dev/null +++ b/k8s/base/litellm-service.yaml @@ -0,0 +1,23 @@ +apiVersion: v1 +kind: Service +metadata: + name: litellm + namespace: egg-system + labels: + app.kubernetes.io/name: litellm + app.kubernetes.io/component: litellm + app.kubernetes.io/part-of: egg +spec: + type: ClusterIP + selector: + app.kubernetes.io/name: litellm + app.kubernetes.io/component: litellm + ports: + # The gateway's UpstreamRegistry resolves to + # http://litellm.egg-system.svc.cluster.local:4000 by default + # (LITELLM_BASE_URL env override). No agent pods reach this Service — + # the gateway is the only client (issue #2769 cq-1). + - name: api + port: 4000 + targetPort: 4000 + protocol: TCP diff --git a/orchestrator/agent_model_resolution.py b/orchestrator/agent_model_resolution.py new file mode 100644 index 0000000000..196031b8d4 --- /dev/null +++ b/orchestrator/agent_model_resolution.py @@ -0,0 +1,196 @@ +"""Per-agent model resolution for the SDLC pipeline. + +This module owns the precedence and classification logic that decides, +for any given agent role, which model the agent should run on and which +upstream the gateway should route its ``/v1/messages`` traffic to. + +Precedence (highest first), matching #2769 task-2-3: + +1. ``PipelineConfig.agent_models[role]`` — per-pipeline override the + operator passes on submission. Keys are validated against + :class:`AgentRole` at construction time. +2. ``repositories.yaml`` ``default_agent_model`` — repository-level + default surfaced via :func:`config.repo_config.get_default_agent_model`. +3. Built-in ``"opus"`` default — preserves today's Claude-only behaviour. + +Classifier (cq-5 mitigation): a model string matching one of the +recognised Claude aliases (``opus``, ``opus[1m]``, ``sonnet``, +``sonnet[1m]``, ``haiku``, ``claude-*``) routes through the Anthropic +upstream and the agent's ``--model`` flag is set to that alias +verbatim. Any other string is treated as a LiteLLM-side model name: +the upstream is ``"litellm"``, the upstream-side model name is +preserved (the gateway rewrites the request body before forwarding — +see ``_rewrite_upstream_model`` in ``gateway/gateway.py``), and Claude +Code is handed the recognised alias ``"opus"`` so its compaction math +stays sane. + +The resolver is a pure function over its three inputs (role, +PipelineConfig, repo) so callers can use it from spawn, restart, and +test paths without further plumbing. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +from egg_contracts.agent_roles import AgentRole + +# Built-in fallback when neither PipelineConfig.agent_models nor the +# repository-level default_agent_model is set. Matches today's hardcoded +# default in ``orchestrator/consensus_wrapper.py::build_consensus_wrapped_command``. +DEFAULT_AGENT_MODEL = "opus" + +# Upstream identifiers used by the gateway's UpstreamRegistry +# (gateway/upstream_registry.py). +UPSTREAM_ANTHROPIC = "anthropic" +UPSTREAM_LITELLM = "litellm" + +# Claude alias presented to Claude Code when the resolved model is a +# non-Claude model routed through LiteLLM (cq-5 mitigation). +LITELLM_CLAUDE_CODE_ALIAS = "opus" + +# Recognised Claude aliases that route through the Anthropic upstream. +# Exact-match set plus a regex for the version-pinned ``claude-*`` family +# (e.g. ``claude-3-5-sonnet-20241022``, ``claude-opus-4-20250514``). +_CLAUDE_EXACT_ALIASES = frozenset( + { + "opus", + "opus[1m]", + "sonnet", + "sonnet[1m]", + "haiku", + } +) +_CLAUDE_VERSIONED_RE = re.compile(r"^claude-") + + +@dataclass(frozen=True) +class AgentModelDecision: + """Resolved per-agent model decision. + + Attributes: + claude_code_alias: The string passed to ``python3 -m egg_agent + --model`` inside the sandbox. For Anthropic-routed models + this is the resolved model name verbatim. For LiteLLM-routed + models this is always :data:`LITELLM_CLAUDE_CODE_ALIAS` (cq-5 + mitigation) so Claude Code's compaction heuristics stay + calibrated against a known Claude family. + upstream: One of :data:`UPSTREAM_ANTHROPIC` or + :data:`UPSTREAM_LITELLM`. The gateway's UpstreamRegistry + keys per-request ``httpx.Client`` + credential by this name. + upstream_model: The upstream-side model name to rewrite the + request body's ``model`` field to (gateway-side + ``_rewrite_upstream_model``). ``None`` on the Anthropic path + — the body is forwarded byte-for-byte unchanged. + """ + + claude_code_alias: str + upstream: str + upstream_model: str | None + + +def _is_claude_alias(model: str) -> bool: # noqa: EGG201 - docstring example shows versioned model ID format + """Return True when *model* is a recognised Claude family alias. + + Matches the explicit aliases the Claude Code harness understands + (``opus``, ``opus[1m]``, ``sonnet``, ``sonnet[1m]``, ``haiku``) + plus the versioned ``claude-*`` family (e.g. + ``claude-3-5-sonnet-20241022``). Used by the classifier in + :func:`resolve_agent_model` to pick the upstream. + """ + if model in _CLAUDE_EXACT_ALIASES: + return True + return bool(_CLAUDE_VERSIONED_RE.match(model)) + + +def classify_model(model: str) -> AgentModelDecision: + """Classify a raw model string into an :class:`AgentModelDecision`. + + Separated from :func:`resolve_agent_model` so callers that already + hold a resolved model string can reuse the classifier without + re-running precedence resolution. + """ + if _is_claude_alias(model): + return AgentModelDecision( + claude_code_alias=model, + upstream=UPSTREAM_ANTHROPIC, + upstream_model=None, + ) + return AgentModelDecision( + claude_code_alias=LITELLM_CLAUDE_CODE_ALIAS, + upstream=UPSTREAM_LITELLM, + upstream_model=model, + ) + + +def resolve_agent_model( + role: AgentRole | str, + pipeline_config: object | None, + repo: str | None, +) -> AgentModelDecision: + """Resolve the model decision for *role* per the precedence rules. + + Args: + role: The :class:`AgentRole` (or its raw value) being spawned. + pipeline_config: A ``PipelineConfig`` instance (typed loosely as + ``object`` to avoid an import cycle with ``orchestrator.models``). + ``None`` is treated as an empty config — the resolver falls + through to the repo-level default and then the built-in. + repo: Repository in ``owner/repo`` format, or ``None`` when the + caller has no repo context. Used to look up + ``default_agent_model`` from ``repositories.yaml``. + + Returns: + An :class:`AgentModelDecision` with the Claude-Code-facing alias, + the chosen upstream name, and the upstream-side model name (or + ``None`` on the Anthropic path). + """ + role_value = role.value if isinstance(role, AgentRole) else role + + # Tier 1: per-pipeline override. + if pipeline_config is not None: + agent_models = getattr(pipeline_config, "agent_models", None) + if isinstance(agent_models, dict): + override = agent_models.get(role_value) + if override: + return classify_model(override) + + # Tier 2: repository-level default. + if repo: + # Lazy import: ``config.repo_config`` reads from disk on first + # call and we want to defer that until a caller actually needs + # repo-level resolution. Also avoids pulling the config module + # into every test that exercises the classifier directly. + # + # Dual-import with fallback: the orchestrator Dockerfile flattens + # ``config/repo_config.py`` to ``/app/repo_config.py`` at the top + # level (``orchestrator/Dockerfile:66``), so the production + # container has no ``config/`` package — only the source-tree + # layout does. This mirrors the established pattern at + # ``shared/egg_restrictions/patterns.py:913-916`` and + # ``orchestrator/routes/signals.py:961-964``. + try: + from config.repo_config import get_default_agent_model + except ImportError: + from repo_config import ( # type: ignore[import-not-found, no-redef] + get_default_agent_model, + ) + + repo_default = get_default_agent_model(repo) + if repo_default: + return classify_model(repo_default) + + # Tier 3: built-in default. + return classify_model(DEFAULT_AGENT_MODEL) + + +__all__ = [ + "AgentModelDecision", + "DEFAULT_AGENT_MODEL", + "LITELLM_CLAUDE_CODE_ALIAS", + "UPSTREAM_ANTHROPIC", + "UPSTREAM_LITELLM", + "classify_model", + "resolve_agent_model", +] diff --git a/orchestrator/concurrent_executor.py b/orchestrator/concurrent_executor.py index 83861f7904..6b1bb74553 100644 --- a/orchestrator/concurrent_executor.py +++ b/orchestrator/concurrent_executor.py @@ -27,6 +27,13 @@ def get_logger(name: str, **kwargs) -> logging.Logger: # type: ignore[misc] return logging.getLogger(name) +from agent_model_resolution import ( + DEFAULT_AGENT_MODEL, + UPSTREAM_ANTHROPIC, + AgentModelDecision, + classify_model, + resolve_agent_model, +) from consensus_wrapper import build_consensus_wrapped_command from events import EventType, emit_event from message_store import Message, MessageType, get_message_store @@ -449,16 +456,53 @@ def _spawn_agent(self, role: AgentRole, prompt_text: str = "") -> AgentExecution branch = self.get_worktree_branch(role, slice_id=self._slice_id) env = self.get_agent_env(role) + # Per-agent model resolution (#2769 slice-2). The decision is a + # pure function over (role, pipeline_config, repo); when no + # override is configured the resolver returns the built-in + # ``opus`` Anthropic decision so the wire shape stays identical + # to the pre-#2769 path. + # + # Defensive wrap: a future regression in the resolver (e.g. a + # broken lazy import for the repo-default tier) would otherwise + # bring down agent spawn for every pipeline. Mirror the restart + # path's ``classify_model(DEFAULT_AGENT_MODEL)`` fallback at + # ``routes/pipelines.py:2683-2699`` so spawn degrades to the + # built-in opus / anthropic decision and logs the resolver + # failure rather than crashing. + try: + decision: AgentModelDecision = resolve_agent_model( + role=role, + pipeline_config=self.pipeline.config, + repo=self.pipeline.repo, + ) + except Exception as resolve_err: + logger.warning( + "Failed to resolve per-agent model decision for spawn, " + "falling back to built-in opus / anthropic default", + role=role, + error=str(resolve_err), + ) + decision = classify_model(DEFAULT_AGENT_MODEL) + command: list[str] | None = None if prompt_text: - command = build_consensus_wrapped_command(prompt_text) + command = build_consensus_wrapped_command(prompt_text, model=decision.claude_code_alias) + + # Forward the upstream/upstream_model kwargs to the spawner only + # when they would change behavior — the default Anthropic decision + # is omitted so test mocks and legacy spawn paths see the same + # call signature they did before #2769 slice-2 (#2769 task-2-4). + spawn_kwargs: dict[str, Any] = { + "role": role, + "branch": branch, + "extra_env": env, + "command": command, + } + if decision.upstream != UPSTREAM_ANTHROPIC or decision.upstream_model is not None: + spawn_kwargs["upstream"] = decision.upstream + spawn_kwargs["upstream_model"] = decision.upstream_model - result = self.spawn_fn( - role=role, - branch=branch, - extra_env=env, - command=command, - ) + result = self.spawn_fn(**spawn_kwargs) # container_id works for both Docker containers and k8s Jobs/pods. # The KubernetesClient returns the Job UID as container_id. diff --git a/orchestrator/gateway_client.py b/orchestrator/gateway_client.py index c252f4b934..85d5c87601 100644 --- a/orchestrator/gateway_client.py +++ b/orchestrator/gateway_client.py @@ -618,6 +618,8 @@ def register_session( worktree_container_id: str | None = None, jira_ticket: str | None = None, synthetic: bool = False, + upstream: str | None = None, + upstream_model: str | None = None, ) -> SessionInfo: """Register a session for a container. @@ -643,6 +645,20 @@ def register_session( call. When provided, the gateway reuses those worktrees instead of re-creating them — avoids a second ``git worktree add`` racing on ``.git/config.lock`` (#1857). + upstream: Optional per-session upstream selector — ``"anthropic"`` + (default behavior when omitted) or ``"litellm"`` to route the + session's ``/v1/messages`` traffic through the LiteLLM proxy + in egg-system (issue #2769). Omitted callers produce a + request body byte-identical to today. The gateway is + authoritative: it validates this against its + ``UpstreamRegistry`` and rejects an unknown value with + HTTP 400, so a slice-2 resolution bug fails fast at + session-create rather than producing a bogus-upstream + session. + upstream_model: Optional upstream-side model name used by the + slice-2 body-rewrite path on the gateway. Only meaningful + when ``upstream="litellm"``; the gateway leaves the request + body's ``model`` field untouched when this is omitted. Returns: SessionInfo with the created session @@ -688,6 +704,12 @@ def register_session( request_data["jira_ticket"] = jira_ticket if synthetic: request_data["synthetic"] = True + if upstream is not None: + # Only include when caller opts in — omitting the field keeps + # the wire shape byte-identical for pre-#2769 callers. + request_data["upstream"] = upstream + if upstream_model is not None: + request_data["upstream_model"] = upstream_model result = self._make_request( "/api/v1/sessions/create", method="POST", diff --git a/orchestrator/kubernetes_spawner.py b/orchestrator/kubernetes_spawner.py index e473b7920c..2d26c2fdd1 100644 --- a/orchestrator/kubernetes_spawner.py +++ b/orchestrator/kubernetes_spawner.py @@ -493,6 +493,8 @@ def spawn_agent_job( spawn_retry_initial_backoff_seconds: float = (DEFAULT_SPAWN_RETRY_INITIAL_BACKOFF_SECONDS), jira_ticket: str | None = None, slice_id: str | None = None, + upstream: str | None = None, + upstream_model: str | None = None, ) -> SpawnedContainer: """Spawn a Kubernetes Job for an agent. @@ -516,6 +518,13 @@ def spawn_agent_job( worktree-creation failures. ``0`` disables retry (#1839). spawn_retry_initial_backoff_seconds: Initial backoff between retries; subsequent attempts scale by ``_SPAWN_RETRY_BACKOFF_MULTIPLIER``. + upstream: Per-agent upstream identifier (#2769 slice-2). + Forwarded to the gateway session-create call only when + set; ``None`` keeps the default Anthropic routing. + upstream_model: Upstream-side model name to rewrite the + request body's ``model`` field to (#2769 slice-2). + ``None`` on the Anthropic path — the body is forwarded + unchanged. Returns: SpawnedContainer with Job and session info @@ -754,6 +763,12 @@ def spawn_agent_job( worktree_container_id=( agent_worktree_id if worktree_created_this_call else None ), + # Per-agent upstream routing (#2769 slice-2). Both fields + # are forwarded to the gateway only when set, so the + # default-Claude case keeps the request body byte- + # identical to the pre-#2769 wire shape. + upstream=upstream, + upstream_model=upstream_model, ) session_token = session_info.session_token @@ -1253,6 +1268,8 @@ def restart_agent_job( spawn_retry_initial_backoff_seconds: float = (DEFAULT_SPAWN_RETRY_INITIAL_BACKOFF_SECONDS), slice_id: str | None = None, wait_for_gateway: bool = True, + upstream: str | None = None, + upstream_model: str | None = None, ) -> SpawnedContainer: """Restart an agent Job: delete and respawn preserving worktree. @@ -1285,6 +1302,15 @@ def restart_agent_job( slice gets an independent budget. wait_for_gateway: Wait for gateway health before respawning. Forwarded to ``spawn_agent_job``. + upstream: Per-agent upstream identifier (#2769 slice-2), + forwarded to ``spawn_agent_job`` so the restarted Job + registers its gateway session against the same upstream + as the initial spawn. ``None`` keeps the default + Anthropic routing. + upstream_model: Upstream-side model name to rewrite the + request body's ``model`` field to (#2769 slice-2), + forwarded to ``spawn_agent_job``. ``None`` on the + Anthropic path — the body is forwarded unchanged. Returns: SpawnedContainer with new Job info. @@ -1412,6 +1438,13 @@ def restart_agent_job( spawn_max_retries=spawn_max_retries, spawn_retry_initial_backoff_seconds=spawn_retry_initial_backoff_seconds, slice_id=slice_id, + # Per-agent upstream routing (#2769 slice-2). Forwarded so a + # restart picks the same upstream as the initial spawn — the + # gateway session is otherwise rebuilt against the + # ``anthropic`` default and would silently route the + # restarted agent to the wrong upstream. + upstream=upstream, + upstream_model=upstream_model, ) logger.info( @@ -1684,6 +1717,8 @@ def _spawn( branch: str | None = None, extra_env: dict[str, str] | None = None, command: list[str] | None = None, + upstream: str | None = None, + upstream_model: str | None = None, ) -> SpawnedContainer: merged_env = {**(sandbox_env or {}), **(extra_env or {})} return self.spawn_agent_job( @@ -1702,6 +1737,8 @@ def _spawn( spawn_max_retries=spawn_max_retries, spawn_retry_initial_backoff_seconds=(spawn_retry_initial_backoff_seconds), slice_id=slice_id, + upstream=upstream, + upstream_model=upstream_model, ) return _spawn diff --git a/orchestrator/models.py b/orchestrator/models.py index 198a1ac65d..4ad890a49c 100644 --- a/orchestrator/models.py +++ b/orchestrator/models.py @@ -754,6 +754,56 @@ class PipelineConfig(BaseModel): "Subsequent attempts scale by 3x (e.g. 30s, 90s). See #1879." ), ) + agent_models: dict[str, str] = Field( + default_factory=dict, + description=( + "Per-role model overrides keyed by AgentRole value (e.g. " + "{'refiner': 'qwen3-coder-30b'}). Keys are restricted to the " + "SDLC phase producer and reviewer roles the resolver honors " + "(agent_roles.MODEL_OVERRIDE_ROLES) — utility/interface roles " + "such as overseer or autofixer are rejected at construction " + "time because their spawn paths never consult the resolver. " + "The value is the upstream-side model name: a Claude alias " + "(opus / opus[1m] / sonnet / sonnet[1m] / haiku / claude-*) " + "routes through the Anthropic upstream, anything else routes " + "through the in-cluster LiteLLM proxy with the recognised " + "alias 'opus' presented to Claude Code (cq-5 mitigation). When a " + "role is absent from this mapping the resolver falls back to the " + "repository-level default_agent_model setting and then to the " + "built-in 'opus' default. See #2769." + ), + ) + + @field_validator("agent_models") + @classmethod + def _validate_agent_models_roles(cls, v: dict[str, str]) -> dict[str, str]: + """Reject ``agent_models`` keys the per-agent model resolver never honors. + + ``resolve_agent_model`` is consulted only by the spawn/restart + paths that cover the SDLC phase producers and reviewers + (``MODEL_OVERRIDE_ROLES``). Utility roles (autofixer, + conflict_resolver) and interface roles (overseer, inspector) spawn + through paths that never call the resolver, so an override naming + one of them would be silently dropped at spawn. Rejecting both + typos and these unhonored-but-real roles at PipelineConfig + construction time surfaces the misconfiguration immediately + instead of letting it silently no-op. See #2769 task-2-1. + """ + if not v: + return v + # Lazy import to avoid a circular dependency with shared.egg_contracts + # when PipelineConfig is imported during package init. + from egg_contracts.agent_roles import MODEL_OVERRIDE_ROLES + + valid = {role.value for role in MODEL_OVERRIDE_ROLES} + invalid = sorted(role for role in v if role not in valid) + if invalid: + raise ValueError( + f"Invalid agent_models role keys: {invalid}. agent_models is " + f"honored only for SDLC phase producer and reviewer roles: " + f"{sorted(valid)}" + ) + return v @model_validator(mode="after") def _validate_post_consensus_budgets(self) -> PipelineConfig: diff --git a/orchestrator/routes/pipelines.py b/orchestrator/routes/pipelines.py index f93812523a..8d2c79c646 100644 --- a/orchestrator/routes/pipelines.py +++ b/orchestrator/routes/pipelines.py @@ -2661,6 +2661,49 @@ def restart_agent(pipeline_id: str, agent_role: str) -> tuple[Response, int]: # and role-specific environment variables to function properly. command = None extra_env: dict[str, str] = {} + + # Per-agent model resolution for the restart path (#2769 task-2-5). + # The default Anthropic decision keeps every restart kwarg byte- + # identical to the pre-#2769 wire shape; non-default decisions wire + # the right ``--model`` into the consensus wrapper and the right + # upstream/upstream_model into the gateway session. + try: + try: + from agent_model_resolution import ( + UPSTREAM_ANTHROPIC, + resolve_agent_model, + ) + except ImportError: + from ..agent_model_resolution import ( # type: ignore[import-not-found, no-redef] + UPSTREAM_ANTHROPIC, + resolve_agent_model, + ) + _model_decision = resolve_agent_model( + role=role, + pipeline_config=pipeline.config, + repo=pipeline.repo, + ) + except Exception as resolve_err: + # Resolution is pure data over already-validated inputs; logging + # and falling back to the built-in opus default preserves restart + # availability even if a future regression breaks the resolver. + logger.warning( + "Failed to resolve per-agent model decision for restart, " + "falling back to built-in opus default", + pipeline_id=pipeline_id, + agent_role=agent_role, + error=str(resolve_err), + ) + try: + from agent_model_resolution import UPSTREAM_ANTHROPIC, classify_model + except ImportError: + from ..agent_model_resolution import ( # type: ignore[import-not-found, no-redef] + UPSTREAM_ANTHROPIC, + classify_model, + ) + + _model_decision = classify_model("opus") + try: try: from concurrent_executor import ConcurrentPhaseExecutor, is_concurrent_execution @@ -2701,7 +2744,9 @@ def restart_agent(pipeline_id: str, agent_role: str) -> tuple[Response, int]: if prompt_text: from consensus_wrapper import build_consensus_wrapped_command - command = build_consensus_wrapped_command(prompt_text) + command = build_consensus_wrapped_command( + prompt_text, model=_model_decision.claude_code_alias + ) except Exception as prompt_err: logger.warning( "Failed to reconstruct agent prompt for restart " @@ -2720,6 +2765,14 @@ def restart_agent(pipeline_id: str, agent_role: str) -> tuple[Response, int]: error=str(e), ) + # Forward upstream/upstream_model only on non-default decisions so a + # restart on the Claude default keeps the spawner kwargs byte- + # identical to today's pre-#2769 shape (regression guard). + restart_upstream_kwargs: dict[str, str | None] = {} + if _model_decision.upstream != UPSTREAM_ANTHROPIC or _model_decision.upstream_model is not None: + restart_upstream_kwargs["upstream"] = _model_decision.upstream + restart_upstream_kwargs["upstream_model"] = _model_decision.upstream_model + try: spawned = spawner.restart_agent_container( pipeline_id=pipeline_id, @@ -2736,6 +2789,7 @@ def restart_agent(pipeline_id: str, agent_role: str) -> tuple[Response, int]: spawn_max_retries=pipeline.config.spawn_max_retries, spawn_retry_initial_backoff_seconds=pipeline.config.spawn_retry_initial_backoff_seconds, slice_id=slice_id, + **restart_upstream_kwargs, ) except (ContainerSpawnError, KubernetesSpawnError) as e: # Revert early status update — the agent is not actually running. diff --git a/orchestrator/tests/test_agent_model_resolution.py b/orchestrator/tests/test_agent_model_resolution.py new file mode 100644 index 0000000000..17f46420fe --- /dev/null +++ b/orchestrator/tests/test_agent_model_resolution.py @@ -0,0 +1,572 @@ +"""Tests for ``orchestrator/agent_model_resolution.py`` — slice-2 of #2769. + +The resolver decides, for a given agent role, which Claude-Code-facing +model alias to pass to ``build_consensus_wrapped_command``, which +upstream to register on the gateway session, and which upstream-side +model name (if any) to put in ``session.upstream_model``. + +Precedence (highest wins): + + 1. ``pipeline_config.agent_models.get(role.value)`` — per-pipeline + 2. ``get_default_agent_model(repo)`` — per-repo + 3. ``"opus"`` built-in default + +Classifier (model string → upstream): + + * ``"opus"``, ``"opus[1m]"``, ``"sonnet"``, ``"sonnet[1m]"``, + ``"haiku"``, ``"claude-*"`` → ``upstream="anthropic"``, + ``claude_code_alias=``, ``upstream_model=None``. + * anything else → ``upstream="litellm"``, + ``claude_code_alias="opus"`` (cq-5 mitigation: Claude Code keeps + seeing a recognized alias so its compaction math stays sane), + ``upstream_model=``. + +Plan reference: ``.egg-state/drafts/2769-plan.md`` TASK-2-3 / TASK-2-7. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +# Add orchestrator to sys.path the same way test_concurrent_executor.py does. +_orchestrator_path = Path(__file__).parent.parent +if str(_orchestrator_path) not in sys.path: + sys.path.insert(0, str(_orchestrator_path)) + + +def _resolver(): + """Return the slice-2 ``resolve_agent_model`` symbol or skip the test. + + Slice-2 has not landed yet from the coder's side when the tester + scaffolds — keep the suite green until the symbol exists. + """ + try: + from agent_model_resolution import resolve_agent_model # type: ignore[import-not-found] + + return resolve_agent_model + except ImportError: + pytest.skip( + "agent_model_resolution.resolve_agent_model not yet implemented " + "(waiting on coder for slice-2)" + ) + + +def _decision_cls(): + """Return the slice-2 ``AgentModelDecision`` symbol or skip.""" + try: + from agent_model_resolution import AgentModelDecision # type: ignore[import-not-found] + + return AgentModelDecision + except ImportError: + pytest.skip( + "agent_model_resolution.AgentModelDecision not yet implemented " + "(waiting on coder for slice-2)" + ) + + +def _agent_role(): + """Return the ``AgentRole`` enum (canonical source: ``egg_contracts.agent_roles``).""" + from egg_contracts.agent_roles import AgentRole + + return AgentRole + + +def _pipeline_config(**overrides): + """Build a ``PipelineConfig`` with the given overrides. + + Mirrors the helper pattern at ``test_concurrent_executor.py:20`` but + keeps the slice-2 ``agent_models`` field exposed. + """ + from models import PipelineConfig + + return PipelineConfig(**overrides) + + +# ============================================================================= +# AgentModelDecision dataclass shape +# ============================================================================= + + +class TestAgentModelDecisionShape: + """The decision triple is the resolver's contract. + + Downstream callers (``concurrent_executor`` / ``routes.pipelines``) + unpack it into ``--model`` (Claude-Code-facing) and the + ``register_session(upstream=..., upstream_model=...)`` payload. + """ + + def test_decision_has_three_named_fields(self): + decision_cls = _decision_cls() + d = decision_cls( + claude_code_alias="opus", + upstream="anthropic", + upstream_model=None, + ) + assert d.claude_code_alias == "opus" + assert d.upstream == "anthropic" + assert d.upstream_model is None + + def test_decision_accepts_upstream_model_string(self): + decision_cls = _decision_cls() + d = decision_cls( + claude_code_alias="opus", + upstream="litellm", + upstream_model="qwen3-coder-30b", + ) + assert d.upstream_model == "qwen3-coder-30b" + + +# ============================================================================= +# Precedence rules: pipeline > repo > built-in +# ============================================================================= + + +class TestResolutionPrecedence: + """Per-pipeline ``agent_models`` wins; then repo default; then ``opus``.""" + + def test_builtin_default_is_opus_when_no_override(self): + """No pipeline override, no repo default → ``opus`` (Anthropic).""" + resolve_agent_model = _resolver() + AgentRole = _agent_role() + config = _pipeline_config() + + with patch("config.repo_config.get_default_agent_model", return_value=None): + d = resolve_agent_model(AgentRole.CODER, config, None) + + assert d.claude_code_alias == "opus" + assert d.upstream == "anthropic" + assert d.upstream_model is None + + def test_repo_default_used_when_no_pipeline_override(self): + """Pipeline empty, repo sets ``sonnet`` → ``sonnet`` (Anthropic).""" + resolve_agent_model = _resolver() + AgentRole = _agent_role() + config = _pipeline_config() + + with patch( + "config.repo_config.get_default_agent_model", + return_value="sonnet", + ): + d = resolve_agent_model(AgentRole.CODER, config, "owner/repo") + + assert d.claude_code_alias == "sonnet" + assert d.upstream == "anthropic" + assert d.upstream_model is None + + def test_pipeline_override_beats_repo_default(self): + """Pipeline ``refiner=qwen3-coder-30b`` wins even if repo says ``sonnet``.""" + resolve_agent_model = _resolver() + AgentRole = _agent_role() + config = _pipeline_config(agent_models={"refiner": "qwen3-coder-30b"}) + + with patch( + "config.repo_config.get_default_agent_model", + return_value="sonnet", + ): + d = resolve_agent_model(AgentRole.REFINER, config, "owner/repo") + + # Pipeline override wins: + assert d.upstream == "litellm" + assert d.upstream_model == "qwen3-coder-30b" + assert d.claude_code_alias == "opus" # cq-5 mitigation + + def test_pipeline_override_does_not_leak_to_other_roles(self): + """``agent_models={"refiner": "qwen3-coder-30b"}`` MUST NOT change + the coder's resolution — slice-2 is per-role, not per-pipeline-wide. + """ + resolve_agent_model = _resolver() + AgentRole = _agent_role() + config = _pipeline_config(agent_models={"refiner": "qwen3-coder-30b"}) + + with patch("config.repo_config.get_default_agent_model", return_value=None): + coder_decision = resolve_agent_model(AgentRole.CODER, config, None) + + # Coder is not overridden → built-in default. + assert coder_decision.upstream == "anthropic" + assert coder_decision.claude_code_alias == "opus" + assert coder_decision.upstream_model is None + + def test_none_repo_skips_repo_lookup(self): + """``repo is None`` MUST NOT raise — overseer / unsliced callers + sometimes pass ``None`` here (no per-repo default to consult). + """ + resolve_agent_model = _resolver() + AgentRole = _agent_role() + config = _pipeline_config() + + # If the resolver consults the repo lookup it should be guarded; + # the simplest correct implementation skips the call entirely when + # ``repo is None``. We patch defensively in either case. + with patch( + "config.repo_config.get_default_agent_model", + return_value=None, + ): + d = resolve_agent_model(AgentRole.CODER, config, None) + + assert d.claude_code_alias == "opus" + assert d.upstream == "anthropic" + + def test_non_string_repo_default_raises_value_error(self): + """A non-string ``default_agent_model`` in ``repositories.yaml`` + MUST raise a clear ``ValueError`` rather than degrade silently. + + Without the ``isinstance`` guard in ``get_default_agent_model`` a + non-string (e.g. ``default_agent_model: 4``) reaches + ``classify_model``, where the ``claude-*`` regex raises an opaque + ``TypeError`` from its internals — which the spawn-path + ``except Exception`` then swallows into the opus fallback with no + actionable log line. The guard converts that into an explicit, + operator-readable error naming the bad value. + """ + from config.repo_config import get_default_agent_model + + with patch("config.repo_config.get_repo_setting", return_value=4): + with pytest.raises(ValueError, match="must be a string"): + get_default_agent_model("owner/repo") + + +# ============================================================================= +# Classifier: Anthropic vs LiteLLM dispatch by model name +# ============================================================================= + + +class TestAnthropicClassification: + """Known Claude aliases → ``upstream="anthropic"``, + ``claude_code_alias=``, ``upstream_model=None``. + + The Anthropic path preserves the original alias on ``--model`` so + today's Claude routing is byte-identical to a default-config + pipeline. + """ + + @pytest.mark.parametrize( + "model", + ["opus", "opus[1m]", "sonnet", "sonnet[1m]", "haiku"], + ) + def test_short_claude_alias_is_anthropic(self, model): + resolve_agent_model = _resolver() + AgentRole = _agent_role() + config = _pipeline_config(agent_models={"coder": model}) + + with patch("config.repo_config.get_default_agent_model", return_value=None): + d = resolve_agent_model(AgentRole.CODER, config, None) + + assert d.upstream == "anthropic", f"alias {model!r} should be anthropic" + assert d.claude_code_alias == model, ( + f"alias should pass through, got {d.claude_code_alias!r}" + ) + assert d.upstream_model is None + + @pytest.mark.parametrize( + "model", + [ + "claude-3-5-sonnet-20241022", + "claude-3-opus-20240229", + "claude-sonnet-4-5", + "claude-haiku-4-5-20251001", + ], + ) + def test_claude_prefixed_full_name_is_anthropic(self, model): + resolve_agent_model = _resolver() + AgentRole = _agent_role() + config = _pipeline_config(agent_models={"coder": model}) + + with patch("config.repo_config.get_default_agent_model", return_value=None): + d = resolve_agent_model(AgentRole.CODER, config, None) + + assert d.upstream == "anthropic" + assert d.claude_code_alias == model + assert d.upstream_model is None + + +class TestLiteLLMClassification: + """Anything that does not look like a Claude alias → LiteLLM. + + On the LiteLLM path the cq-5 mitigation pins + ``claude_code_alias="opus"`` regardless of the upstream model + name — Claude Code's compaction math is name-derived, so it + must keep seeing a recognized alias. + """ + + @pytest.mark.parametrize( + "model", + [ + "qwen3-coder-30b", + "qwen2.5-72b-instruct", + "qwen-max", + "llama-3-70b-instruct", + "mistral-large-2", + "deepseek-v3", + "gpt-4o", + ], + ) + def test_non_claude_model_routes_to_litellm(self, model): + resolve_agent_model = _resolver() + AgentRole = _agent_role() + config = _pipeline_config(agent_models={"coder": model}) + + with patch("config.repo_config.get_default_agent_model", return_value=None): + d = resolve_agent_model(AgentRole.CODER, config, None) + + assert d.upstream == "litellm", f"{model!r} should route to litellm" + # cq-5 mitigation: Claude Code MUST keep seeing a recognized alias. + assert d.claude_code_alias == "opus", ( + f"LiteLLM-routed agents MUST present 'opus' to Claude Code so " + f"compaction math stays sane (cq-5); got {d.claude_code_alias!r}" + ) + assert d.upstream_model == model + + def test_litellm_alias_pin_is_opus_not_upstream_model(self): + """Adversarial guard for the cq-5 mitigation: the + Claude-Code-facing alias for a LiteLLM-routed agent is + ALWAYS ``"opus"``, never the upstream model name — even + when the upstream model is something like ``"opus-7b"`` + whose substring matches the Claude alias. + """ + resolve_agent_model = _resolver() + AgentRole = _agent_role() + + for tricky_model in ("opus-7b", "claude-clone-12b", "claudia-9000"): + config = _pipeline_config(agent_models={"coder": tricky_model}) + + with patch("config.repo_config.get_default_agent_model", return_value=None): + d = resolve_agent_model(AgentRole.CODER, config, None) + + # Either the classifier correctly routes to anthropic + # (only if the name matches the documented exact aliases), + # or it routes to litellm with ``claude_code_alias="opus"``. + # What MUST NOT happen: routing to litellm with a non-opus + # alias, because Claude Code would then receive an unknown + # model name in its --model flag. + if d.upstream == "litellm": + assert d.claude_code_alias == "opus", ( + f"litellm route MUST pin claude_code_alias='opus' for " + f"cq-5, got {d.claude_code_alias!r} for " + f"upstream_model={tricky_model!r}" + ) + + +# ============================================================================= +# PipelineConfig.agent_models validation (TASK-2-1) +# ============================================================================= + + +def _agent_models_field_exists() -> bool: + """Return True if ``PipelineConfig`` exposes the slice-2 + ``agent_models`` field — slice-2 may not have landed yet. + """ + from models import PipelineConfig + + return "agent_models" in PipelineConfig.model_fields + + +class TestAgentModelsValidation: + """``PipelineConfig.agent_models`` keys MUST be valid ``AgentRole`` + values; values are free-form strings validated downstream. + """ + + def test_default_is_empty_dict(self): + """No behavioral change for existing pipelines — slice-2 + regression guard. + """ + if not _agent_models_field_exists(): + pytest.skip("PipelineConfig.agent_models not yet implemented") + config = _pipeline_config() + assert config.agent_models == {}, ( + f"Default agent_models MUST be an empty dict (regression " + f"guard); got {config.agent_models!r}" + ) + + def test_known_role_constructs_successfully(self): + if not _agent_models_field_exists(): + pytest.skip("PipelineConfig.agent_models not yet implemented") + config = _pipeline_config(agent_models={"refiner": "qwen3-coder-30b"}) + assert config.agent_models == {"refiner": "qwen3-coder-30b"} + + def test_unknown_role_raises_validation_error(self): + """Pydantic validator MUST reject unknown roles at construction + time — names that drift from the canonical ``AgentRole`` enum + would silently fail at spawn time otherwise. + """ + if not _agent_models_field_exists(): + pytest.skip("PipelineConfig.agent_models not yet implemented") + from pydantic import ValidationError + + with pytest.raises(ValidationError) as excinfo: + _pipeline_config(agent_models={"bogus_role": "qwen3-coder-30b"}) + + # Error message MUST name the offending role so operators can + # debug from logs alone. + assert "bogus_role" in str(excinfo.value), ( + f"ValidationError MUST cite the unknown role; got: {excinfo.value}" + ) + + def test_unhonored_real_role_raises_validation_error(self): + """Roles that exist in ``AgentRole`` but are never threaded + through ``resolve_agent_model`` (overseer, autofixer, + conflict_resolver, inspector) MUST be rejected — accepting them + would let a deliberate override silently no-op at spawn, which is + exactly the silent-ignore trap the validator exists to prevent. + """ + if not _agent_models_field_exists(): + pytest.skip("PipelineConfig.agent_models not yet implemented") + from pydantic import ValidationError + + for unhonored in ("overseer", "autofixer", "conflict_resolver", "inspector"): + with pytest.raises(ValidationError) as excinfo: + _pipeline_config(agent_models={unhonored: "qwen3-coder-30b"}) + assert unhonored in str(excinfo.value), ( + f"ValidationError MUST cite the unhonored role {unhonored!r}; got: {excinfo.value}" + ) + + def test_multiple_known_roles_accepted(self): + """Many roles can be overridden simultaneously. + + ``applier`` (the producer of the ``apply`` phase, threaded + through ``resolve_agent_model`` via ``_PHASE_ROLES["apply"]``) is + included deliberately: it is an honored role and must be accepted + as a key, even though it is easy to mistake for an unhonored + utility role. + """ + if not _agent_models_field_exists(): + pytest.skip("PipelineConfig.agent_models not yet implemented") + config = _pipeline_config( + agent_models={ + "refiner": "qwen3-coder-30b", + "coder": "claude-3-5-sonnet-20241022", + "tester": "sonnet", + "applier": "haiku", + } + ) + assert config.agent_models["refiner"] == "qwen3-coder-30b" + assert config.agent_models["coder"] == "claude-3-5-sonnet-20241022" + assert config.agent_models["tester"] == "sonnet" + assert config.agent_models["applier"] == "haiku" + + def test_one_bad_role_in_a_mix_still_fails(self): + """Mixed dict — one bad role kills the whole construction.""" + if not _agent_models_field_exists(): + pytest.skip("PipelineConfig.agent_models not yet implemented") + from pydantic import ValidationError + + with pytest.raises(ValidationError): + _pipeline_config( + agent_models={ + "refiner": "qwen3-coder-30b", + "not_a_real_role_xyz": "x", + } + ) + + +# ============================================================================= +# Regression: default config produces byte-identical Anthropic decision +# ============================================================================= + + +class TestDualImportRepoConfig: + """Regression for the prod-container import topology (slice-2 v2). + + `orchestrator/Dockerfile:66` flattens ``config/repo_config.py`` to + ``/app/repo_config.py`` (no ``/app/config/`` package). The resolver + must therefore fall back to a top-level ``repo_config`` import when + the ``config.`` package is unavailable. This regression test asserts + the fallback works — without it, every Anthropic-default spawn in + production would have raised ``ModuleNotFoundError`` on the first + pipeline submit. + + Mirrors the dual-import pattern already used at + ``shared/egg_restrictions/patterns.py:913-916`` and + ``orchestrator/routes/signals.py:961-964``. + """ + + def test_top_level_repo_config_fallback_resolves(self, monkeypatch): + """Simulate the prod-container layout: ``config`` package absent, + ``repo_config`` available at the top level. The resolver MUST + still produce the built-in opus / anthropic decision instead of + propagating ``ModuleNotFoundError``. + """ + if not _agent_models_field_exists(): + pytest.skip("PipelineConfig.agent_models not yet implemented") + + import types + + resolve_agent_model = _resolver() + AgentRole = _agent_role() + config = _pipeline_config() + + # Build a top-level ``repo_config`` shim whose + # ``get_default_agent_model`` returns None (no override). + repo_config_shim = types.ModuleType("repo_config") + repo_config_shim.get_default_agent_model = lambda repo: None # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "repo_config", repo_config_shim) + + # Remove the ``config.repo_config`` module so the primary import + # raises ImportError and the fallback fires. Need to remove + # ``config`` itself too so ``from config.repo_config import ...`` + # actually raises rather than re-importing the existing module. + for name in ("config.repo_config", "config"): + sys.modules.pop(name, None) + + # Block import of the ``config`` package by inserting a meta-path + # finder that raises ImportError specifically for it — leaves + # other imports alone. + import importlib.abc + import importlib.machinery + + class _BlockConfigFinder(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "config" or fullname.startswith("config."): + raise ImportError(f"Simulating prod-container layout: {fullname} unavailable") + return None + + finder = _BlockConfigFinder() + sys.meta_path.insert(0, finder) + try: + d = resolve_agent_model(AgentRole.CODER, config, "owner/repo") + finally: + sys.meta_path.remove(finder) + + # Built-in opus / anthropic falls through. The shim returned None + # so neither pipeline override (empty) nor repo default fired. + assert d.claude_code_alias == "opus" + assert d.upstream == "anthropic" + assert d.upstream_model is None + + +class TestDefaultPathRegression: + """With ``agent_models={}`` and no repo default, EVERY role resolves + to the Anthropic default — this is the slice-2 no-op invariant. + """ + + def test_every_assigned_role_defaults_to_anthropic_opus(self): + resolve_agent_model = _resolver() + AgentRole = _agent_role() + config = _pipeline_config() + + with patch("config.repo_config.get_default_agent_model", return_value=None): + for role in ( + AgentRole.CODER, + AgentRole.TESTER, + AgentRole.DOCUMENTER, + AgentRole.REFINER, + AgentRole.ARCHITECT, + AgentRole.TASK_PLANNER, + AgentRole.RISK_ANALYST, + AgentRole.REVIEWER_CODE, + AgentRole.REVIEWER_CONTRACT, + AgentRole.REVIEWER_CODE_HOLISTIC, + AgentRole.REVIEWER_SECURITY, + AgentRole.REVIEWER_CONCURRENCY, + ): + d = resolve_agent_model(role, config, "any/repo") + assert d.claude_code_alias == "opus" + assert d.upstream == "anthropic" + assert d.upstream_model is None, ( + f"Default path for {role.value} produced " + f"upstream_model={d.upstream_model!r} — slice-2 " + f"regression: default config must be Anthropic-only" + ) diff --git a/orchestrator/tests/test_concurrent_executor.py b/orchestrator/tests/test_concurrent_executor.py index 041cde68a1..c36bf9235e 100644 --- a/orchestrator/tests/test_concurrent_executor.py +++ b/orchestrator/tests/test_concurrent_executor.py @@ -9,6 +9,8 @@ from pathlib import Path from unittest.mock import MagicMock, patch +import pytest # noqa: F401 # used in slice-2 (#2769) test additions below + # Add orchestrator to path _orchestrator_path = Path(__file__).parent.parent if str(_orchestrator_path) not in sys.path: @@ -933,3 +935,422 @@ def _evaluate_incomplete(): assert result["fallback"] == "message_bus" finally: remove_peer_consensus_tracker("KORE-1234") + + +# ============================================================================= +# Per-agent model wiring — slice-2 of issue #2769 (TASK-2-4 / TASK-2-7) +# ============================================================================= +# +# Slice 2 threads ``resolve_agent_model(role, pipeline.config, pipeline.repo)`` +# through the spawn path so: +# +# - ``build_consensus_wrapped_command`` receives the resolved +# Claude-Code-facing alias as ``model=`` (e.g. always ``"opus"`` on +# LiteLLM-routed agents, per cq-5). +# - ``spawn_fn`` (which dispatches to ``KubernetesSpawner.spawn_agent`` +# → ``GatewayClient.register_session``) receives ``upstream=`` and +# ``upstream_model=`` so the gateway session is registered with the +# right per-agent routing decision. +# +# Default-config pipelines (``agent_models == {}``) MUST still hit the +# old call shape — no new register_session kwargs, no change to the +# ``--model`` flag. That's the slice-2 no-op invariant. +# ============================================================================= + + +def _kubernetes_spawn_result(role_value: str = "coder"): + """Build a minimal ``SpawnedContainer`` so ``_spawn_agent`` is happy.""" + from kubernetes_spawner import SpawnedContainer + from models import ContainerInfo, ContainerStatus + + info = ContainerInfo( + container_id=f"uid-{role_value}", + container_name=f"issue-999-{role_value}", + status=ContainerStatus.PENDING, + namespace="egg-sandbox", + job_name=f"issue-999-{role_value}", + ) + return SpawnedContainer( + container_info=info, + session_info=None, + agent_role=None, + pipeline_id="issue-999", + environment={}, + ) + + +def _slice_2_available() -> bool: + """Return True if the slice-2 resolver has landed. + + Tests below skip when False — the coder hasn't pushed yet. + """ + try: + import agent_model_resolution # type: ignore[import-not-found] # noqa: F401 + + return True + except ImportError: + return False + + +class TestSpawnDefaultAgentModelPath: + """Regression guard: default-config pipelines spawn EXACTLY as + before slice-2 — same ``build_consensus_wrapped_command`` args, no + new ``register_session`` kwargs.""" + + def test_default_config_passes_opus_to_consensus_wrapper(self): + """``agent_models == {}`` → coder spawn passes ``model="opus"`` + (or whatever today's built-in default is) to + ``build_consensus_wrapped_command``. + """ + from concurrent_executor import ConcurrentPhaseExecutor + from egg_orchestrator.types import AgentRole + + pipeline = _make_pipeline() + # No agent_models, no repo default → built-in path. + captured: dict[str, object] = {} + + def _capture_command(prompt_text, **kwargs): + captured["prompt_text"] = prompt_text + # No default — the test must distinguish "model='opus' was + # passed" from "model was not passed at all". + captured["model"] = kwargs.get("model") + return ["bash", "-c", "true"] + + mock_spawn = MagicMock(return_value=_kubernetes_spawn_result()) + executor = ConcurrentPhaseExecutor(pipeline, spawn_fn=mock_spawn) + + with patch( + "concurrent_executor.build_consensus_wrapped_command", + side_effect=_capture_command, + ): + executor._spawn_agent(AgentRole.CODER, prompt_text="run task") + + assert captured.get("model") == "opus", ( + f"Default-config spawn MUST pass model='opus' to the " + f"consensus wrapper; got {captured.get('model')!r}" + ) + + def test_default_config_omits_upstream_kwargs_on_spawn_fn(self): + """The spawn_fn must NOT receive upstream/upstream_model kwargs + when ``agent_models == {}`` — that would be a wire-shape change + on the default Claude path (regression guard). + """ + from concurrent_executor import ConcurrentPhaseExecutor + from egg_orchestrator.types import AgentRole + + pipeline = _make_pipeline() + mock_spawn = MagicMock(return_value=_kubernetes_spawn_result()) + executor = ConcurrentPhaseExecutor(pipeline, spawn_fn=mock_spawn) + + with patch( + "concurrent_executor.build_consensus_wrapped_command", + return_value=["bash", "-c", "true"], + ): + executor._spawn_agent(AgentRole.CODER, prompt_text="run task") + + # spawn_fn was called once + assert mock_spawn.call_count == 1, mock_spawn.call_args_list + _args, kwargs = mock_spawn.call_args + + # Slice-2 invariant: when no override is configured, the + # default-Anthropic case omits the new kwargs entirely OR + # passes them as None — either is acceptable wire-shape; + # what's NOT acceptable is sending ``upstream="litellm"`` or + # a non-None ``upstream_model``. + assert kwargs.get("upstream") in (None, "anthropic"), ( + f"Default config should not send upstream='litellm'; got {kwargs.get('upstream')!r}" + ) + assert kwargs.get("upstream_model") is None, ( + f"Default config MUST NOT send upstream_model; got {kwargs.get('upstream_model')!r}" + ) + + +class TestSpawnLiteLLMConfiguredPath: + """``agent_models={"refiner": "qwen3-coder-30b"}`` MUST: + + - Pass ``model="opus"`` (cq-5 mitigation) to the consensus wrapper. + - Pass ``upstream="litellm"`` and + ``upstream_model="qwen3-coder-30b"`` to the spawn_fn. + + Other roles (e.g. coder) in the same pipeline MUST keep the + default Anthropic shape — the override is per-role. + """ + + def test_refiner_override_passes_opus_to_consensus_wrapper(self): + if not _slice_2_available(): + pytest.skip("agent_model_resolution not yet implemented") + + from concurrent_executor import ConcurrentPhaseExecutor + from egg_orchestrator.types import AgentRole + + pipeline = _make_pipeline() + try: + pipeline.config.agent_models = {"refiner": "qwen3-coder-30b"} + except AttributeError, ValueError: + pipeline.config.__dict__["agent_models"] = {"refiner": "qwen3-coder-30b"} + + captured: dict[str, object] = {} + + def _capture_command(prompt_text, **kwargs): + captured["model"] = kwargs.get("model") + return ["bash", "-c", "true"] + + mock_spawn = MagicMock(return_value=_kubernetes_spawn_result("refiner")) + executor = ConcurrentPhaseExecutor(pipeline, spawn_fn=mock_spawn) + + with patch( + "concurrent_executor.build_consensus_wrapped_command", + side_effect=_capture_command, + ): + executor._spawn_agent(AgentRole.REFINER, prompt_text="run task") + + # cq-5: Claude Code's ``--model`` flag MUST be a recognized + # Claude alias, NEVER the LiteLLM-side upstream model name. + assert captured.get("model") == "opus", ( + f"LiteLLM-routed refiner MUST present model='opus' to " + f"Claude Code (cq-5); got {captured.get('model')!r}" + ) + + def test_refiner_override_passes_upstream_to_spawn_fn(self): + if not _slice_2_available(): + pytest.skip("agent_model_resolution not yet implemented") + + from concurrent_executor import ConcurrentPhaseExecutor + from egg_orchestrator.types import AgentRole + + pipeline = _make_pipeline() + try: + pipeline.config.agent_models = {"refiner": "qwen3-coder-30b"} + except AttributeError, ValueError: + pipeline.config.__dict__["agent_models"] = {"refiner": "qwen3-coder-30b"} + + mock_spawn = MagicMock(return_value=_kubernetes_spawn_result("refiner")) + executor = ConcurrentPhaseExecutor(pipeline, spawn_fn=mock_spawn) + + with patch( + "concurrent_executor.build_consensus_wrapped_command", + return_value=["bash", "-c", "true"], + ): + executor._spawn_agent(AgentRole.REFINER, prompt_text="run task") + + _args, kwargs = mock_spawn.call_args + assert kwargs.get("upstream") == "litellm", ( + f"refiner spawn_fn MUST receive upstream='litellm'; got {kwargs.get('upstream')!r}" + ) + assert kwargs.get("upstream_model") == "qwen3-coder-30b", ( + f"refiner spawn_fn MUST receive upstream_model='qwen3-coder-30b'; " + f"got {kwargs.get('upstream_model')!r}" + ) + + def test_refiner_override_does_not_affect_coder_spawn(self): + """Per-role override — coder spawn MUST stay on the default + Anthropic shape even when refiner is overridden. + """ + if not _slice_2_available(): + pytest.skip("agent_model_resolution not yet implemented") + + from concurrent_executor import ConcurrentPhaseExecutor + from egg_orchestrator.types import AgentRole + + pipeline = _make_pipeline() + try: + pipeline.config.agent_models = {"refiner": "qwen3-coder-30b"} + except AttributeError, ValueError: + pipeline.config.__dict__["agent_models"] = {"refiner": "qwen3-coder-30b"} + + captured: dict[str, object] = {} + + def _capture_command(prompt_text, **kwargs): + captured["model"] = kwargs.get("model") + return ["bash", "-c", "true"] + + mock_spawn = MagicMock(return_value=_kubernetes_spawn_result("coder")) + executor = ConcurrentPhaseExecutor(pipeline, spawn_fn=mock_spawn) + + with patch( + "concurrent_executor.build_consensus_wrapped_command", + side_effect=_capture_command, + ): + executor._spawn_agent(AgentRole.CODER, prompt_text="run task") + + # Coder is not overridden → built-in default → "opus" alias + # to Anthropic. + assert captured.get("model") == "opus" + _args, kwargs = mock_spawn.call_args + assert kwargs.get("upstream") in (None, "anthropic"), ( + f"Coder spawn_fn must not carry upstream='litellm' when " + f"only refiner is overridden; got {kwargs.get('upstream')!r}" + ) + assert kwargs.get("upstream_model") is None, ( + f"Coder spawn_fn must have upstream_model=None when only " + f"refiner is overridden; got {kwargs.get('upstream_model')!r}" + ) + + +class TestSpawnClaudeAliasOverride: + """Override to a different Claude alias (e.g. ``"sonnet"``) — the + upstream stays Anthropic but ``--model`` gets the new alias. + """ + + def test_sonnet_override_passes_sonnet_to_consensus_wrapper(self): + if not _slice_2_available(): + pytest.skip("agent_model_resolution not yet implemented") + + from concurrent_executor import ConcurrentPhaseExecutor + from egg_orchestrator.types import AgentRole + + pipeline = _make_pipeline() + try: + pipeline.config.agent_models = {"coder": "sonnet"} + except AttributeError, ValueError: + pipeline.config.__dict__["agent_models"] = {"coder": "sonnet"} + + captured: dict[str, object] = {} + + def _capture_command(prompt_text, **kwargs): + captured["model"] = kwargs.get("model") + return ["bash", "-c", "true"] + + mock_spawn = MagicMock(return_value=_kubernetes_spawn_result("coder")) + executor = ConcurrentPhaseExecutor(pipeline, spawn_fn=mock_spawn) + + with patch( + "concurrent_executor.build_consensus_wrapped_command", + side_effect=_capture_command, + ): + executor._spawn_agent(AgentRole.CODER, prompt_text="run task") + + # Anthropic-classified — alias passes through as-is. + assert captured.get("model") == "sonnet" + _args, kwargs = mock_spawn.call_args + # Upstream is Anthropic → no LiteLLM routing decision. + assert kwargs.get("upstream") in (None, "anthropic") + assert kwargs.get("upstream_model") is None + + +class TestSpawnResolverFailureFallback: + """If ``resolve_agent_model`` ever raises at spawn time, the spawner + MUST degrade to the built-in opus / anthropic default instead of + bringing down the pipeline (defensive guard added in slice-2 v2). + + Mirrors the existing restart-path fallback at + ``routes/pipelines.py:2683-2699``. + """ + + def test_resolver_exception_falls_back_to_opus_anthropic(self): + if not _slice_2_available(): + pytest.skip("agent_model_resolution not yet implemented") + + from concurrent_executor import ConcurrentPhaseExecutor + from egg_orchestrator.types import AgentRole + + pipeline = _make_pipeline() + captured: dict[str, object] = {} + + def _capture_command(prompt_text, **kwargs): + captured["model"] = kwargs.get("model") + return ["bash", "-c", "true"] + + mock_spawn = MagicMock(return_value=_kubernetes_spawn_result("coder")) + executor = ConcurrentPhaseExecutor(pipeline, spawn_fn=mock_spawn) + + # Force the resolver to raise. The spawner's defensive wrap + # MUST catch this and fall back to the built-in opus default. + with ( + patch( + "concurrent_executor.resolve_agent_model", + side_effect=RuntimeError("simulated resolver bug"), + ), + patch( + "concurrent_executor.build_consensus_wrapped_command", + side_effect=_capture_command, + ), + ): + # MUST NOT raise — the defensive guard catches and falls back. + execution = executor._spawn_agent(AgentRole.CODER, prompt_text="run task") + + # Spawn completed → execution returned. + assert execution is not None + # Fallback decision is built-in opus / anthropic; no kwargs added. + assert captured.get("model") == "opus", ( + f"Resolver-failure fallback MUST pass model='opus' to the " + f"consensus wrapper; got {captured.get('model')!r}" + ) + _args, kwargs = mock_spawn.call_args + # No new wire kwargs on the fallback path — preserves the + # pre-#2769 spawn_fn signature for legacy spawners. + assert kwargs.get("upstream") in (None, "anthropic"), ( + f"Resolver-failure fallback added upstream='{kwargs.get('upstream')!r}' " + f"to spawn_fn — must omit on the default-Anthropic path" + ) + assert kwargs.get("upstream_model") is None + + def test_resolver_exception_still_calls_spawn_fn(self): + """Defensive guard MUST NOT short-circuit the spawn — the + agent's container/job must still come up on the built-in + default. (Adversarial probe: a broken `except Exception: raise` + path here would silently break every spawn.) + """ + if not _slice_2_available(): + pytest.skip("agent_model_resolution not yet implemented") + + from concurrent_executor import ConcurrentPhaseExecutor + from egg_orchestrator.types import AgentRole + + pipeline = _make_pipeline() + mock_spawn = MagicMock(return_value=_kubernetes_spawn_result("coder")) + executor = ConcurrentPhaseExecutor(pipeline, spawn_fn=mock_spawn) + + with ( + patch( + "concurrent_executor.resolve_agent_model", + side_effect=ImportError("simulated lazy-import bug"), + ), + patch( + "concurrent_executor.build_consensus_wrapped_command", + return_value=["bash", "-c", "true"], + ), + ): + executor._spawn_agent(AgentRole.CODER, prompt_text="run task") + + assert mock_spawn.call_count == 1, ( + f"Resolver failure MUST NOT short-circuit spawn; spawn_fn " + f"called {mock_spawn.call_count} times (expected 1)" + ) + + +class TestResolverMissingRepoConfigDoesNotCrash: + """Adversarial probe: the v1 NACK called out a FileNotFoundError + leak when ``repositories.yaml`` is absent. v2 ``get_default_agent_model`` + catches that and returns ``None`` so default pipelines still spawn. + + This test exercises the path end-to-end via ``_spawn_agent`` with + a test pipeline that has a non-None ``repo`` (the v1 trigger + condition) and confirms no exception propagates. + """ + + def test_spawn_with_missing_repositories_yaml_does_not_raise(self): + if not _slice_2_available(): + pytest.skip("agent_model_resolution not yet implemented") + + from concurrent_executor import ConcurrentPhaseExecutor + from egg_orchestrator.types import AgentRole + + pipeline = _make_pipeline() + # ``test/repo`` is the _make_pipeline default; the resolver + # would call get_default_agent_model("test/repo") which used + # to raise. After v2 the helper returns None on FileNotFoundError. + assert pipeline.repo == "test/repo" + + mock_spawn = MagicMock(return_value=_kubernetes_spawn_result("coder")) + executor = ConcurrentPhaseExecutor(pipeline, spawn_fn=mock_spawn) + + with patch( + "concurrent_executor.build_consensus_wrapped_command", + return_value=["bash", "-c", "true"], + ): + # MUST NOT raise FileNotFoundError — the v2 fix in + # config.repo_config.get_default_agent_model catches it. + executor._spawn_agent(AgentRole.CODER, prompt_text="run task") + + assert mock_spawn.call_count == 1 diff --git a/orchestrator/tests/test_gateway_client.py b/orchestrator/tests/test_gateway_client.py index 9c709ebb1f..74f0313c58 100644 --- a/orchestrator/tests/test_gateway_client.py +++ b/orchestrator/tests/test_gateway_client.py @@ -2028,3 +2028,138 @@ def test_get_gateway_client_returns_singleton(self): client2 = get_gateway_client() assert client1 is client2 + + +# ============================================================================ +# register_session upstream kwargs — slice-1 of issue #2769 (TASK-1-7) +# ============================================================================ +# +# Slice 1 extends ``GatewayClient.register_session`` with two optional +# kwargs: +# - ``upstream: str | None = None`` +# - ``upstream_model: str | None = None`` +# +# When set, they're included in the POST body to ``/api/v1/sessions/create``. +# When omitted, the body is unchanged from today's shape (back-compat +# guard — no caller in slice 1 actually sends them). +# ============================================================================ + + +class TestRegisterSessionUpstreamKwargs: + """Slice-1 wire-shape extensions to ``register_session``.""" + + def test_omits_upstream_keys_when_not_provided(self, gateway_client): + """Back-compat guard: no slice-1 caller passes the new kwargs. + The wire body must be byte-identical to today's shape. + """ + captured: dict = {} + + def fake_make_request(endpoint, method, data, use_launcher_auth): + captured["data"] = data + return { + "success": True, + "data": { + "session_token": "tok-1", + "created_at": datetime.now().isoformat(), + "expires_at": datetime.now().isoformat(), + }, + } + + with patch.object(gateway_client, "_make_request", side_effect=fake_make_request): + gateway_client.register_session( + container_id="abc", + container_ip="172.18.0.5", + mode="public", + ) + + assert "upstream" not in captured["data"], ( + "register_session must omit 'upstream' from the POST body when " + "not provided — slice-1 no-op invariant" + ) + assert "upstream_model" not in captured["data"], ( + "register_session must omit 'upstream_model' from the POST body " + "when not provided — slice-1 no-op invariant" + ) + + def test_includes_upstream_litellm_when_provided(self, gateway_client): + """Explicit LiteLLM kwargs land in the POST body.""" + captured: dict = {} + + def fake_make_request(endpoint, method, data, use_launcher_auth): + captured["data"] = data + return { + "success": True, + "data": { + "session_token": "tok-1", + "created_at": datetime.now().isoformat(), + "expires_at": datetime.now().isoformat(), + }, + } + + with patch.object(gateway_client, "_make_request", side_effect=fake_make_request): + gateway_client.register_session( + container_id="qwen-agent", + mode="private", + upstream="litellm", + upstream_model="qwen3-coder-30b", + ) + + assert captured["data"].get("upstream") == "litellm" + assert captured["data"].get("upstream_model") == "qwen3-coder-30b" + + def test_includes_upstream_anthropic_explicit_when_provided(self, gateway_client): + """Explicit ``upstream='anthropic'`` is a valid no-op that still + lands in the body (so the audit log records the per-session + upstream — TASK-1-5 AC). + """ + captured: dict = {} + + def fake_make_request(endpoint, method, data, use_launcher_auth): + captured["data"] = data + return { + "success": True, + "data": { + "session_token": "tok-1", + "created_at": datetime.now().isoformat(), + "expires_at": datetime.now().isoformat(), + }, + } + + with patch.object(gateway_client, "_make_request", side_effect=fake_make_request): + gateway_client.register_session( + container_id="explicit-anthropic-agent", + mode="private", + upstream="anthropic", + ) + + assert captured["data"].get("upstream") == "anthropic" + # ``upstream_model`` not provided — should be omitted. + assert "upstream_model" not in captured["data"] + + def test_upstream_only_without_upstream_model_emits_only_upstream(self, gateway_client): + """Asymmetric kwargs: ``upstream='litellm'`` without + ``upstream_model`` is valid (LiteLLM's default model_list will + be consulted upstream). + """ + captured: dict = {} + + def fake_make_request(endpoint, method, data, use_launcher_auth): + captured["data"] = data + return { + "success": True, + "data": { + "session_token": "tok-1", + "created_at": datetime.now().isoformat(), + "expires_at": datetime.now().isoformat(), + }, + } + + with patch.object(gateway_client, "_make_request", side_effect=fake_make_request): + gateway_client.register_session( + container_id="litellm-default-agent", + mode="private", + upstream="litellm", + ) + + assert captured["data"].get("upstream") == "litellm" + assert "upstream_model" not in captured["data"] diff --git a/scripts/await-egg-deploy.sh b/scripts/await-egg-deploy.sh index 64ea0d5d00..2859920f41 100755 --- a/scripts/await-egg-deploy.sh +++ b/scripts/await-egg-deploy.sh @@ -57,12 +57,27 @@ while :; do exit 0 fi - # Fast-fail: a pod can't pull its image. Almost always tag drift — - # HEAD moved since the last build+import, so `make deploy` references - # egg-*:$TAG which was never imported into k3s. - if kubectl -n "$NS" get pods \ - -o jsonpath='{range .items[*]}{range .status.containerStatuses[*]}{.state.waiting.reason}{"\n"}{end}{end}' \ - 2>/dev/null | grep -qE 'ImagePullBackOff|ErrImagePull'; then + # Fast-fail: an egg-owned pod can't pull its image. Almost always tag + # drift — HEAD moved since the last build+import, so `make deploy` + # references egg-*:$TAG which was never imported into k3s. + # + # The scan is scoped by label to egg's own deployments + # (orchestrator/gateway) — those images are egg-built and tag-rewritten + # by `make deploy`. Third-party pods in egg-system (e.g. the LiteLLM + # proxy, whose image is pulled from an external registry) have nothing + # to do with EGG_IMAGE_TAG: their ImagePullBackOff must not abort the + # deploy or mis-blame egg's tag. + egg_image_pull_failed=0 + for d in "${DEPLOYMENTS[@]}"; do + if kubectl -n "$NS" get pods \ + -l "app.kubernetes.io/component=$d" \ + -o jsonpath='{range .items[*]}{range .status.containerStatuses[*]}{.state.waiting.reason}{"\n"}{end}{end}' \ + 2>/dev/null | grep -qE 'ImagePullBackOff|ErrImagePull'; then + egg_image_pull_failed=1 + break + fi + done + if [ "$egg_image_pull_failed" -eq 1 ]; then echo "ERROR: egg-system pods cannot pull image tag '${TAG}' — it is not in k3s." >&2 echo " A commit, pull, or rebase since your last build moved EGG_IMAGE_TAG." >&2 echo " 'make deploy' alone only deploys; run 'make redeploy' to rebuild +" >&2 diff --git a/shared/egg_contracts/agent_roles.py b/shared/egg_contracts/agent_roles.py index 36bbef4d11..2f86b32d84 100644 --- a/shared/egg_contracts/agent_roles.py +++ b/shared/egg_contracts/agent_roles.py @@ -1209,6 +1209,23 @@ def can_retry(self, max_retries: int = 2) -> bool: } +# Roles whose per-pipeline ``PipelineConfig.agent_models`` override is +# actually honored. ``orchestrator.agent_model_resolution.resolve_agent_model`` +# is consulted only by the concurrent-executor spawn path and the +# ``restart_agent`` route, which between them spawn every phase producer +# and reviewer in the two maps above. Utility roles (AUTOFIXER, +# CONFLICT_RESOLVER) and interface roles (OVERSEER, INSPECTOR) spawn +# through dedicated paths that never call the resolver, so an +# ``agent_models`` entry naming one of them would be silently dropped at +# spawn — ``PipelineConfig``'s validator rejects such keys up front. See +# #2769. +MODEL_OVERRIDE_ROLES: frozenset[AgentRole] = frozenset( + role + for role_group in (*_PHASE_ROLES.values(), *_PHASE_REVIEWERS.values()) + for role in role_group +) + + EGG_REPO = "jwbron/egg" # Reviewer roles that only apply to the egg repo itself diff --git a/tests/gateway/test_anthropic_credentials.py b/tests/gateway/test_anthropic_credentials.py index 98250b53a8..74a9f12803 100644 --- a/tests/gateway/test_anthropic_credentials.py +++ b/tests/gateway/test_anthropic_credentials.py @@ -206,3 +206,125 @@ def test_short_oauth_token_rejected(self, tmp_path): cred = manager.get_credential() assert cred is None # Rejected as too short + + +# ============================================================================= +# LiteLLM credential resolver (issue #2769 slice-1, TASK-1-2) +# ============================================================================= +# +# The LiteLLM credential resolver lives alongside the Anthropic one — same +# secrets.env file, same parse_env_file helper, same mtime-invalidated +# cache. It reads ``LITELLM_MASTER_KEY`` and returns an +# ``AnthropicCredential`` shaped ``header_name="x-api-key"``, +# ``header_value=""``. When the env var is absent the resolver +# returns ``None`` (no-op default — matches today's behavior when no +# Anthropic credentials are configured). +# +# The resolver is the parallel ``LiteLLMCredentialsManager`` class in +# ``anthropic_credentials`` — same secrets.env file, same +# ``parse_env_file`` helper, and same mtime-invalidated cache as the +# Anthropic manager. + + +class TestLiteLLMCredentialResolver: + """Tests for the LiteLLM master-key resolver (TASK-1-2).""" + + @pytest.fixture + def _resolve(self): + """Return ``resolve(secrets_path)`` yielding the + ``AnthropicCredential | None`` produced by ``LiteLLMCredentialsManager``. + """ + from anthropic_credentials import LiteLLMCredentialsManager + + def _do(secrets_path: Path): + return LiteLLMCredentialsManager(secrets_path=secrets_path).get_credential() + + return _do + + def test_returns_x_api_key_credential_when_key_present(self, tmp_path, _resolve): + secrets_file = tmp_path / "secrets.env" + secrets_file.write_text('LITELLM_MASTER_KEY="litellm-master-key-1234567890"') + + cred = _resolve(secrets_file) + assert cred is not None + assert cred.header_name == "x-api-key" + assert cred.header_value == "litellm-master-key-1234567890" + assert cred.is_api_key + assert not cred.is_oauth + + def test_returns_none_when_key_missing(self, tmp_path, _resolve): + """No-op default — no warning, no error, just None.""" + secrets_file = tmp_path / "secrets.env" + secrets_file.write_text('ANTHROPIC_API_KEY="sk-ant-12345678901234567890"') + + cred = _resolve(secrets_file) + assert cred is None + + def test_returns_none_when_secrets_file_missing(self, tmp_path, _resolve): + """File absent — fail closed, return None.""" + cred = _resolve(tmp_path / "nonexistent.env") + assert cred is None + + def test_empty_master_key_returns_none(self, tmp_path, _resolve): + """``LITELLM_MASTER_KEY=""`` is the documented disable signal.""" + secrets_file = tmp_path / "secrets.env" + secrets_file.write_text('LITELLM_MASTER_KEY=""') + + cred = _resolve(secrets_file) + assert cred is None + + def test_litellm_resolver_independent_of_anthropic_keys(self, tmp_path, _resolve): + """The LiteLLM resolver MUST NOT fall back to ANTHROPIC_API_KEY + when LITELLM_MASTER_KEY is absent — that would silently route + agent traffic through the wrong credential. + """ + secrets_file = tmp_path / "secrets.env" + secrets_file.write_text( + 'ANTHROPIC_API_KEY="sk-ant-12345678901234567890123456789012345678901234567890"\n' + 'CLAUDE_CODE_OAUTH_TOKEN="oauth-1234567890abcdef"\n' + # No LITELLM_MASTER_KEY entry + ) + + cred = _resolve(secrets_file) + assert cred is None, ( + "LiteLLM resolver must not silently inject the Anthropic credential " + "when LITELLM_MASTER_KEY is absent" + ) + + +class TestLiteLLMResolverCachingBehavior: + """Cache invalidation: changing the file mtime invalidates the cache, + matching ``AnthropicCredentialsManager``'s contract (TASK-1-2 AC). + """ + + def test_mtime_change_invalidates_cache(self, tmp_path): + import time + + from anthropic_credentials import LiteLLMCredentialsManager # type: ignore[attr-defined] + + secrets_file = tmp_path / "secrets.env" + secrets_file.write_text('LITELLM_MASTER_KEY="initial-key-1234567890"') + + manager = LiteLLMCredentialsManager(secrets_path=secrets_file) + cred1 = manager.get_credential() + assert cred1 is not None + assert cred1.header_value == "initial-key-1234567890" + + time.sleep(0.1) + secrets_file.write_text('LITELLM_MASTER_KEY="rotated-key-0987654321"') + + cred2 = manager.get_credential() + assert cred2 is not None + assert cred2.header_value == "rotated-key-0987654321" + + def test_unchanged_file_uses_cached_credential(self, tmp_path): + from anthropic_credentials import LiteLLMCredentialsManager # type: ignore[attr-defined] + + secrets_file = tmp_path / "secrets.env" + secrets_file.write_text('LITELLM_MASTER_KEY="stable-key-1234567890"') + + manager = LiteLLMCredentialsManager(secrets_path=secrets_file) + cred1 = manager.get_credential() + cred2 = manager.get_credential() + # Same object instance — cache hit + assert cred1 is cred2 or cred1 == cred2 diff --git a/tests/gateway/test_anthropic_proxy.py b/tests/gateway/test_anthropic_proxy.py index 3bb9fd18c7..aa40d2c371 100644 --- a/tests/gateway/test_anthropic_proxy.py +++ b/tests/gateway/test_anthropic_proxy.py @@ -1332,3 +1332,1419 @@ def test_capture_handles_invalid_response(self, tmp_path): buffer = TranscriptBuffer(container_id, buffer_dir=tmp_path) entries = buffer.read_entries() assert len(entries) == 0 + + +# ============================================================================= +# Upstream routing — slice-1 of issue #2769 (TASK-1-3 / TASK-1-6) +# ============================================================================= +# +# Slice 1 wires the gateway's two proxy routes (``/v1/messages`` and +# ``/v1/messages/count_tokens``) through a per-request ``UpstreamRegistry`` +# lookup keyed on ``session.upstream``. When the session is absent or +# ``session.upstream == "anthropic"`` the routes MUST behave +# byte-identically to today's hard-wired Anthropic path — that's the +# slice-1 no-op invariant. When ``session.upstream == "litellm"`` the +# routes MUST hit the LiteLLM client and inject the LiteLLM credential +# instead. +# +# The tests below patch the registry / credential resolvers and drive +# the Flask test client to assert the routing decision end-to-end +# without needing a live upstream. +# ============================================================================= + + +def _build_mock_session(upstream: str | None = None, upstream_model: str | None = None): + """Build a MagicMock Session with the upstream + upstream_model fields + needed by the slice-1 routing decision. Falls back to ``"anthropic"`` + when ``upstream is None`` to mirror the production default in the + Session dataclass. + """ + session = MagicMock() + session.mode = "public" + session.container_id = "test-container-routing" + session.upstream = "anthropic" if upstream is None else upstream + session.upstream_model = upstream_model + return session + + +class TestUpstreamRoutingMessages: + """``proxy_anthropic_messages`` routes per ``session.upstream``.""" + + @pytest.fixture + def client(self): + from gateway.gateway import app + + app.config["TESTING"] = True + with app.test_client() as client: + yield client + + def test_no_session_routes_to_anthropic(self, client): + """Backwards-compat: when no session exists for the remote IP, + the request still routes to the Anthropic upstream — the + slice-1 no-op invariant. + """ + from httpx import Headers + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch("gateway.gateway.get_anthropic_client") as mock_anthropic_get, + ): + cred = MagicMock(header_name="x-api-key", header_value="sk-ant-test") + mock_creds_get.return_value.get_credential.return_value = cred + + sm = MagicMock() + sm.get_session_by_ip.return_value = None + mock_sm_get.return_value = sm + + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.content = json.dumps({"content": "ok"}).encode() + mock_response.status_code = 200 + mock_response.headers = Headers([("content-type", "application/json")]) + mock_client.post.return_value = mock_response + mock_anthropic_get.return_value = mock_client + + response = client.post( + "/v1/messages", + data=json.dumps({"model": "claude-3"}), + content_type="application/json", + ) + + assert response.status_code == 200 + # When no session is found, defaulting to anthropic is the + # slice-1 invariant: the Anthropic httpx client MUST be used. + assert mock_client.post.called or mock_client.send.called, ( + "Anthropic client was not invoked for a no-session request" + ) + + def test_anthropic_session_uses_anthropic_upstream(self, client): + """Explicit ``session.upstream == "anthropic"`` still routes to + the Anthropic upstream. + """ + from httpx import Headers + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch("gateway.gateway.get_anthropic_client") as mock_anthropic_get, + ): + cred = MagicMock(header_name="x-api-key", header_value="sk-ant-test") + mock_creds_get.return_value.get_credential.return_value = cred + + sm = MagicMock() + sm.get_session_by_ip.return_value = _build_mock_session(upstream="anthropic") + mock_sm_get.return_value = sm + + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.content = json.dumps({"content": "ok"}).encode() + mock_response.status_code = 200 + mock_response.headers = Headers([("content-type", "application/json")]) + mock_client.post.return_value = mock_response + mock_anthropic_get.return_value = mock_client + + response = client.post( + "/v1/messages", + data=json.dumps({"model": "claude-3"}), + content_type="application/json", + ) + + assert response.status_code == 200 + assert mock_client.post.called or mock_client.send.called + + def test_litellm_session_uses_litellm_upstream(self, client): + """``session.upstream == "litellm"`` routes to the LiteLLM client + from the upstream registry; the Anthropic client is NOT used. + """ + from httpx import Headers + + litellm_client = MagicMock() + mock_response = MagicMock() + mock_response.content = json.dumps({"content": "ok"}).encode() + mock_response.status_code = 200 + mock_response.headers = Headers([("content-type", "application/json")]) + litellm_client.post.return_value = mock_response + + anthropic_client = MagicMock() + # Anthropic client should NOT be called. + anthropic_client.post.side_effect = AssertionError( + "Anthropic client must not be invoked for a LiteLLM-routed request" + ) + + # Build a registry that dispatches per upstream name. + def _registry_get(upstream): + if upstream == "anthropic": + return ( + anthropic_client, + lambda: MagicMock(header_name="x-api-key", header_value="sk-ant-test"), + ) + if upstream == "litellm": + return ( + litellm_client, + lambda: MagicMock(header_name="x-api-key", header_value="litellm-key"), + ) + raise KeyError(upstream) + + fake_registry = MagicMock() + fake_registry.get.side_effect = _registry_get + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch("gateway.gateway.get_upstream_registry", return_value=fake_registry, create=True), + patch("gateway.gateway.get_anthropic_client", return_value=anthropic_client), + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-test" + ) + mock_litellm_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="litellm-key" + ) + sm = MagicMock() + sm.get_session_by_ip.return_value = _build_mock_session( + upstream="litellm", upstream_model="qwen3-coder-30b" + ) + mock_sm_get.return_value = sm + + response = client.post( + "/v1/messages", + data=json.dumps({"model": "opus"}), # cq-5 alias on the wire + content_type="application/json", + ) + + assert response.status_code == 200 + assert litellm_client.post.called or litellm_client.send.called, ( + "LiteLLM client was not invoked for a LiteLLM-routed request" + ) + + def test_litellm_request_injects_litellm_credential(self, client): + """The header injected on a LiteLLM-routed request must be the + LiteLLM ``x-api-key``, not the Anthropic one — otherwise the + gateway leaks the Anthropic credential to LiteLLM (and vice + versa). + """ + from httpx import Headers + + captured_headers: dict[str, str] = {} + + def _post_capture(*_args, **kwargs): + captured_headers.update(kwargs.get("headers", {})) + mock_response = MagicMock() + mock_response.content = b"{}" + mock_response.status_code = 200 + mock_response.headers = Headers([("content-type", "application/json")]) + return mock_response + + litellm_client = MagicMock() + litellm_client.post.side_effect = _post_capture + anthropic_client = MagicMock() + + def _registry_get(upstream): + if upstream == "anthropic": + return ( + anthropic_client, + lambda: MagicMock(header_name="x-api-key", header_value="sk-ant-shouldnotleak"), + ) + if upstream == "litellm": + return ( + litellm_client, + lambda: MagicMock( + header_name="x-api-key", header_value="litellm-key-only-this" + ), + ) + raise KeyError(upstream) + + fake_registry = MagicMock() + fake_registry.get.side_effect = _registry_get + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch("gateway.gateway.get_upstream_registry", return_value=fake_registry, create=True), + patch("gateway.gateway.get_anthropic_client", return_value=anthropic_client), + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-shouldnotleak" + ) + mock_litellm_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="litellm-key-only-this" + ) + sm = MagicMock() + sm.get_session_by_ip.return_value = _build_mock_session( + upstream="litellm", upstream_model="qwen3-coder-30b" + ) + mock_sm_get.return_value = sm + + client.post( + "/v1/messages", + data=json.dumps({"model": "opus"}), + content_type="application/json", + ) + + assert captured_headers.get("x-api-key") == "litellm-key-only-this", ( + f"LiteLLM-routed request did not inject LiteLLM credential; " + f"got headers: {captured_headers}" + ) + assert "sk-ant-shouldnotleak" not in captured_headers.values(), ( + "Anthropic credential leaked into LiteLLM-routed request headers" + ) + + +class TestUpstreamRoutingCountTokens: + """``proxy_count_tokens`` mirrors the routing decision.""" + + @pytest.fixture + def client(self): + from gateway.gateway import app + + app.config["TESTING"] = True + with app.test_client() as client: + yield client + + def test_anthropic_session_count_tokens_uses_anthropic_upstream(self, client): + from httpx import Headers + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch("gateway.gateway.get_anthropic_client") as mock_anthropic_get, + ): + cred = MagicMock(header_name="x-api-key", header_value="sk-ant-test") + mock_creds_get.return_value.get_credential.return_value = cred + + sm = MagicMock() + sm.get_session_by_ip.return_value = _build_mock_session(upstream="anthropic") + mock_sm_get.return_value = sm + + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.content = json.dumps({"input_tokens": 1}).encode() + mock_response.status_code = 200 + mock_response.headers = Headers([("content-type", "application/json")]) + mock_client.post.return_value = mock_response + mock_anthropic_get.return_value = mock_client + + response = client.post( + "/v1/messages/count_tokens", + data=json.dumps({"model": "claude-3", "messages": []}), + content_type="application/json", + ) + + assert response.status_code == 200 + assert mock_client.post.called + + def test_litellm_session_count_tokens_uses_litellm_upstream(self, client): + from httpx import Headers + + litellm_client = MagicMock() + mock_response = MagicMock() + mock_response.content = json.dumps({"input_tokens": 7}).encode() + mock_response.status_code = 200 + mock_response.headers = Headers([("content-type", "application/json")]) + litellm_client.post.return_value = mock_response + + anthropic_client = MagicMock() + anthropic_client.post.side_effect = AssertionError( + "Anthropic client must not be invoked for a LiteLLM count_tokens request" + ) + + def _registry_get(upstream): + if upstream == "anthropic": + return ( + anthropic_client, + lambda: MagicMock(header_name="x-api-key", header_value="sk-ant-test"), + ) + if upstream == "litellm": + return ( + litellm_client, + lambda: MagicMock(header_name="x-api-key", header_value="litellm-key"), + ) + raise KeyError(upstream) + + fake_registry = MagicMock() + fake_registry.get.side_effect = _registry_get + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch("gateway.gateway.get_upstream_registry", return_value=fake_registry, create=True), + patch("gateway.gateway.get_anthropic_client", return_value=anthropic_client), + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-test" + ) + mock_litellm_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="litellm-key" + ) + sm = MagicMock() + sm.get_session_by_ip.return_value = _build_mock_session( + upstream="litellm", upstream_model="qwen3-coder-30b" + ) + mock_sm_get.return_value = sm + + response = client.post( + "/v1/messages/count_tokens", + data=json.dumps({"model": "opus", "messages": []}), + content_type="application/json", + ) + + assert response.status_code == 200 + assert litellm_client.post.called, ( + "LiteLLM client was not invoked for a LiteLLM count_tokens request" + ) + + +class TestInjectUpstreamCredentials: + """``_inject_upstream_credentials(headers, upstream)`` dispatches per upstream. + + Back-compat: ``_inject_anthropic_credentials`` is preserved as a thin + alias calling through with ``upstream="anthropic"`` (TASK-1-3 AC). + """ + + @pytest.fixture + def _inject_fn(self): + """Return the upstream-aware credential injector.""" + from gateway.gateway import _inject_upstream_credentials + + return _inject_upstream_credentials + + def test_anthropic_dispatch_matches_legacy_helper(self, _inject_fn): + """For ``upstream="anthropic"``, the new helper behaves + byte-identically to today's ``_inject_anthropic_credentials``. + """ + from gateway.gateway import _inject_anthropic_credentials, app + + with patch("gateway.gateway.get_credentials_manager") as mock_get: + cred = MagicMock(header_name="x-api-key", header_value="sk-ant-byte-identical") + mock_get.return_value.get_credential.return_value = cred + + headers_a, error_a = _inject_anthropic_credentials({"Content-Type": "application/json"}) + headers_b, error_b = _inject_fn({"Content-Type": "application/json"}, "anthropic") + + assert error_a is None + assert error_b is None + assert headers_a == headers_b + + # Also assert the 401 shape matches for the no-credential path. + with patch("gateway.gateway.get_credentials_manager") as mock_get: + mock_get.return_value.get_credential.return_value = None + with app.app_context(): + _h_a, err_a = _inject_anthropic_credentials({"Content-Type": "application/json"}) + _h_b, err_b = _inject_fn({"Content-Type": "application/json"}, "anthropic") + assert err_a is not None and err_b is not None + assert err_a[1] == err_b[1] == 401 + + def test_litellm_dispatch_injects_litellm_credential(self, _inject_fn): + """``upstream="litellm"`` injects the LiteLLM ``x-api-key`` from + the LiteLLM credential resolver, NOT the Anthropic one. + """ + # Patch both resolvers; assert only LiteLLM's value lands in the + # headers. The Anthropic resolver MUST NOT be consulted on this + # path. + with ( + patch("gateway.gateway.get_credentials_manager") as mock_anthropic_get, + patch("gateway.gateway.get_litellm_credentials_manager") as mock_litellm_get, + ): + # Make Anthropic resolver explosive — if it's called, the + # test fails loudly. + mock_anthropic_get.return_value.get_credential.side_effect = AssertionError( + "Anthropic resolver consulted on LiteLLM-routed request" + ) + mock_litellm_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", + header_value="litellm-master-key-1234567890", + ) + headers, error = _inject_fn({"Content-Type": "application/json"}, "litellm") + assert error is None + assert headers["x-api-key"] == "litellm-master-key-1234567890" + + def test_litellm_no_credential_returns_401(self, _inject_fn): + """Same 401 shape as today's Anthropic-no-credential path.""" + from gateway.gateway import app + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_anthropic_get, + patch("gateway.gateway.get_litellm_credentials_manager") as mock_litellm_get, + ): + mock_anthropic_get.return_value.get_credential.return_value = None + mock_litellm_get.return_value.get_credential.return_value = None + with app.app_context(): + _headers, error = _inject_fn({"Content-Type": "application/json"}, "litellm") + assert error is not None + assert error[1] == 401 + + def test_unknown_upstream_returns_502_without_anthropic_fallthrough(self, _inject_fn): + """An upstream the registry does not serve must fail closed with a + 502 — never silently treated as Anthropic. Regression guard for + the pre-review behavior where an unknown upstream's error code + depended on unrelated Anthropic-credential state (401 vs 502). + """ + from gateway.gateway import app + + with patch("gateway.gateway.get_credentials_manager") as mock_anthropic_get: + # No Anthropic credential configured — the old fall-through + # would have produced a 401 here instead of a 502. + mock_anthropic_get.return_value.get_credential.return_value = None + with app.app_context(): + _headers, error = _inject_fn({"Content-Type": "application/json"}, "bogus_upstream") + assert error is not None + assert error[1] == 502, ( + f"Unknown upstream must fail closed with 502 regardless of " + f"Anthropic-credential state; got {error[1]}" + ) + + +# ============================================================================= +# Adversarial probes — issue #2769 slice-1 +# ============================================================================= +# +# The probes below target seams that are easy to get wrong: +# +# - Unknown-upstream defense in the proxy route (TASK-1-6) — the +# code path that fires when a session somehow ends up with an upstream +# the registry does not serve (corrupted persistence, slice-2 +# misconfig). Must fail closed with 502, not crash or silently +# forward to Anthropic. +# +# - The two proxy routes (proxy_anthropic_messages and +# proxy_count_tokens) MUST agree on the upstream for a given session +# — split-brain (messages routed to LiteLLM, count_tokens routed to +# Anthropic) would break Claude Code's token accounting badly. +# +# - LiteLLM upstream must NOT trigger the Anthropic "client-supplied +# auth fall-through" — i.e. with no LITELLM_MASTER_KEY configured, +# a Claude Code request that happens to carry an Authorization +# header must NOT silently route to LiteLLM with that header. +# ============================================================================= + + +class TestUnknownUpstreamDefense: + """Defensive 502 when session.upstream is unknown at proxy time.""" + + @pytest.fixture + def client(self): + from gateway.gateway import app + + app.config["TESTING"] = True + with app.test_client() as client: + yield client + + def test_unknown_upstream_on_session_returns_5xx(self, client): + """If a session escapes session-create validation with an unknown + upstream (e.g. corrupted persistence, manual edit), the proxy + MUST fail closed — never silently fall back to Anthropic. + + ``_inject_upstream_credentials`` checks ``UpstreamRegistry.is_known`` + before any per-upstream branch, so an unregistered upstream is + rejected with a deterministic 502 ahead of the client-resolution + block (the ``except UnknownUpstreamError`` there is now unreachable + defensive code). We still assert on the 5xx range rather than the + exact 502 because the fail-closed contract — not the specific + code — is what this test guards. + """ + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-test" + ) + # Build a session whose upstream is one the registry will not + # serve. The session-create endpoint validates upstream before + # storing it, but this guards the in-flight error case where a + # session somehow lands with an unknown value (slice-2 misconfig, + # or persistence corruption). + session = _build_mock_session(upstream="bogus_upstream") + sm = MagicMock() + sm.get_session_by_ip.return_value = session + mock_sm_get.return_value = sm + + response = client.post( + "/v1/messages", + data=json.dumps({"model": "claude-3"}), + content_type="application/json", + ) + + # cq-8: Fail closed on LiteLLM unreachable; the + # unknown-upstream branch is the most-likely-to-fire variant + # of "we can't reach what the session told us to reach". + assert 500 <= response.status_code < 600, ( + f"Unknown upstream MUST fail closed (5xx), got " + f"{response.status_code} ({response.data!r})" + ) + # The response MUST NOT be a 2xx — Anthropic must not have + # been hit as a silent fallback. + assert response.status_code != 200, ( + "Unknown upstream silently fell back to Anthropic — " + "this is the cq-8 fail-closed contract violation" + ) + + +class TestRoutingConsistencyAcrossProxyRoutes: + """``/v1/messages`` and ``/v1/messages/count_tokens`` MUST agree on + upstream for any given session. A split-brain (messages -> LiteLLM, + count_tokens -> Anthropic) silently corrupts Claude Code's token + accounting and is hard to detect post-hoc. + """ + + @pytest.fixture + def client(self): + from gateway.gateway import app + + app.config["TESTING"] = True + with app.test_client() as client: + yield client + + def test_messages_and_count_tokens_agree_on_litellm_routing(self, client): + """Both routes should land on the LiteLLM client when the + session is upstream=='litellm'. + """ + from httpx import Headers + + litellm_client = MagicMock() + mock_response = MagicMock() + mock_response.content = b"{}" + mock_response.status_code = 200 + mock_response.headers = Headers([("content-type", "application/json")]) + litellm_client.post.return_value = mock_response + + anthropic_client = MagicMock() + + def _registry_get(upstream): + if upstream == "anthropic": + return (anthropic_client, lambda: None) + if upstream == "litellm": + return (litellm_client, lambda: None) + raise KeyError(upstream) + + fake_registry = MagicMock() + fake_registry.get.side_effect = _registry_get + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch("gateway.gateway.get_upstream_registry", return_value=fake_registry, create=True), + patch("gateway.gateway.get_anthropic_client", return_value=anthropic_client), + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-test" + ) + mock_litellm_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="litellm-key" + ) + session = _build_mock_session(upstream="litellm", upstream_model="qwen3-coder-30b") + sm = MagicMock() + sm.get_session_by_ip.return_value = session + mock_sm_get.return_value = sm + + # Hit both routes back-to-back with the same session. + client.post( + "/v1/messages", + data=json.dumps({"model": "opus"}), + content_type="application/json", + ) + client.post( + "/v1/messages/count_tokens", + data=json.dumps({"model": "opus", "messages": []}), + content_type="application/json", + ) + + # Both must have hit the LiteLLM client. An Anthropic call + # here indicates the split-brain bug. + assert litellm_client.post.call_count == 2, ( + f"Split-brain routing: LiteLLM hit " + f"{litellm_client.post.call_count} times across messages " + f"and count_tokens — expected 2. anthropic_client.post " + f"called {anthropic_client.post.call_count} times " + f"(should be 0)." + ) + assert anthropic_client.post.call_count == 0 + + +class TestLiteLLMNoFallbackToClientAuth: + """The LiteLLM upstream MUST NOT honour client-supplied + Authorization / x-api-key headers as a fallback when + LITELLM_MASTER_KEY is unset. That fall-through is the Anthropic + path's OAuth-mode escape hatch; on the LiteLLM path it would route + a Claude Code Anthropic OAuth token to a third-party LiteLLM + backend, leaking the credential. + + TASK-1-3 AC: ``Missing credentials for either upstream return a + 401 with the same JSON body shape as today.`` + """ + + @pytest.fixture + def client(self): + from gateway.gateway import app + + app.config["TESTING"] = True + with app.test_client() as client: + yield client + + def test_litellm_no_master_key_does_not_honour_client_auth(self, client): + """With LITELLM_MASTER_KEY unset and the Claude Code client + sending an Authorization header, the LiteLLM path returns 401 + — does NOT silently forward the Anthropic OAuth token to + LiteLLM. + """ + with ( + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + # LiteLLM has NO credential. + mock_litellm_get.return_value.get_credential.return_value = None + + session = _build_mock_session(upstream="litellm", upstream_model="qwen3-coder-30b") + sm = MagicMock() + sm.get_session_by_ip.return_value = session + mock_sm_get.return_value = sm + + response = client.post( + "/v1/messages", + data=json.dumps({"model": "opus"}), + content_type="application/json", + headers={"Authorization": "Bearer claude-oauth-should-not-leak"}, + ) + + assert response.status_code == 401, ( + f"LiteLLM path with no master key MUST return 401 even " + f"when client carries Authorization header. Got " + f"{response.status_code} ({response.data!r}) — this " + f"would leak the Claude OAuth token to a third-party " + f"LiteLLM backend." + ) + + def test_litellm_no_master_key_does_not_honour_client_api_key(self, client): + """Same guard for x-api-key — client-supplied API keys must + not bypass the LiteLLM credential gate. + """ + with ( + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + mock_litellm_get.return_value.get_credential.return_value = None + + session = _build_mock_session(upstream="litellm", upstream_model="qwen3-coder-30b") + sm = MagicMock() + sm.get_session_by_ip.return_value = session + mock_sm_get.return_value = sm + + response = client.post( + "/v1/messages", + data=json.dumps({"model": "opus"}), + content_type="application/json", + headers={"x-api-key": "sk-ant-should-not-leak"}, + ) + + assert response.status_code == 401 + + +class TestSessionCreateUpstreamValidation: + """Slice-1 session-create endpoint MUST reject unknown ``upstream`` + values with 400 (TASK-1-5 AC). + + These patch ``get_launcher_secret`` directly and send the matching + bearer token, so the validation branch is exercised deterministically + — they do not depend on ``EGG_LAUNCHER_SECRET`` being set in the + environment. + """ + + _SECRET = "test-launcher-secret-upstream-validation" + + @pytest.fixture + def client(self): + from gateway.gateway import app + + app.config["TESTING"] = True + with app.test_client() as client: + yield client + + def _auth(self): + return {"Authorization": f"Bearer {self._SECRET}"} + + def test_bogus_upstream_returns_400(self, client): + """POST /api/v1/sessions/create with ``upstream='bogus'`` returns + a 400 with a descriptive error. The upstream check fires before + any worktree/session work, so patching launcher auth is the only + stub the rejection path needs. + """ + with patch("gateway.gateway.get_launcher_secret", return_value=self._SECRET): + response = client.post( + "/api/v1/sessions/create", + data=json.dumps( + { + "container_id": "test-container", + "container_ip": "172.18.0.5", + "mode": "private", + "pipeline_id": "test-pipeline", + "upstream": "bogus_upstream_name", + } + ), + content_type="application/json", + headers=self._auth(), + ) + + assert response.status_code == 400, ( + f"Bogus upstream MUST return 400; got {response.status_code} ({response.data!r})" + ) + body = json.loads(response.data) + # The error message should mention the rejected upstream so the + # operator can debug. The exact phrasing is flexible. + msg = body.get("message", body.get("error", {}).get("message", "")) + assert "upstream" in str(msg).lower() or "bogus_upstream_name" in str(body) + + def test_known_upstreams_pass_validation(self, client, tmp_path): + """The registry's two known upstreams pass session-create + validation and reach a clean 200. A real SessionManager is wired + in so ``register_session`` actually runs; ``repos`` is omitted so + no worktree machinery is touched. + """ + from session_manager import SessionManager + + for upstream in ("anthropic", "litellm"): + manager = SessionManager(persistence_file=tmp_path / f"sessions-{upstream}.json") + with ( + patch("gateway.gateway.get_launcher_secret", return_value=self._SECRET), + patch("gateway.gateway.get_session_manager", return_value=manager), + ): + response = client.post( + "/api/v1/sessions/create", + data=json.dumps( + { + "container_id": f"test-{upstream}", + "container_ip": "172.18.0.5", + "mode": "private", + "pipeline_id": f"test-pipeline-{upstream}", + "upstream": upstream, + } + ), + content_type="application/json", + headers=self._auth(), + ) + + assert response.status_code == 200, ( + f"Valid upstream '{upstream}' was rejected: " + f"{response.status_code} ({response.data!r})" + ) + + +# ============================================================================= +# Upstream body rewrite — slice-2 of issue #2769 (TASK-2-6 / TASK-2-8) +# ============================================================================= +# +# Slice 2 adds ``_rewrite_upstream_model(request_body, upstream_model)`` +# next to ``_filter_blocked_tools`` in gateway.py. On LiteLLM-routed +# requests with ``session.upstream_model`` set, the helper REPLACES the +# top-level ``"model"`` field in the JSON body with the upstream-side +# model name (e.g. ``"qwen3-coder-30b"``). This is the cq-5 mitigation +# on the wire: Claude Code is presented a recognized Claude alias +# (``"opus"``) as ``--model``, so its compaction math stays sane; the +# gateway rewrites the body just before forwarding to LiteLLM so the +# upstream actually receives the right model name. +# +# Invariants: +# +# - ``upstream == "anthropic"`` → body is byte-identical (regression). +# - ``upstream == "litellm"`` and ``upstream_model is None`` → body +# unchanged (slice-1 no-op state). +# - ``upstream == "litellm"`` and ``upstream_model="qwen3-coder-30b"`` +# → forwarded body has ``"model": "qwen3-coder-30b"`` regardless of +# incoming ``"model"`` value. +# - Invalid JSON → original body returned unchanged (proxy MUST NOT +# crash on a malformed body — slice-1 ``_filter_blocked_tools`` +# matches this contract). +# - The rewrite happens AFTER ``_filter_blocked_tools`` and BEFORE the +# upstream request is built, so blocked-tool stripping in private +# mode still works. +# ============================================================================= + + +def _capture_upstream_body(captured_holder: dict, status: int = 200, response_body: bytes = b"{}"): + """Build an httpx-mock side_effect that captures the body forwarded + to the upstream client. Stores under ``captured_holder["body"]``. + Mirrors the slice-1 ``_post_capture`` helper used by the + credential-leak tests. + """ + from httpx import Headers + + def _capture(*args, **kwargs): + # The proxy builds requests one of two ways depending on + # streaming: ``client.post(url, content=request_body, ...)`` or + # ``client.build_request("POST", url, content=request_body, ...)``. + # Both routes feed ``request_body`` via the ``content`` kwarg. + body = kwargs.get("content") + if body is None: + # Fall back to positional inspection for safety. + for arg in args: + if isinstance(arg, (bytes, bytearray, str)): + body = arg + break + captured_holder["body"] = body + + mock_response = MagicMock() + mock_response.content = response_body + mock_response.status_code = status + mock_response.headers = Headers([("content-type", "application/json")]) + return mock_response + + return _capture + + +class TestRewriteUpstreamModelHelper: + """Direct unit tests on the ``_rewrite_upstream_model`` helper. + + Skips when the helper has not landed yet (waiting on coder). + """ + + @pytest.fixture + def _rewrite_fn(self): + try: + from gateway.gateway import _rewrite_upstream_model # type: ignore[attr-defined] + + return _rewrite_upstream_model + except ImportError: + pytest.skip("_rewrite_upstream_model not yet implemented (waiting on coder)") + + def test_no_op_when_upstream_model_is_none(self, _rewrite_fn): + body = json.dumps({"model": "opus", "messages": []}).encode() + out = _rewrite_fn(body, None) + # Byte-identical when no rewrite is requested. + assert out == body + + def test_rewrites_top_level_model_field(self, _rewrite_fn): + body = json.dumps({"model": "opus", "messages": []}).encode() + out = _rewrite_fn(body, "qwen3-coder-30b") + parsed = json.loads(out) + assert parsed["model"] == "qwen3-coder-30b" + # Other fields preserved. + assert parsed["messages"] == [] + + def test_preserves_other_top_level_keys(self, _rewrite_fn): + """The rewrite must not drop other body fields — system prompt, + tools, max_tokens, etc. If it does, Claude Code's request + shape silently changes shape across the gateway. + """ + body = json.dumps( + { + "model": "opus", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 4096, + "system": "You are a helpful assistant.", + "tools": [{"name": "bash"}], + "stream": True, + } + ).encode() + out = _rewrite_fn(body, "qwen3-coder-30b") + parsed = json.loads(out) + assert parsed["model"] == "qwen3-coder-30b" + assert parsed["max_tokens"] == 4096 + assert parsed["system"] == "You are a helpful assistant." + assert parsed["tools"] == [{"name": "bash"}] + assert parsed["stream"] is True + assert parsed["messages"] == [{"role": "user", "content": "hi"}] + + def test_invalid_json_returns_original_body(self, _rewrite_fn): + """Slice-1 ``_filter_blocked_tools`` matches this contract — + on JSONDecodeError the helper returns the input unchanged so + the proxy doesn't crash on a malformed body. Slice-2 must do + the same. + """ + body = b"not valid json {{" + out = _rewrite_fn(body, "qwen3-coder-30b") + assert out == body + + def test_empty_body_returns_unchanged(self, _rewrite_fn): + body = b"" + out = _rewrite_fn(body, "qwen3-coder-30b") + # Either byte-identical or a degenerate ``{}``-rewrite is + # acceptable; what's not acceptable is a crash. + assert isinstance(out, (bytes, bytearray)) + + def test_body_without_model_key_is_handled(self, _rewrite_fn): + """Adversarial: incoming body has no ``model`` key. The helper + either inserts the upstream model (so LiteLLM still gets the + right model name) or returns the body unchanged — what MUST + NOT happen is a KeyError crash. + """ + body = json.dumps({"messages": [{"role": "user", "content": "hi"}]}).encode() + out = _rewrite_fn(body, "qwen3-coder-30b") + # Should not raise, and should not corrupt the rest of the body. + # If the implementation chooses to inject the model field, the + # result MUST be valid JSON with ``messages`` preserved. + try: + parsed = json.loads(out) + assert parsed.get("messages") == [{"role": "user", "content": "hi"}] + except json.JSONDecodeError: + # Returning the body unchanged is also acceptable. + assert out == body + + def test_unicode_model_name_round_trips(self, _rewrite_fn): + """The helper must not break on non-ASCII model names — even + though LiteLLM model names are ASCII in practice, the helper + shouldn't impose a stricter encoding constraint than the + original body parser. + """ + body = json.dumps({"model": "opus", "messages": []}).encode() + out = _rewrite_fn(body, "qwen-✨-30b") + parsed = json.loads(out) + assert parsed["model"] == "qwen-✨-30b" + + +def _rewrite_helper_exists() -> bool: + """Return True if slice-2's ``_rewrite_upstream_model`` has landed + on the gateway side. Tests below skip when False. + """ + try: + # noqa: F401 — import is for existence check, the symbol is unused. + from gateway.gateway import _rewrite_upstream_model # type: ignore[attr-defined] # noqa: F401, I001 + + return True + except ImportError: + return False + + +class TestRewriteUpstreamModelOnProxyRoute: + """End-to-end: ``proxy_anthropic_messages`` with a LiteLLM session + forwards the rewritten body to the LiteLLM upstream. + + These tests drive the Flask app like slice-1's + ``TestUpstreamRoutingMessages`` and assert on the body forwarded + to the LiteLLM-side httpx client. + """ + + @pytest.fixture + def client(self): + from gateway.gateway import app + + app.config["TESTING"] = True + with app.test_client() as client: + yield client + + def _build_full_registry_patch(self, captured_body_holder): + """Build the registry / credential patches used by the + body-rewrite tests. Returns a list of context managers ready + to ``with ... :``. + """ + try: + import upstream_registry # type: ignore[import-not-found] # noqa: F401 + except ImportError: + pytest.skip("upstream_registry not yet implemented (waiting on coder)") + + litellm_client = MagicMock() + litellm_client.post.side_effect = _capture_upstream_body(captured_body_holder) + # Also wire build_request → send → iter_bytes for the streaming + # path. Slice-2 tests target the non-streaming path, but we set + # up both to avoid spurious AttributeErrors. + litellm_client.build_request.return_value = MagicMock() + + anthropic_client = MagicMock() + anthropic_client.post.side_effect = AssertionError( + "Anthropic client must not be invoked for a LiteLLM-routed request" + ) + + def _registry_get(upstream): + if upstream == "anthropic": + return ( + anthropic_client, + lambda: MagicMock(header_name="x-api-key", header_value="sk-ant-test"), + ) + if upstream == "litellm": + return ( + litellm_client, + lambda: MagicMock(header_name="x-api-key", header_value="litellm-key"), + ) + raise KeyError(upstream) + + fake_registry = MagicMock() + fake_registry.get.side_effect = _registry_get + return fake_registry, litellm_client, anthropic_client + + def test_litellm_session_with_upstream_model_rewrites_body(self, client): + """With ``upstream="litellm"`` and + ``upstream_model="qwen3-coder-30b"``, the body forwarded to the + LiteLLM client has ``"model": "qwen3-coder-30b"`` regardless of + what Claude Code sent (it sends ``"opus"`` per the cq-5 + mitigation). + """ + if not _rewrite_helper_exists(): + pytest.skip("_rewrite_upstream_model not yet implemented") + captured: dict = {} + fake_registry, litellm_client, _ = self._build_full_registry_patch(captured) + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch( + "gateway.gateway.get_upstream_registry", + return_value=fake_registry, + create=True, + ), + patch("gateway.gateway.get_anthropic_client") as mock_anthropic_get, + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-test" + ) + mock_litellm_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="litellm-key" + ) + # ``get_anthropic_client`` is also called in the no-session + # path; wire it to a harmless mock to be safe. + mock_anthropic_get.return_value = MagicMock() + + sm = MagicMock() + sm.get_session_by_ip.return_value = _build_mock_session( + upstream="litellm", upstream_model="qwen3-coder-30b" + ) + mock_sm_get.return_value = sm + + response = client.post( + "/v1/messages", + # Claude Code presents 'opus' on the wire (cq-5). + data=json.dumps({"model": "opus", "messages": []}), + content_type="application/json", + ) + + assert response.status_code == 200 + body = captured.get("body") + assert body is not None, "LiteLLM client was not called with a body to capture" + parsed = json.loads(body) if isinstance(body, (bytes, bytearray)) else json.loads(body) + assert parsed["model"] == "qwen3-coder-30b", ( + f"LiteLLM-routed body MUST have model='qwen3-coder-30b' " + f"(rewritten from incoming 'opus'); got {parsed.get('model')!r}" + ) + + def test_litellm_session_without_upstream_model_preserves_body_model(self, client): + """Slice-1 no-op state: a LiteLLM session with + ``upstream_model=None`` MUST NOT rewrite the body — the proxy + passes whatever model name the client sent. + """ + captured: dict = {} + fake_registry, _litellm_client, _ = self._build_full_registry_patch(captured) + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch( + "gateway.gateway.get_upstream_registry", + return_value=fake_registry, + create=True, + ), + patch("gateway.gateway.get_anthropic_client"), + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-test" + ) + mock_litellm_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="litellm-key" + ) + + sm = MagicMock() + sm.get_session_by_ip.return_value = _build_mock_session( + upstream="litellm", upstream_model=None + ) + mock_sm_get.return_value = sm + + client.post( + "/v1/messages", + data=json.dumps({"model": "opus", "messages": []}), + content_type="application/json", + ) + + body = captured.get("body") + assert body is not None + parsed = json.loads(body) + assert parsed["model"] == "opus", ( + f"upstream_model=None must NOT rewrite the body; got model={parsed.get('model')!r}" + ) + + def test_anthropic_session_body_is_byte_identical(self, client): + """Regression guard for the Anthropic path — with no session, + or ``upstream="anthropic"``, the body forwarded to Anthropic is + BYTE-identical to the body the client sent. This is the + slice-2 no-op invariant. + """ + from httpx import Headers + + captured: dict = {} + + def _post_capture(*args, **kwargs): + captured["body"] = kwargs.get("content") or (args[1] if len(args) > 1 else None) + mock_response = MagicMock() + mock_response.content = b'{"ok": true}' + mock_response.status_code = 200 + mock_response.headers = Headers([("content-type", "application/json")]) + return mock_response + + mock_anthropic_client = MagicMock() + mock_anthropic_client.post.side_effect = _post_capture + + # Build a registry where 'anthropic' returns our client — and + # 'litellm' is never expected to be called on this test. Use + # ``get_anthropic_client`` so the proxy's fast path picks up + # the mock (the slice-1 code shape calls that directly for the + # anthropic upstream). + original_body = json.dumps({"model": "claude-3-5-sonnet-20241022", "messages": []}) + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch("gateway.gateway.get_anthropic_client", return_value=mock_anthropic_client), + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-test" + ) + sm = MagicMock() + sm.get_session_by_ip.return_value = _build_mock_session(upstream="anthropic") + mock_sm_get.return_value = sm + + response = client.post( + "/v1/messages", + data=original_body, + content_type="application/json", + ) + + assert response.status_code == 200 + body = captured.get("body") + # Byte-identical regression guard: Anthropic path MUST NOT + # rewrite anything. The incoming model name (a full + # ``claude-3-5-sonnet-...`` alias) must survive verbatim. + assert body == original_body.encode(), ( + f"Anthropic path body MUST be byte-identical " + f"(regression guard for slice-2); incoming " + f"{original_body!r}, forwarded {body!r}" + ) + + def test_litellm_count_tokens_rewrites_body(self, client): + """The body-rewrite must mirror across BOTH proxy routes — + otherwise Claude Code's token accounting for a LiteLLM-routed + agent silently drifts from the actual upstream model. + """ + if not _rewrite_helper_exists(): + pytest.skip("_rewrite_upstream_model not yet implemented") + captured: dict = {} + fake_registry, litellm_client, _ = self._build_full_registry_patch(captured) + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch( + "gateway.gateway.get_upstream_registry", + return_value=fake_registry, + create=True, + ), + patch("gateway.gateway.get_anthropic_client"), + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-test" + ) + mock_litellm_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="litellm-key" + ) + sm = MagicMock() + sm.get_session_by_ip.return_value = _build_mock_session( + upstream="litellm", upstream_model="qwen3-coder-30b" + ) + mock_sm_get.return_value = sm + + response = client.post( + "/v1/messages/count_tokens", + data=json.dumps({"model": "opus", "messages": []}), + content_type="application/json", + ) + + assert response.status_code == 200 + body = captured.get("body") + assert body is not None + parsed = json.loads(body) + assert parsed["model"] == "qwen3-coder-30b", ( + f"count_tokens body rewrite missing on LiteLLM route; got " + f"{parsed.get('model')!r} — Claude Code token accounting " + f"will drift from the upstream model on the next request" + ) + + def test_rewrite_runs_after_tool_filtering(self, client): + """``_rewrite_upstream_model`` MUST run AFTER + ``_filter_blocked_tools`` per the plan — if it runs before, a + body that gets re-serialized by the rewrite changes the byte + shape the tool-filter sees, which would silently break the + private-mode tool strip. + + Probe: send a private-mode session with a blocked tool AND a + LiteLLM upstream. Assert both the tool was stripped AND the + model was rewritten in the forwarded body. + """ + if not _rewrite_helper_exists(): + pytest.skip("_rewrite_upstream_model not yet implemented") + captured: dict = {} + fake_registry, _litellm_client, _ = self._build_full_registry_patch(captured) + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch( + "gateway.gateway.get_upstream_registry", + return_value=fake_registry, + create=True, + ), + patch("gateway.gateway.get_anthropic_client"), + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-test" + ) + mock_litellm_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="litellm-key" + ) + + session = _build_mock_session(upstream="litellm", upstream_model="qwen3-coder-30b") + session.mode = "private" # trigger tool-stripping + sm = MagicMock() + sm.get_session_by_ip.return_value = session + mock_sm_get.return_value = sm + + client.post( + "/v1/messages", + data=json.dumps( + { + "model": "opus", + "messages": [], + # WebSearch is blocked in private mode per + # BLOCKED_TOOLS_PRIVATE_MODE — if rewrite runs + # before filter, this tool may survive into + # the forwarded body. + "tools": [ + {"name": "WebSearch"}, + {"name": "Read"}, + ], + } + ), + content_type="application/json", + ) + + body = captured.get("body") + assert body is not None + parsed = json.loads(body) + # Model was rewritten: + assert parsed["model"] == "qwen3-coder-30b" + # Blocked tool was stripped (tool-filter ran): + tool_names = {t.get("name") for t in parsed.get("tools", [])} + assert "WebSearch" not in tool_names, ( + f"Blocked tool 'WebSearch' survived in private mode + " + f"LiteLLM upstream — tool-filter / rewrite ordering " + f"may be broken. Forwarded tools: {tool_names}" + ) + assert "Read" in tool_names, ( + f"Tool-filter stripped a non-blocked tool ('Read') — got {tool_names}" + ) + + +class TestRewriteUpstreamModelMalformedBodyResilience: + """Adversarial probes: the body-rewrite helper MUST NOT crash the + proxy on adversarial inputs (matches the slice-1 + ``_filter_blocked_tools`` resilience contract). + """ + + @pytest.fixture + def client(self): + from gateway.gateway import app + + app.config["TESTING"] = True + with app.test_client() as client: + yield client + + def test_invalid_json_request_does_not_crash_proxy(self, client): + """A LiteLLM session with an invalid-JSON body MUST NOT crash + the proxy — the helper returns the body unchanged and the + upstream gets the malformed body, mirroring today's + Anthropic-route behavior. + """ + try: + import upstream_registry # type: ignore[import-not-found] # noqa: F401 + except ImportError: + pytest.skip("upstream_registry not yet implemented") + + captured: dict = {} + litellm_client = MagicMock() + litellm_client.post.side_effect = _capture_upstream_body(captured) + litellm_client.build_request.return_value = MagicMock() + anthropic_client = MagicMock() + + def _registry_get(upstream): + if upstream == "anthropic": + return (anthropic_client, lambda: MagicMock()) + if upstream == "litellm": + return (litellm_client, lambda: MagicMock()) + raise KeyError(upstream) + + fake_registry = MagicMock() + fake_registry.get.side_effect = _registry_get + + with ( + patch("gateway.gateway.get_credentials_manager") as mock_creds_get, + patch("gateway.gateway.get_session_manager") as mock_sm_get, + patch( + "gateway.gateway.get_upstream_registry", + return_value=fake_registry, + create=True, + ), + patch("gateway.gateway.get_anthropic_client", return_value=anthropic_client), + patch( + "gateway.gateway.get_litellm_credentials_manager", create=True + ) as mock_litellm_get, + ): + mock_creds_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="sk-ant-test" + ) + mock_litellm_get.return_value.get_credential.return_value = MagicMock( + header_name="x-api-key", header_value="litellm-key" + ) + sm = MagicMock() + sm.get_session_by_ip.return_value = _build_mock_session( + upstream="litellm", upstream_model="qwen3-coder-30b" + ) + mock_sm_get.return_value = sm + + # Malformed JSON body — proxy should pass it through + # unchanged, and we should NOT see a 500. + response = client.post( + "/v1/messages", + data=b"not valid json {{", + content_type="application/json", + ) + + # Either the proxy forwarded the malformed body (200 from + # the mock) or returned a structured error — what MUST NOT + # happen is a 500 / unhandled exception. + assert response.status_code != 500, ( + f"Malformed JSON crashed the LiteLLM-routed proxy " + f"(should pass through unchanged like the Anthropic " + f"path); got 500 with {response.data!r}" + ) diff --git a/tests/gateway/test_upstream_registry.py b/tests/gateway/test_upstream_registry.py new file mode 100644 index 0000000000..9dea1a8789 --- /dev/null +++ b/tests/gateway/test_upstream_registry.py @@ -0,0 +1,201 @@ +"""Tests for the gateway UpstreamRegistry (issue #2769 slice-1). + +The registry replaces the lone ``get_anthropic_client()`` singleton at +``gateway/gateway.py:9320`` with a per-upstream-name lookup +(``"anthropic"`` -> the Anthropic httpx client + Anthropic credential +resolver, ``"litellm"`` -> the LiteLLM client + LiteLLM credential +resolver). Slice 1 wires it through ``proxy_anthropic_messages`` and +``proxy_count_tokens`` but no agent yet asks for ``"litellm"``, so the +registry is exercised primarily by these unit tests in slice 1. + +Coverage targets (from plan slice 1 acceptance criteria for TASK-1-1): + +- ``UpstreamRegistry.get("anthropic")`` returns a client with + ``base_url == "https://api.anthropic.com"`` and the existing + Anthropic credential resolver. +- ``UpstreamRegistry.get("litellm")`` returns a client whose + ``base_url`` is sourced from ``LITELLM_BASE_URL`` (default + ``http://litellm.egg-system.svc.cluster.local:4000``) and the + LiteLLM credential resolver. +- ``UpstreamRegistry.get("unknown")`` raises ``UnknownUpstreamError``. +- Both clients share the same timeout / pooling characteristics as + today's ``_anthropic_client`` (regression guard for the singleton's + semantics). +""" + +from __future__ import annotations + +import importlib +import sys +from pathlib import Path + +import httpx +import pytest + +# Add gateway to path for imports — mirror the pattern used by +# tests/gateway/test_anthropic_credentials.py / test_anthropic_proxy.py. +sys.path.insert(0, str(Path(__file__).parent.parent.parent / "gateway")) + + +@pytest.fixture(autouse=True) +def _reset_upstream_registry(monkeypatch): + """Reset the global registry between tests so env-var-driven config + changes (LITELLM_BASE_URL) are observed. + + The registry exposes ``reset_upstream_registry()`` for test + isolation, mirroring ``reset_credentials_manager()`` at + ``gateway/anthropic_credentials.py:226``. + """ + from upstream_registry import reset_upstream_registry + + reset_upstream_registry() + yield + reset_upstream_registry() + + +class TestUpstreamRegistryAnthropic: + """``UpstreamRegistry.get("anthropic")`` keeps today's behavior.""" + + def test_anthropic_returns_client_with_anthropic_base_url(self): + from upstream_registry import get_upstream_registry + + registry = get_upstream_registry() + client, _credential_resolver = registry.get("anthropic") + assert isinstance(client, httpx.Client) + # noqa: EGG200 annotation in production code permits the literal — + # tests assert it explicitly so a future refactor cannot silently + # change the Anthropic upstream URL. + assert str(client.base_url).rstrip("/") == "https://api.anthropic.com" + + def test_anthropic_credential_resolver_is_anthropic(self): + """Anthropic upstream uses the existing AnthropicCredentialsManager.""" + from anthropic_credentials import AnthropicCredentialsManager + from upstream_registry import get_upstream_registry + + registry = get_upstream_registry() + _client, credential_resolver = registry.get("anthropic") + # Resolver is callable / returns an AnthropicCredential or None; + # the resolver instance itself MUST come from the Anthropic + # credentials manager — not the LiteLLM one — so existing + # ANTHROPIC_API_KEY / CLAUDE_CODE_OAUTH_TOKEN behavior is + # untouched. + manager = getattr(credential_resolver, "__self__", None) + assert manager is None or isinstance(manager, AnthropicCredentialsManager), ( + "Anthropic upstream credential resolver must be backed by AnthropicCredentialsManager" + ) + + +class TestUpstreamRegistryLiteLLM: + """``UpstreamRegistry.get("litellm")`` returns a fresh LiteLLM client.""" + + def test_litellm_default_base_url(self, monkeypatch): + """When LITELLM_BASE_URL is unset, defaults to the cluster Service DNS.""" + monkeypatch.delenv("LITELLM_BASE_URL", raising=False) + # Force re-import so the module-level default re-reads the env. + from upstream_registry import ( + get_upstream_registry, + reset_upstream_registry, + ) + + reset_upstream_registry() + registry = get_upstream_registry() + client, _credential_resolver = registry.get("litellm") + assert isinstance(client, httpx.Client) + assert str(client.base_url).rstrip("/") == ( + "http://litellm.egg-system.svc.cluster.local:4000" + ) + + def test_litellm_custom_base_url_from_env(self, monkeypatch): + """LITELLM_BASE_URL overrides the default.""" + monkeypatch.setenv("LITELLM_BASE_URL", "http://litellm-custom:5555") + # Re-import upstream_registry to pick up new env var + if "upstream_registry" in sys.modules: + importlib.reload(sys.modules["upstream_registry"]) + from upstream_registry import ( + get_upstream_registry, + reset_upstream_registry, + ) + + reset_upstream_registry() + registry = get_upstream_registry() + client, _credential_resolver = registry.get("litellm") + assert str(client.base_url).rstrip("/") == "http://litellm-custom:5555" + + def test_litellm_credential_resolver_is_litellm(self): + """LiteLLM upstream uses the LiteLLM credential resolver, not Anthropic.""" + from anthropic_credentials import AnthropicCredentialsManager + from upstream_registry import get_upstream_registry + + registry = get_upstream_registry() + _client, credential_resolver = registry.get("litellm") + manager = getattr(credential_resolver, "__self__", None) + assert manager is None or not isinstance(manager, AnthropicCredentialsManager), ( + "LiteLLM upstream credential resolver must NOT be the AnthropicCredentialsManager" + ) + + +class TestUpstreamRegistryUnknown: + """Unknown upstream names raise a typed error.""" + + def test_unknown_upstream_raises_typed_error(self): + from upstream_registry import ( + UnknownUpstreamError, + get_upstream_registry, + ) + + registry = get_upstream_registry() + with pytest.raises(UnknownUpstreamError): + registry.get("unknown") + + def test_unknown_upstream_error_names_the_upstream(self): + """The error message should name the offending upstream so the + gateway operator can debug a misconfigured session quickly.""" + from upstream_registry import ( + UnknownUpstreamError, + get_upstream_registry, + ) + + registry = get_upstream_registry() + with pytest.raises(UnknownUpstreamError) as exc_info: + registry.get("bogus_upstream_name") + assert "bogus_upstream_name" in str(exc_info.value) + + +class TestUpstreamRegistryClientSemantics: + """Both upstream clients share today's singleton's timeout / pooling + characteristics so neither regresses connection-pool reuse under + concurrent load (issue #1907 retry policy). + """ + + def test_both_clients_are_singletons_within_registry(self): + """Calling ``registry.get("anthropic")`` twice returns the same + client instance — pooling / connection reuse must not be broken + by per-request lookup. + """ + from upstream_registry import get_upstream_registry + + registry = get_upstream_registry() + client_a, _ = registry.get("anthropic") + client_b, _ = registry.get("anthropic") + assert client_a is client_b + + def test_litellm_client_is_singleton(self): + """LiteLLM client is similarly cached so the SSE retry loop's + pool semantics match the Anthropic path.""" + from upstream_registry import get_upstream_registry + + registry = get_upstream_registry() + client_a, _ = registry.get("litellm") + client_b, _ = registry.get("litellm") + assert client_a is client_b + + def test_anthropic_and_litellm_are_distinct_clients(self): + """The two upstreams must NOT share an httpx.Client — different + base_url / credential semantics. + """ + from upstream_registry import get_upstream_registry + + registry = get_upstream_registry() + anthropic_client, _ = registry.get("anthropic") + litellm_client, _ = registry.get("litellm") + assert anthropic_client is not litellm_client