diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_add_autorouter_session_baseline_models/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_add_autorouter_session_baseline_models/migration.sql new file mode 100644 index 00000000000..e7ce1a3180b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_add_autorouter_session_baseline_models/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "baseline_models" JSONB NOT NULL DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 817df082d8c..7d521d54791 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1514,6 +1514,7 @@ model LiteLLM_AutoRouterSession { classifier_cost Float @default(0) classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") + baseline_models Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/litellm/litellm_core_utils/private_json.py b/litellm/litellm_core_utils/private_json.py index 30f64c8fc27..4cd4a9b4f82 100644 --- a/litellm/litellm_core_utils/private_json.py +++ b/litellm/litellm_core_utils/private_json.py @@ -36,6 +36,21 @@ def stage_private_json(path: str, data: Mapping[str, object]) -> str: return tmp_path +def stage_private_bytes(path: str, data: bytes) -> str: + parent: Final = Path(path).parent + parent.mkdir(parents=True, exist_ok=True) + fd, tmp_path = tempfile.mkstemp(dir=str(parent), prefix=".tmp-") + try: + with os.fdopen(fd, "wb") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + except BaseException: + Path(tmp_path).unlink(missing_ok=True) + raise + return tmp_path + + def commit_staged_json(staged: str, path: str) -> None: """Move a staged file into place, replacing whatever is there in one step""" try: @@ -68,3 +83,8 @@ def discard_staged_json(staged: str) -> None: def write_private_json(path: str, data: Mapping[str, object]) -> None: """Atomically write JSON to path with owner-only permissions (0600)""" commit_staged_json(stage_private_json(path, data), path) + + +def write_private_bytes(path: str, data: bytes) -> None: + """Atomically write bytes to path with owner-only permissions (0600); a reader holding the old file keeps it whole""" + commit_staged_json(stage_private_bytes(path, data), path) diff --git a/litellm/models/__init__.py b/litellm/models/__init__.py index 07d1ffa743d..50eb6f8af4f 100644 --- a/litellm/models/__init__.py +++ b/litellm/models/__init__.py @@ -3,6 +3,7 @@ Domain models for LiteLLM backend. """ from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.models.autorouter_session import LiteLLM_AutoRouterSession from litellm.models.budget import ( LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, @@ -40,6 +41,7 @@ __all__ = [ "CredentialBase", "CredentialItem", "LiteLLM_AccessGroupTable", + "LiteLLM_AutoRouterSession", "LiteLLM_BudgetTable", "LiteLLM_BudgetTableFull", "LiteLLM_Config", diff --git a/litellm/models/autorouter_session.py b/litellm/models/autorouter_session.py new file mode 100644 index 00000000000..c7126236ec3 --- /dev/null +++ b/litellm/models/autorouter_session.py @@ -0,0 +1,39 @@ +""" +Auto-router per-session rollup model. + +Canonical definition for ``litellm_autoroutersession``, the row the spend flush +maintains per (api_key, session_id, router_name). +""" + +from collections.abc import Mapping +from datetime import datetime + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_AutoRouterSession(LiteLLMPydanticObjectBase): + api_key: str + session_id: str + router_name: str + router_type: str + first_turn_at: datetime + last_turn_at: datetime + last_model: str + turns: int + spend: float + saved_spend: float + classifier_cost: float + tier_turns: Mapping[str, int] + baseline_models: Mapping[str, int] + + @property + def baseline_model(self) -> str | None: + """The baseline most of this session's turns were priced against, or None when no turn recorded one. + + A router reconfigured mid-session leaves turns priced against two baselines; the row keeps both + counts, and the label is the one that priced the most money-carrying turns rather than whatever the + router is configured with now. + """ + if not self.baseline_models: + return None + return max(self.baseline_models, key=lambda model: (self.baseline_models[model], model)) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 382ea384275..7b0e92d81ae 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -880,6 +880,8 @@ class LiteLLMRoutes(enum.Enum): # proxy admin, or team admin naming their own team via team_id "/auto_router/test_routing", "/auto_router/validate_complexity_router_config", + # Per-session auto-router read - the endpoint scopes the row to the caller's own key hash + "/auto_router/session", # Agent registry - reads are role-scoped and writes are proxy-admin-gated # inside agent_endpoints/endpoints.py *agent_management_routes, diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 4045857a237..2071576a943 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -508,9 +508,9 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_ ### Route Every Claude Code Session Through the Proxy -`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, writes the key it resolved (your fresh `lite login`, or an explicit `--api-key`) into `env.ANTHROPIC_AUTH_TOKEN` as a static token, drops any stray `ANTHROPIC_API_KEY` or `apiKeyHelper` so nothing fights that token, and leaves every other setting in the file untouched. It backs up the original file before patching it. Nothing here writes an `apiKeyHelper`: Claude Code would spawn `lite` (and its keychain check) on every credential refresh, so the key is copied in instead and `lite up` restores the file when it stops. -Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. +Two things need to already be true: you've run `lite login` (or passed a key), and the proxy is already reachable, since `lite up` does not start one for you. ```bash lite login @@ -526,21 +526,19 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi #### Making It Permanent at Login -`lite up` holds the patch only for as long as it runs. To wire Claude Code up once and leave it that way, pass `--config-claude` to `lite login`: +`lite up` holds its patch only for as long as it runs. To wire Claude Code up at login and leave it that way, pass `--config-claude` to `lite login`: ```bash lite --base-url https://your-proxy.example.com login --config-claude ``` -It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: no foreground process to keep alive, and `lite unconfigure claude` restores what it changed (see below). Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. +It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and the key this login minted as `env.ANTHROPIC_AUTH_TOKEN`, but persistently: no foreground process to keep alive, and `lite unconfigure claude` restores what it changed (see below). Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag -Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`. +The key in the file is the login's own, so it expires with it (24h by default): run `lite login --config-claude` again after that, which rewrites the key in place. Earlier versions wrote an `apiKeyHelper` that ran `lite auth print-token` instead, so a later login refreshed Claude Code by itself; that meant Claude Code spawning a full `lite` start, keychain check included, on every credential refresh, so the helper is no longer written and a stale one is stripped by the next `--config-claude` or `configure claude`. Like `lite up`, the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first -Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops. +#### Configuring Claude Code Once, With a Virtual Key -#### Configuring Claude Code Once, With a Virtual Key or Your Login - -`lite configure claude` wires Claude Code up persistently and `lite unconfigure claude` puts things back. It is what `lite login --config-claude` does, plus a pinned model and an undo, and it also takes a long-lived virtual key when that is what you have: +`lite configure claude` wires Claude Code up persistently with a long-lived virtual key, a pinned model and an undo, and `lite unconfigure claude` puts things back: ```bash curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh @@ -548,11 +546,25 @@ lite --base-url https://your-proxy.example.com configure claude --api-key sk-... claude ``` -With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way the command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (under `claude-router-` for a group whose id contains neither `claude` nor `anthropic`, since Claude Code lists only those) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key, which has to be on `/v1/models` for the key. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window (the proxy appends `[1m]` for a group whose configured or known input window reaches 1M) and sends no thinking parameters for it, so name the group like a Claude model id to change that. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control +The key comes from `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) and is written into `env.ANTHROPIC_AUTH_TOKEN`; without one the command refuses, since a `lite login` credential expires within a day and keeping it fresh would mean Claude Code running `lite` through `apiKeyHelper` on every credential refresh. The command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (under `claude-router-` for a group whose id contains neither `claude` nor `anthropic`, since Claude Code lists only those) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key and as `env.ANTHROPIC_MODEL`, both of which have to be on `/v1/models` for the key. The second one matters for `claude -c` and `claude --resume`: a resumed session otherwise re-sends the model its transcript recorded, which behind an auto-router with `return_raw_model_name: true` is the tier model that answered, and a key scoped to the router alias gets a 403 for it; `ANTHROPIC_MODEL` outranks the transcript on resume. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window (the proxy appends `[1m]` for a group whose configured or known input window reaches 1M) and sends no thinking parameters for it, so name the group like a Claude model id to change that. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt -What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Like `--config-claude`, both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any login prompt or request +What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any request + +#### Routed model and savings in the status line + +`lite configure claude`, `lite login --config-claude`, `lite up` and `lite autoroute up` also install a status line (`~/.litellm/statusline.py`, registered as `statusLine` in `~/.claude/settings.json` unless you already run one) that shows which model the auto-router actually served the last turn and, once the proxy has recorded the session, what the session cost against the router's savings baseline: + +``` +claude-auto · Routed to: claude-haiku-4-5 -63% vs Claude Opus 5 +LiteLLM ████████░░░░░░░░░░░░░░░░ $0.14 +Claude Opus 5 ████████████████████████ $0.38 +``` + +The routed model comes from Claude Code's own transcript, so it only names the tier model when the auto-router deployment sets `return_raw_model_name: true` (the `lite autoroute` wizard does); otherwise it shows the alias you requested. The cost lines come from `GET /auto_router/session?session_id=...`, which any virtual key may call for its own sessions, and are cached for five seconds under a per-user `$TMPDIR/litellm-statusline-` directory. The baseline is the priciest model in the router's hardest tier, the same counterfactual the auto-router's savings reports use. `lite unconfigure claude` removes the `statusLine` entry only while it still points at that script. + +`lite codex` registers the same script as a Codex `Stop` hook for the launch, so after each turn Codex prints the same block as a system message. Codex asks once to trust the hook; the answer is remembered for later launches. ### QA Complexity-Based Auto-Routing Against Your Real Proxy diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index bf2a784f590..ea1eed65505 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -1,4 +1,6 @@ +import json import os +import re import shutil import subprocess import sys @@ -13,7 +15,7 @@ import requests from pydantic import BaseModel, TypeAdapter, ValidationError from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login -from .claude_settings import claude_settings_path, lite_api_key_helper_configured +from .claude_settings import ClaudeSettingsError, install_statusline_script from .cmd_quoting import quote_for_cmd from .pi import ( LITELLM_PROXY_API_KEY_ENV, @@ -86,8 +88,6 @@ def build_agent_env( base_url: str, api_key: str, profiles: frozenset[str], - *, - export_anthropic_token: bool = True, ) -> dict[str, str]: """Return a copy of base_env wired to route the agent through the proxy. @@ -102,19 +102,12 @@ def build_agent_env( proxy's /v1/models; likewise left alone when already set. pi ignores both base URL variables and instead resolves $LITELLM_PROXY_API_KEY from its synced models.json provider entry. - - With export_anthropic_token=False the bearer is left out (and any inherited - one dropped) so Claude Code asks its configured apiKeyHelper instead; Claude - Code prefers ANTHROPIC_AUTH_TOKEN over the helper and warns when both are set. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") if PROFILE_ANTHROPIC in profiles: env[ANTHROPIC_BASE_URL_ENV] = root - if export_anthropic_token: - env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key - else: - env.pop(ANTHROPIC_AUTH_TOKEN_ENV, None) + env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key env.pop(ANTHROPIC_API_KEY_ENV, None) if ENABLE_TOOL_SEARCH_ENV not in env: env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE @@ -189,10 +182,52 @@ def prepare_pi( return ("--model", f"{PI_PROVIDER_NAME}/{ids[0]}") +def _warn(message: str) -> None: + click.echo(message, err=True) + + +_CODEX_STOP_HOOKS_DECLARED: Final = re.compile( + r"^\s*(\[\[\s*\"?hooks\"?\s*\.\s*\"?Stop\"?\s*\]\]|\"?hooks\"?(?:\s*\.\s*\"?Stop\"?)?\s*=|\[\s*\"?hooks\"?\s*\])", + re.MULTILINE, +) + + +def codex_config_path(base_env: Mapping[str, str]) -> Path: + return Path(base_env.get("CODEX_HOME") or Path.home() / ".codex") / "config.toml" + + +def codex_declares_stop_hooks(config_path: Path) -> bool: + """A config that cannot be read or decoded declares nothing we can see; Codex reports its own + TOML failure at launch, so the pre-check must not be the thing that stops `lite codex`.""" + try: + return _CODEX_STOP_HOOKS_DECLARED.search(config_path.read_text(encoding="utf-8")) is not None + except (OSError, UnicodeDecodeError): + return False + + +def prepare_codex( + base_url: str, + api_key: str, + base_env: Mapping[str, str], + *, + install: Callable[[], str] = install_statusline_script, + warn: Callable[[str], None] = _warn, +) -> tuple[str, ...]: + """A `-c hooks.Stop=` session flag replaces the user's whole Stop list, so their own hooks win over ours.""" + if codex_declares_stop_hooks(codex_config_path(base_env)): + warn("litellm: your Codex config already declares hooks; not adding the routed-model Stop hook") + return () + try: + command: Final = install() + except ClaudeSettingsError as e: + raise AgentRunError(str(e)) from e + return ("-c", f'hooks.Stop=[{{hooks=[{{type="command",command={json.dumps(command)}}}]}}]') + + _Preparer: TypeAlias = Callable[[str, str, Mapping[str, str]], Sequence[str]] _PREPARERS: Final[Mapping[str, _Preparer]] = MappingProxyType( - {"pi": prepare_pi} # mutable-ok: MappingProxyType freezes the provider registry + {"pi": prepare_pi, "codex": prepare_codex} # mutable-ok: MappingProxyType freezes the provider registry ) @@ -454,10 +489,6 @@ def _restore_controlling_terminal() -> None: os.close(fd) -def _warn(message: str) -> None: - click.echo(message, err=True) - - def run_agent( base_url: str, api_key: str, @@ -474,7 +505,6 @@ def run_agent( launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off, reattach_terminal: Callable[[], None] | None = None, preparers: Mapping[str, _Preparer] = MappingProxyType(_PREPARERS), - export_anthropic_token: bool = True, ) -> None: """Validate, wire the environment, and hand off to the agent. @@ -506,9 +536,7 @@ def run_agent( env: Final = MappingProxyType( { - **build_agent_env( - env_before_sync, base_url, api_key, profiles, export_anthropic_token=export_anthropic_token - ), + **build_agent_env(env_before_sync, base_url, api_key, profiles), **(_NO_EXTRA_ENV if isinstance(synced, ModelSyncSkipped) else synced), } ) @@ -546,26 +574,14 @@ def resolve_api_key(ctx: click.Context) -> str: _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy." -def _helper_supplies_token( - ctx_obj: CliContextObj, base_url: str, profiles: frozenset[str], settings_path: Path -) -> bool: - if PROFILE_ANTHROPIC not in profiles or not ctx_obj.get("api_key_from_token_file"): - return False - return lite_api_key_helper_configured(base_url, settings_path) - - def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] started_interactive: Final = _is_interactive() api_key: Final = resolve_api_key(ctx) - display_name, profiles = agent_profile(binary) - settings_path: Final = claude_settings_path(os.environ) - helper_supplies_token: Final = _helper_supplies_token(ctx_obj, base_url, profiles, settings_path) + display_name, _profiles = agent_profile(binary) click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}") - if helper_supplies_token: - click.echo(f"litellm: {display_name} reads its key from the apiKeyHelper in {settings_path}") try: run_agent( @@ -574,7 +590,6 @@ def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify [binary, *args], skip_verify=skip_verify, reattach_terminal=(_restore_controlling_terminal if started_interactive else None), - export_anthropic_token=not helper_supplies_token, ) except AgentRunError as e: raise click.ClickException(str(e)) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 4f704afe9d6..98af32fa7aa 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -42,14 +42,13 @@ from litellm.litellm_core_utils.cli_token_utils import ( from .claude_settings import ( STARTING_MODEL_ROLE, - ApiKeyHelper, ClaudeSettingsError, KeepModel, + StaticToken, claude_settings_path, configure_claude_settings, configure_state_path, refuse_while_owned, - resolve_api_key_helper, settings_file_owners, ) from .pkce_login import ( @@ -784,13 +783,16 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None: return None -def _configure_claude_code(base_url: str) -> None: - """Point Claude Code at base_url by patching the settings.json it reads, undoable with `lite unconfigure claude`.""" +def _configure_claude_code(base_url: str, api_key: str) -> None: + """Write the key this login just minted into Claude Code's settings.json as a static token, undoable with + `lite unconfigure claude`. The key expires with the login, so the flag is the re-wire step of each login + rather than a one-time setup: no apiKeyHelper is written, since Claude Code would spawn `lite` (and its + keychain probe) on every credential refresh to keep one fresh.""" settings_path: Final = claude_settings_path(os.environ) try: configure_claude_settings( base_url, - ApiKeyHelper(resolve_api_key_helper(base_url)), + StaticToken(api_key), KeepModel(), settings_path, configure_state_path(settings_path), @@ -800,22 +802,28 @@ def _configure_claude_code(base_url: str) -> None: raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}") click.echo(f"\nConfigured Claude Code: {settings_path} now routes through {base_url.rstrip('/')}.") click.echo( + "This login's key is stored in the file, so run `lite login --config-claude` again after it expires. " "Your other Claude Code settings were left untouched. Restart Claude Code to pick this up. " f"Undo with `lite unconfigure claude`; `lite configure claude --model` sets {STARTING_MODEL_ROLE}." ) def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: SecretSave) -> None: + """Claude Code is configured from the key in hand, so it does not wait on the CLI's own store: a login whose + token file or keychain refused it still has a usable key, and `--config-claude` asked for exactly that + key to be written into settings.json.""" from litellm.proxy.client.cli.interface import show_commands click.echo("\nLogin successful!") click.echo(f"JWT Token: {api_key[:20]}...") click.echo(storage_notice(stored)) + if config_claude: + _configure_claude_code(base_url, api_key) if isinstance(stored, (CredentialNotSaved, CredentialNotRecorded)): + if config_claude: + click.echo("Claude Code was configured with this key even though the CLI itself could not keep it.") return click.echo("You can now use the CLI without specifying --api-key") - if config_claude: - _configure_claude_code(base_url) click.echo("\n" + "=" * 60) show_commands() @@ -850,8 +858,8 @@ def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault) -> None: is_flag=True, default=False, help=( - "After logging in, update ~/.claude/settings.json so Claude Code routes through this proxy. " - "Unrelated settings are preserved." + "After logging in, write this login's key into ~/.claude/settings.json so Claude Code routes through " + "this proxy; run it again after the key expires. Unrelated settings are preserved." ), ) @click.option( diff --git a/litellm/proxy/client/cli/commands/autoroute/commands.py b/litellm/proxy/client/cli/commands/autoroute/commands.py index 86701f186fd..5d91fc81350 100644 --- a/litellm/proxy/client/cli/commands/autoroute/commands.py +++ b/litellm/proxy/client/cli/commands/autoroute/commands.py @@ -1,5 +1,4 @@ import atexit -import json import secrets import signal import threading @@ -15,8 +14,10 @@ from ..claude_settings import ( CLAUDE_SETTINGS_PATH, ClaudeSettingsError, StaticToken, + install_statusline_script, load_json_or_empty, merge_claude_settings, + write_claude_settings, ) from ..up import BackupRecord as ClaudeBackupRecord from ..up import restore_claude_settings, write_backup @@ -151,6 +152,7 @@ def up(port: int) -> None: raise click.ClickException(str(e)) try: + status_line: Final = install_statusline_script() original_existed: Final = CLAUDE_SETTINGS_PATH.exists() original_settings: Final = load_json_or_empty(CLAUDE_SETTINGS_PATH) write_backup( @@ -158,11 +160,15 @@ def up(port: int) -> None: AUTOROUTE_BACKUP_PATH, ) merged: Final = merge_claude_settings( - original_settings, base_url, StaticToken(master_key), AUTOROUTER_MODEL_NAME, AUTOROUTER_MODEL_NAME + original_settings, + base_url, + StaticToken(master_key), + AUTOROUTER_MODEL_NAME, + AUTOROUTER_MODEL_NAME, + status_line=status_line, ) CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True) - with secure_create(CLAUDE_SETTINGS_PATH) as f: - json.dump(merged, f, indent=2) + write_claude_settings(CLAUDE_SETTINGS_PATH, merged) except ClaudeSettingsError as e: terminate(process.pid) clear_pid_record() diff --git a/litellm/proxy/client/cli/commands/autoroute/config.py b/litellm/proxy/client/cli/commands/autoroute/config.py index 9dfc4ad079b..1f3ad34e3d9 100644 --- a/litellm/proxy/client/cli/commands/autoroute/config.py +++ b/litellm/proxy/client/cli/commands/autoroute/config.py @@ -162,6 +162,7 @@ def build_generated_model_list(config: AutorouteConfig) -> list[JsonValue]: complexity_router_config: Final[dict[str, JsonValue]] = { "tiers": {tier: list(models) for tier, models in config.tiers.items()}, "default_model": config.default_model, + "return_raw_model_name": True, } if isinstance(config.classifier, LLMClassifier): complexity_router_config["classifier_type"] = "llm" diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index d0ee507f0b2..e6231f3cac9 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -1,16 +1,18 @@ """Shared handling of Claude Code's ~/.claude/settings.json. `lite up` and `lite autoroute up` patch this file temporarily and restore it on -exit; `lite login --config-claude` and `lite configure claude` patch it -persistently and record how to undo it. All of them need the same merge and the -same apiKeyHelper command, and `up` already imports from `auth`, so the shared -parts live here rather than in any one command module. +exit; `lite configure claude` patches it persistently and records how to undo it. +All of them need the same merge, and `up` already imports from `auth`, so the +shared parts live here rather than in any one command module. The credential is +always a static token in `env.ANTHROPIC_AUTH_TOKEN`: Claude Code's `apiKeyHelper` +would spawn a `lite` process on every credential refresh, and that process touches +the keychain, so nothing here writes one; a helper left by an earlier version is +owned like any other key and stripped. """ import hashlib import json import shlex -import shutil import sys from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass @@ -27,13 +29,16 @@ from litellm.litellm_core_utils.private_json import ( discard_staged_json, ensure_private_dir, stage_private_json, + write_private_bytes, ) +from . import statusline_script from .cmd_quoting import quote_for_cmd ENV_KEY: Final = "env" API_KEY_HELPER_KEY: Final = "apiKeyHelper" MODEL_KEY: Final = "model" +STATUS_LINE_KEY: Final = "statusLine" ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN" ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" @@ -41,6 +46,7 @@ ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" ENABLE_TOOL_SEARCH_VALUE: Final = "true" ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1" +ANTHROPIC_MODEL_KEY: Final = "ANTHROPIC_MODEL" ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = ( "ANTHROPIC_DEFAULT_SONNET_MODEL", "ANTHROPIC_DEFAULT_HAIKU_MODEL", @@ -53,19 +59,22 @@ OWNED_ENV_KEYS: Final = ( ANTHROPIC_BASE_URL_KEY, ANTHROPIC_AUTH_TOKEN_KEY, ANTHROPIC_API_KEY_KEY, + ANTHROPIC_MODEL_KEY, ) -OWNED_TOP_LEVEL_KEYS: Final = (API_KEY_HELPER_KEY, MODEL_KEY) +OWNED_TOP_LEVEL_KEYS: Final = (API_KEY_HELPER_KEY, MODEL_KEY, STATUS_LINE_KEY) OWNED_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in OWNED_ENV_KEYS), *OWNED_TOP_LEVEL_KEYS) _CREDENTIAL_ENV_KEYS: Final = frozenset((ANTHROPIC_API_KEY_KEY, ANTHROPIC_AUTH_TOKEN_KEY)) _CREDENTIAL_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in sorted(_CREDENTIAL_ENV_KEYS)), API_KEY_HELPER_KEY) _BASE_URL_PATH: Final = f"{ENV_KEY}.{ANTHROPIC_BASE_URL_KEY}" -STARTING_MODEL_ROLE: Final = "the /model picker's default row, the model Claude Code starts on" +_MODEL_PATHS: Final = (MODEL_KEY, f"{ENV_KEY}.{ANTHROPIC_MODEL_KEY}") +STARTING_MODEL_ROLE: Final = "the /model picker's default row, the model Claude Code starts and resumes on" CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" CLAUDE_CONFIG_DIR_ENV: Final = "CLAUDE_CONFIG_DIR" BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json" CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json" +STATUSLINE_SCRIPT_PATH: Final = Path.home() / ".litellm" / "statusline.py" @dataclass(frozen=True, slots=True) @@ -123,16 +132,6 @@ class StaticToken: token: str -@dataclass(frozen=True, slots=True) -class ApiKeyHelper: - """A `lite auth print-token` command Claude Code runs per request, so a login renews in place.""" - - command: str - - -ClaudeCredential: TypeAlias = StaticToken | ApiKeyHelper - - @dataclass(frozen=True, slots=True) class KeepModel: """Leave the top-level `model` as it is, the user's or an earlier configure's (a re-login).""" @@ -145,7 +144,9 @@ class UnpinModel: @dataclass(frozen=True, slots=True) class StartOn: - """Pin the top-level `model`, the row Claude Code starts on.""" + """Pin `model` and `env.ANTHROPIC_MODEL`: the row Claude Code starts on, and the one a resumed session stays + on, since resume otherwise re-sends the transcript's served model, which a raw-model router made a tier + model the key may not reach.""" model: str @@ -259,6 +260,17 @@ def _write_target(settings_path: Path) -> Path: raise ClaudeSettingsError(f"Could not resolve {settings_path}: {e}") from e +def write_claude_settings(settings_path: Path, settings: Mapping[str, JsonValue]) -> None: + """The one way a settings document lands on disk: staged owner-only beside the target and renamed into + place, through a symlink rather than over it. Every writer (`configure`, `up`, `autoroute up` and the + restores) may be carrying the credential, so none creates the file under the umask or truncates it.""" + target: Final = _write_target(settings_path) + try: + commit_staged_json(stage_private_json(str(target), settings), str(target)) + except OSError as e: + raise ClaudeSettingsError(f"Could not write {settings_path}: {e}") from e + + def _stage(path: Path, document: Mapping[str, object]) -> str: try: return stage_private_json(str(path), document) @@ -287,20 +299,49 @@ def _land( raise ClaudeSettingsError(f"Could not {'remove' if staged is None else 'write'} {path}: {e}") from e +def statusline_command(script_path: Path, platform: str = sys.platform) -> str: + """This interpreter, not a bare `python3`: it is the one the apiKeyHelper already depends on.""" + quote: Final = quote_for_cmd if platform.startswith("win") else shlex.quote + return " ".join(quote(token) for token in (sys.executable, str(script_path))) + + +def install_statusline_script(script_path: Path | None = None) -> str: + target: Final = script_path or STATUSLINE_SCRIPT_PATH + try: + ensure_private_dir(target.parent) + write_private_bytes(str(target), Path(statusline_script.__file__).read_bytes()) + except OSError as e: + raise ClaudeSettingsError(f"Could not install the status line script at {target}: {e}") from e + return statusline_command(target) + + +def with_status_line(settings: Mapping[str, JsonValue], command: str) -> Mapping[str, JsonValue]: + """Ours is recognised by the script it runs, so a re-install under another interpreter is still ours.""" + existing: Final = settings.get(STATUS_LINE_KEY) + existing_command: Final = existing.get("command") if isinstance(existing, dict) else None + ours: Final = existing is None or (isinstance(existing_command, str) and command.split()[-1] in existing_command) + if not ours: + return settings + entry: Final = dict((("type", "command"), ("command", command))) # mutable-ok: JSON document + return dict(chain(settings.items(), ((STATUS_LINE_KEY, entry),))) # mutable-ok: JSON document + + def merge_claude_settings( settings: Mapping[str, JsonValue], base_url: str, - credential: ClaudeCredential, + credential: StaticToken, default_model: str | None = None, tier_model: str | None = None, + *, + status_line: str | None = None, ) -> Mapping[str, JsonValue]: """Return a new settings mapping wired to route Claude Code through the proxy. - A StaticToken lands in env.ANTHROPIC_AUTH_TOKEN, an ApiKeyHelper in the top-level apiKeyHelper; - the other credential slots are removed either way, since Claude Code given two credentials may - send the wrong one. ENABLE_TOOL_SEARCH and CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY get their - defaults only when missing. `default_model` is the top-level `model`, the row Claude Code starts - on; `tier_model` is `lite autoroute up`'s knob that points every ANTHROPIC_DEFAULT_*_MODEL at one + The token lands in env.ANTHROPIC_AUTH_TOKEN; the other credential slots (a stray ANTHROPIC_API_KEY, + an apiKeyHelper) are removed, since Claude Code given two credentials may send the wrong one. + ENABLE_TOOL_SEARCH and CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY get their defaults only when + missing. `default_model` is the top-level `model` and env.ANTHROPIC_MODEL (see StartOn); + `tier_model` is `lite autoroute up`'s knob that points every ANTHROPIC_DEFAULT_*_MODEL at one group. Apart from those tier keys, exactly OWNED_PATHS are touched. """ raw_env: Final = settings.get(ENV_KEY, {}) @@ -312,60 +353,24 @@ def merge_claude_settings( (ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE), ), ((key, value) for key, value in current_env.items() if key not in _CREDENTIAL_ENV_KEYS), - ((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")),), - ((ANTHROPIC_AUTH_TOKEN_KEY, credential.token),) if isinstance(credential, StaticToken) else (), + ((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")), (ANTHROPIC_AUTH_TOKEN_KEY, credential.token)), + ((ANTHROPIC_MODEL_KEY, default_model),) if default_model is not None else (), ((key, tier_model) for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS if tier_model is not None), ) ) return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping chain( - ((key, value) for key, value in settings.items() if key not in (API_KEY_HELPER_KEY, ENV_KEY)), + ( + (key, value) + for key, value in (with_status_line(settings, status_line) if status_line else settings).items() + if key not in (API_KEY_HELPER_KEY, ENV_KEY) + ), ((ENV_KEY, env),), - ((API_KEY_HELPER_KEY, credential.command),) if isinstance(credential, ApiKeyHelper) else (), ((MODEL_KEY, default_model),) if default_model is not None else (), ) ) -def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str: - """Build the shell command Claude Code should run for its apiKeyHelper. - - Claude Code hands the string to the system shell, `sh` on POSIX and cmd.exe - on Windows, so every token is quoted for the shell that will read it. - - Resolves `lite` to an absolute path so the helper works regardless of the - PATH visible to whatever subprocess Claude Code spawns it from. Passing - --base-url explicitly (rather than relying on the bare invocation Claude - Code would otherwise use) makes `print-token` enforce that the cached - token was actually issued for this proxy -- without it, a token minted - for a different, previously-logged-into proxy would be handed to - whichever server the settings currently point at. - - --base-url belongs to the top-level `lite` group, so it has to precede the - subcommand; click rejects it outright after `print-token`. - """ - lite_path: Final = shutil.which("lite") - if lite_path is None: - raise ClaudeSettingsError( - "Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs an absolute path to it." - ) - quote: Final = quote_for_cmd if platform.startswith("win") else shlex.quote - return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token")) - - -def lite_api_key_helper_configured(base_url: str, settings_path: Path) -> bool: - """Whether settings_path already carries the apiKeyHelper `lite login --config-claude` writes for base_url. - - Only an exact match counts: a helper for another proxy, a hand-written one, or - settings that cannot be read leave the caller on the env-token path. - """ - try: - configured_helper: Final = load_json_or_empty(settings_path).get(API_KEY_HELPER_KEY) - return configured_helper == resolve_api_key_helper(base_url.rstrip("/")) - except ClaudeSettingsError: - return False - - def _owned(container: Mapping[str, JsonValue], key: str) -> OwnedValue: return OwnedValue(present=key in container, value=container.get(key)) @@ -460,12 +465,13 @@ def read_configure_receipt(state_path: Path) -> ConfigureReceipt | None: def configure_claude_settings( base_url: str, - credential: ClaudeCredential, + credential: StaticToken, model: ModelChoice, settings_path: Path, state_path: Path, owners: Sequence[SettingsFileOwner], commit: Callable[[str, str], None] = commit_staged_json, + script_path: Path | None = None, ) -> None: """Persistently route Claude Code through base_url, recording how to undo it. @@ -474,19 +480,28 @@ def configure_claude_settings( discards the staged settings, and a settings rename that fails after the receipt landed puts the earlier receipt back (or removes the new one), so the receipt on disk never describes settings that were not written. `model`: StartOn pins the starting model, UnpinModel lets go of a pin an - earlier configure made (never of the user's own), KeepModel leaves it alone (a re-login). + earlier configure made (never of the user's own), KeepModel leaves it alone (a re-login). The + status line script is installed and registered under `statusLine` unless the user runs their own; + the receipt owns that key like any other, so unconfigure removes only ours. """ refuse_while_owned(settings_path, owners) current: Final = load_json_or_empty(settings_path) _env_object(current, settings_path) earlier: Final = read_configure_receipt(state_path) - existing: Final = ( - _with(current, MODEL_KEY, earlier.previous[MODEL_KEY]) - if isinstance(model, UnpinModel) and earlier is not None and _ours(current, MODEL_KEY, earlier) - else current + unpinned: Final = MappingProxyType( + { + path: earlier.previous[path] + for path in _MODEL_PATHS + if isinstance(model, UnpinModel) and earlier is not None and _ours(current, path, earlier) + } ) + existing: Final = _with_all(current, unpinned) merged: Final = merge_claude_settings( - existing, base_url, credential, model.model if isinstance(model, StartOn) else None + existing, + base_url, + credential, + model.model if isinstance(model, StartOn) else None, + status_line=install_statusline_script(script_path), ) receipt: Final = _receipt(current, merged, earlier, settings_path.exists()) target: Final = _write_target(settings_path) @@ -581,6 +596,7 @@ __all__ = ( "ANTHROPIC_AUTH_TOKEN_KEY", "ANTHROPIC_BASE_URL_KEY", "ANTHROPIC_DEFAULT_MODEL_ENV_KEYS", + "ANTHROPIC_MODEL_KEY", "API_KEY_HELPER_KEY", "AUTOROUTE_BACKUP_PATH", "BACKUP_PATH", @@ -598,8 +614,8 @@ __all__ = ( "OWNED_TOP_LEVEL_KEYS", "SETTINGS_FILE_OWNERS", "STARTING_MODEL_ROLE", - "ApiKeyHelper", - "ClaudeCredential", + "STATUSLINE_SCRIPT_PATH", + "STATUS_LINE_KEY", "ClaudeSettingsError", "ConfigureReceipt", "KeepModel", @@ -614,12 +630,11 @@ __all__ = ( "claude_settings_path", "configure_claude_settings", "configure_state_path", - "lite_api_key_helper_configured", "load_json_or_empty", "merge_claude_settings", "read_configure_receipt", "refuse_while_owned", - "resolve_api_key_helper", "settings_file_owners", "unconfigure_claude_settings", + "write_claude_settings", ) diff --git a/litellm/proxy/client/cli/commands/configure.py b/litellm/proxy/client/cli/commands/configure.py index 9d329f8d8f2..4acf94e16f9 100644 --- a/litellm/proxy/client/cli/commands/configure.py +++ b/litellm/proxy/client/cli/commands/configure.py @@ -18,11 +18,9 @@ from litellm.proxy.common_utils.model_listing_utils import ( GATEWAY_CLIENT_HEADER, ) -from .auth import CliContextObj, context_secret_vault, get_stored_api_key +from .auth import CliContextObj from .claude_settings import ( STARTING_MODEL_ROLE, - ApiKeyHelper, - ClaudeCredential, ClaudeSettingsError, ModelChoice, StartOn, @@ -33,12 +31,10 @@ from .claude_settings import ( configure_claude_settings, configure_state_path, refuse_while_owned, - resolve_api_key_helper, settings_file_owners, unconfigure_claude_settings, ) from .pi import ListedModel, ListingFailure, PiSyncError, fetch_model_listing -from .up import ensure_fresh_login _LISTED_MODELS_SHOWN: Final = 20 _CLAUDE_TARGET: Final = "claude" @@ -54,24 +50,21 @@ _MODEL_OPTION_HELP: Final = ( ) -def resolve_credential(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, str]: - """The credential to write and the key to check the proxy with. +def resolve_credential(ctx: click.Context, api_key: str | None) -> StaticToken: + """The long-lived key written into settings.json: --api-key, `lite --api-key` or LITELLM_PROXY_API_KEY. - An explicit key (--api-key, `lite --api-key`, LITELLM_PROXY_API_KEY) is long-lived and goes - into settings.json as a static token. Without one, the stored `lite login` credential is used - the way `lite login --config-claude` uses it, through apiKeyHelper, since it expires within a - day and renews in place there; a missing or stale login is refreshed first, as `lite up` does. + A `lite login` credential is never written: it expires within a day, and keeping it fresh would mean + Claude Code running `lite` through `apiKeyHelper` on every credential refresh. """ ctx_obj: Final[CliContextObj] = ctx.obj explicit: Final = api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key")) - if explicit: - return StaticToken(explicit), explicit - base_url: Final = ctx_obj["base_url"] - ensure_fresh_login(ctx) - stored: Final = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx)) - if not stored: - raise ClaudeSettingsError("Login did not produce a usable token.") - return ApiKeyHelper(resolve_api_key_helper(base_url)), stored + if not explicit: + raise ClaudeSettingsError( + "`lite configure claude` needs a long-lived virtual key: pass --api-key, `lite --api-key`, or set " + "LITELLM_PROXY_API_KEY. Your `lite login` credential expires within a day, so it is not written " + "into Claude Code's settings." + ) + return StaticToken(explicit) @dataclass(frozen=True, slots=True) @@ -83,22 +76,22 @@ class _Listing: return tuple(model.id for model in self.models) -def _start(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, _Listing]: +def _start(ctx: click.Context, api_key: str | None) -> tuple[StaticToken, _Listing]: """Every configure path begins the same way: the local ownership check first, so a `lite up` - session is refused before any login prompt or request, then the credential, then the listing.""" + session is refused before any request, then the credential, then the listing.""" settings_path: Final = claude_settings_path(os.environ) try: refuse_while_owned(settings_path, settings_file_owners(settings_path)) - credential, key = resolve_credential(ctx, api_key) + credential: Final = resolve_credential(ctx, api_key) except ClaudeSettingsError as e: raise click.ClickException(str(e)) - return credential, _listed_models(ctx.obj["base_url"], key) + return credential, _listed_models(ctx.obj["base_url"], credential.token) def _listing_error(base_url: str, error: PiSyncError) -> str: """The hint that fits how the listing failed: only an unreachable proxy gets the "is it running" question.""" if error.kind is ListingFailure.REJECTED: - return f"LiteLLM rejected your key (HTTP {error.status}). Run `lite login` to refresh it, or pass a valid --api-key." + return f"LiteLLM rejected your key (HTTP {error.status}). Pass a valid --api-key." if error.kind is ListingFailure.UNREACHABLE: return f"{error.message} Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?" if error.kind is ListingFailure.EMPTY: @@ -122,7 +115,7 @@ def _model_choice(model: str | None) -> ModelChoice: return StartOn(model) if model is not None else UnpinModel() -def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listing: _Listing, model: str | None) -> None: +def _apply_claude(ctx: click.Context, credential: StaticToken, listing: _Listing, model: str | None) -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] listed: Final = listing.ids @@ -148,16 +141,13 @@ def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listing: _Li in_picker: Final = sum(1 for listed_model in listed if CLAUDE_CODE_PICKER_PATTERN.search(listed_model)) click.echo(f"Configured Claude Code: {settings_path} now routes through {base_url}.") - click.echo( - "Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN." - if isinstance(credential, StaticToken) - else "Credential: your `lite login`, read through apiKeyHelper on every request, so a later login renews it." - ) + click.echo("Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN.") click.echo( f"Starting model: {starting} ({STARTING_MODEL_ROLE}); switch any time with /model." if starting is not None else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or " - "pass --model to start on a proxy model." + "pass --model to start on a proxy model. Without a pin, a resumed session re-sends the model its transcript " + "recorded, which behind a raw-model auto-router is the tier model." ) click.echo( f"/model will list all {len(listed)} of the proxy's models." @@ -166,7 +156,7 @@ def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listing: _Li "'claude' or 'anthropic', and this proxy does not list the rest under such names." ) click.echo("Start `claude` from any terminal. Undo with `lite unconfigure claude`.") - if isinstance(credential, StaticToken) and settings_path.is_symlink(): + if settings_path.is_symlink(): click.echo( f"Note: {settings_path} is a symlink to {settings_path.resolve()}, so your key now lives in " "that file; keep it out of version control.", @@ -235,16 +225,16 @@ def unconfigure_group() -> None: "api_key", default=None, help="Long-lived LiteLLM virtual key written into Claude Code's settings. Defaults to the `lite --api-key` / " - "LITELLM_PROXY_API_KEY value; with neither, your `lite login` credential is used through apiKeyHelper.", + "LITELLM_PROXY_API_KEY value; required, since a `lite login` credential expires within a day.", ) @click.option("--model", default=None, help=_MODEL_OPTION_HELP) @click.pass_context def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) -> None: """Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`. - Patches ~/.claude/settings.json in place: the proxy URL, your credential (a virtual key as a - static token, or your `lite login` through apiKeyHelper), and gateway model discovery so - /model lists the proxy's models; --model picks the one Claude Code starts on. Every other + Patches ~/.claude/settings.json in place: the proxy URL, your virtual key as a static token, + and gateway model discovery so /model lists the proxy's models; --model picks the one Claude + Code starts on and resumes with. Every other setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back. Assumes the proxy is already running. """ diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py new file mode 100644 index 00000000000..a8abeb68978 --- /dev/null +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -0,0 +1,392 @@ +"""Claude Code status line and Codex Stop hook for auto-routed sessions. + +`lite` copies this file verbatim to ~/.litellm/statusline.py and registers it as Claude +Code's `statusLine` command and as Codex's `[[hooks.Stop]]` command, so it must stay +standard-library only and must never import litellm. Claude Code re-runs it on every +status refresh (about every 300ms while typing), so the proxy is asked at most once per +TTL per session and every other refresh is served from a small on-disk cache that holds +only the proxy's answer, never the key. + +Claude Code pipes a JSON payload on stdin (session_id, transcript_path, model); the routed +model is the `message.model` of the latest foreground assistant line in the transcript, +which is the proxy's response `model` field. That only names the tier model when the +auto-router deployment sets `return_raw_model_name: true`; otherwise it is the alias the +client requested. Codex pipes its Stop event instead (hook_event_name, session_id) and has +no transcript to read, so the routed model comes from the proxy's session record and the +result is printed as a `systemMessage` for the transcript. The proxy key is read from the +agent's own environment (the static token `lite configure claude` writes); nothing here +spawns a credential helper. + +Cost figures come from GET /auto_router/session on the proxy, which reads the per-session +rollup written by the spend flush. That flush is asynchronous, so a turn's cost lands a +second or two after the turn; the cache TTL absorbs it. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import sys +import tempfile +import time +import urllib.error +import urllib.request +from collections.abc import Callable, Mapping +from pathlib import Path +from types import MappingProxyType +from typing import IO, Final, NamedTuple, Protocol +from urllib.parse import urlencode + +SESSION_ENDPOINT: Final = "/auto_router/session" +CACHE_TTL_SECONDS: Final = 5.0 +FETCH_TIMEOUT_SECONDS: Final = 3 +BAR_WIDTH: Final = 24 +BAR_FULL: Final = "\u2588" +BAR_EMPTY: Final = "\u2591" +SEPARATOR: Final = " \u00b7 " +TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024 +CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",) +CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY") +CODEX_BASE_URL_ENV_KEYS: Final = ("OPENAI_BASE_URL",) +CODEX_API_KEY_ENV_KEYS: Final = ("OPENAI_API_KEY",) +CODEX_STOP_EVENT: Final = "Stop" +SYNTHETIC_MODEL: Final = "" +LITELLM_LABEL: Final = "LiteLLM" +RESET: Final = "\033[0m" +BOLD: Final = "\033[1m" +DIM: Final = "\033[90m" +LITELLM_COLOR: Final = "\033[38;2;79;70;229m" +BASELINE_COLOR: Final = "\033[38;2;217;119;87m" +EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) +EMPTY_ENV: Final[Mapping[str, str]] = MappingProxyType({}) + + +class Session(NamedTuple): + router_name: str + last_model: str + spend: float + baseline_spend: float + baseline_model: str | None + + +class Credentials(NamedTuple): + base_url: str + api_key: str + + @property + def usable(self) -> bool: + return bool(self.base_url and self.api_key) + + +class Fetched(NamedTuple): + session: Session | None + definitive: bool + + +class Fetch(Protocol): + def __call__(self, credentials: Credentials, session_id: str) -> Fetched: ... + + +def as_mapping(value: object) -> Mapping[str, object]: + return value if isinstance(value, dict) else EMPTY + + +def as_str(value: object) -> str: + return value if isinstance(value, str) else "" + + +def printable(value: object) -> str: + """Labels come from the transcript, the proxy, and Claude Code's model cache, none of which this script + controls, and every one is written to a terminal: a control character (ESC, BEL, C1) in a model name + could redraw the screen or set the clipboard, so only printable text survives.""" + return "".join(character for character in as_str(value) if character.isprintable()) + + +def load_json(raw: bytes | str) -> object: + try: + return json.loads(raw) + except ValueError: + return None + + +def resolve_base_url(env: Mapping[str, str], keys: tuple[str, ...]) -> str: + raw: Final = next((env[key] for key in keys if env.get(key)), "").strip().rstrip("/") + return raw.removesuffix("/v1") + + +def resolve_api_key(env: Mapping[str, str], keys: tuple[str, ...]) -> str: + return next((env[key] for key in keys if env.get(key)), "").strip() + + +def claude_credentials(env: Mapping[str, str]) -> Credentials: + """Claude Code's own resolution order, so the key-scoped lookup runs as the principal that wrote the rows: + ANTHROPIC_AUTH_TOKEN, then ANTHROPIC_API_KEY. A `lite` variable such as LITELLM_PROXY_API_KEY is not a + key Claude Code ever sends, so honoring it would ask as someone else. An apiKeyHelper is never run: a + status line refreshes every few hundred milliseconds, and spawning a credential helper that often is + how a keychain prompt ends up on screen a hundred times.""" + return Credentials(resolve_base_url(env, CLAUDE_BASE_URL_ENV_KEYS), resolve_api_key(env, CLAUDE_API_KEY_ENV_KEYS)) + + +def codex_credentials(env: Mapping[str, str]) -> Credentials: + return Credentials(resolve_base_url(env, CODEX_BASE_URL_ENV_KEYS), resolve_api_key(env, CODEX_API_KEY_ENV_KEYS)) + + +def _transcript_line_model(line: bytes) -> str: + """A `` model is Claude Code's own marker for a locally produced message (an API error, a + resume note), not a served model, so it is skipped like a sidechain line.""" + item: Final = as_mapping(load_json(line)) + if item.get("type") != "assistant" or item.get("isSidechain") is True or item.get("agentId"): + return "" + model: Final = printable(as_mapping(item.get("message")).get("model")) + return "" if model == SYNTHETIC_MODEL else model + + +def latest_transcript_model(transcript_path: str) -> str: + if not transcript_path: + return "" + try: + with Path(transcript_path).open("rb") as transcript: + size: Final = transcript.seek(0, os.SEEK_END) + transcript.seek(max(0, size - TRANSCRIPT_SCAN_LIMIT_BYTES)) + tail: Final = transcript.read() + except OSError: + return "" + return next((model for line in reversed(tail.split(b"\n")) if (model := _transcript_line_model(line))), "") + + +def model_label(model: str, config_dir: Path) -> str: + bare: Final = model.rsplit("/", 1)[-1] + try: + raw: Final = (config_dir / "cache" / "gateway-models.json").read_bytes() + except OSError: + return bare + listed: Final = as_mapping(load_json(raw)).get("models") + if not isinstance(listed, list): + return bare + entries: Final = tuple(as_mapping(entry) for entry in listed) + return next( + ( + printable(entry.get("display_name")) + for entry in entries + if entry.get("id") in (model, bare) and printable(entry.get("display_name")) + ), + bare, + ) + + +def baseline_label(model: str, config_dir: Path) -> str: + labelled: Final = model_label(model, config_dir) + if labelled != model.rsplit("/", 1)[-1]: + return labelled + return " ".join(word.capitalize() for word in labelled.replace("-", " ").split()) + + +def fetch_session(credentials: Credentials, session_id: str) -> Fetched: + """Any 4xx is this credential's definite answer (no row, no access, expired login) and is cached for the + TTL; a 5xx or transport failure is not, so the next refresh tries again.""" + query: Final = urlencode((("session_id", session_id),)) + request: Final = urllib.request.Request( + f"{credentials.base_url}{SESSION_ENDPOINT}?{query}", + headers={ # mutable-ok: urllib.request.Request takes a dict + "Authorization": f"Bearer {credentials.api_key}", + "Accept": "application/json", + }, + ) + try: + with urllib.request.urlopen(request, timeout=FETCH_TIMEOUT_SECONDS) as response: + raw: Final[bytes] = response.read() + except urllib.error.HTTPError as error: + return Fetched(session=None, definitive=400 <= error.code < 500) + except (urllib.error.URLError, OSError): + return Fetched(session=None, definitive=False) + session: Final = _session_from_payload(as_mapping(load_json(raw))) + return Fetched(session=session, definitive=session is not None) + + +def _session_from_payload(payload: Mapping[str, object]) -> Session | None: + router_name: Final = printable(payload.get("router_name")) + last_model: Final = printable(payload.get("last_model")) + spend: Final = payload.get("spend") + baseline_spend: Final = payload.get("baseline_spend") + if not router_name or not last_model: + return None + if not isinstance(spend, (int, float)) or not isinstance(baseline_spend, (int, float)): + return None + return Session( + router_name=router_name, + last_model=last_model, + spend=float(spend), + baseline_spend=float(baseline_spend), + baseline_model=printable(payload.get("baseline_model")) or None, + ) + + +def cache_path(cache_dir: Path, credentials: Credentials, session_id: str) -> Path: + identity: Final = "\n".join((credentials.base_url, credentials.api_key, session_id)) + return cache_dir / hashlib.sha256(identity.encode()).hexdigest() + + +def load_session( + credentials: Credentials, + session_id: str, + cache_dir: Path, + fetch: Fetch = fetch_session, + now: Callable[[], float] = time.time, +) -> Session | None: + path: Final = cache_path(cache_dir, credentials, session_id) + cached: Final = _read_cache(path) + fetched_at: Final = cached.get("fetched_at") + if isinstance(fetched_at, (int, float)) and now() - fetched_at < CACHE_TTL_SECONDS: + return _session_from_payload(as_mapping(cached.get("session"))) + fetched: Final = fetch(credentials, session_id) + if fetched.definitive: + _write_cache(path, fetched.session, now()) + return fetched.session + + +NOFOLLOW: Final = getattr(os, "O_NOFOLLOW", 0) + + +def cache_dir_name() -> str: + return f"litellm-statusline-{os.getuid()}" if hasattr(os, "getuid") else "litellm-statusline" + + +def _own_private_dir(directory: Path) -> bool: + """A shared temp root lets another local user pre-create the directory, so it must be ours and private + before anything is read or written under it. Windows has no uids or POSIX mode bits and a per-user temp + directory already, so there it only has to exist and not be a link.""" + try: + directory.mkdir(mode=0o700, parents=True, exist_ok=True) + status: Final = directory.lstat() + except OSError: + return False + if not os.path.isdir(directory) or os.path.islink(directory): + return False + if not hasattr(os, "getuid"): + return True + return status.st_uid == os.getuid() and not status.st_mode & 0o077 + + +def _read_cache(path: Path) -> Mapping[str, object]: + if not _own_private_dir(path.parent): + return EMPTY + try: + descriptor: Final = os.open(path, os.O_RDONLY | NOFOLLOW) + with os.fdopen(descriptor, "rb") as handle: + return as_mapping(load_json(handle.read())) + except OSError: + return EMPTY + + +def _write_cache(path: Path, session: Session | None, fetched_at: float) -> None: + """Staged beside the entry and renamed into place, so a refresh reading the entry never sees a torn write.""" + entry: Final = session._asdict() if session else None + body: Final = json.dumps({"fetched_at": fetched_at, "session": entry}) # mutable-ok: json.dumps takes a dict + if not _own_private_dir(path.parent): + return + try: + descriptor, staged = tempfile.mkstemp(dir=path.parent, prefix=".tmp-") + except OSError: + return + try: + with os.fdopen(descriptor, "w") as handle: + handle.write(body) + os.replace(staged, path) + except OSError: + Path(staged).unlink(missing_ok=True) + + +def _bar(fraction: float, color: str, width: int, use_color: bool) -> str: + filled: Final = round(max(0.0, min(1.0, fraction)) * width) + if not use_color: + return BAR_FULL * filled + BAR_EMPTY * (width - filled) + return f"{color}{BAR_FULL * filled}{DIM}{BAR_EMPTY * (width - filled)}{RESET}" + + +def render(model: str, session: Session | None, config_dir: Path, use_color: bool, bar_width: int = BAR_WIDTH) -> str: + def paint(code: str, text: str) -> str: + return f"{code}{text}{RESET}" if use_color else text + + routed: Final = paint(BOLD, f"Routed to: {model}") + if session is None: + return routed + header: Final = f"{session.router_name}{SEPARATOR}{routed}" + if session.baseline_model is None or session.baseline_spend <= 0: + return header + reference: Final = baseline_label(session.baseline_model, config_dir) + pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100 + delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}") + peak: Final = max(session.spend, session.baseline_spend) + label_width: Final = max(len(LITELLM_LABEL), len(reference)) + rows: Final = ( + (LITELLM_LABEL, session.spend, LITELLM_COLOR), + (reference, session.baseline_spend, BASELINE_COLOR), + ) + lines: Final = ( + f"{paint(DIM, label.ljust(label_width))} {_bar(amount / peak, color, bar_width, use_color)} " + f"{paint(DIM, f'${amount:.2f}')}" + for label, amount, color in rows + ) + return "\n".join((f"{header} {delta}", *lines)) + + +def color_enabled(env: Mapping[str, str]) -> bool: + return env.get("NO_COLOR") is None and env.get("TERM", "") not in ("", "dumb") + + +def status_line( + payload: Mapping[str, object], env: Mapping[str, str], config_dir: Path, cache_dir: Path, fetch: Fetch +) -> str: + fallback: Final = printable(as_mapping(payload.get("model")).get("display_name")) + served: Final = latest_transcript_model(as_str(payload.get("transcript_path"))) + if not served: + return fallback or "claude" + label: Final = model_label(served, config_dir) + session_id: Final = as_str(payload.get("session_id")) + credentials: Final = claude_credentials(env) + if not session_id or not credentials.usable: + return render(label, None, config_dir, color_enabled(env)) + session: Final = load_session(credentials, session_id, cache_dir, fetch) + return render(label, session, config_dir, color_enabled(env)) + + +def codex_stop_message( + payload: Mapping[str, object], env: Mapping[str, str], config_dir: Path, cache_dir: Path, fetch: Fetch +) -> str: + """No cache here: the Stop hook runs once per turn, and a first turn's cached absence would hide the record + the next turn finds.""" + session_id: Final = as_str(payload.get("session_id")) + credentials: Final = codex_credentials(env) + if not session_id or not credentials.usable: + return "" + session: Final = fetch(credentials, session_id).session + if session is None: + return "" + text: Final = render(model_label(session.last_model, config_dir), session, config_dir, use_color=False) + return json.dumps({"systemMessage": f"\n{text}"}) # mutable-ok: json.dumps takes a dict + + +def run(stdin: IO[str], stdout: IO[str], env: Mapping[str, str], fetch: Fetch = fetch_session) -> None: + """A failure renders each mode's own quiet fallback: Claude Code gets the label it already knows, Codex gets + nothing at all rather than a bare string it would reject as hook JSON.""" + body: Final = as_mapping(load_json(stdin.read())) + codex: Final = body.get("hook_event_name") == CODEX_STOP_EVENT + config_dir: Final = Path(env.get("CLAUDE_CONFIG_DIR") or Path.home() / ".claude") + cache_dir: Final = ( + Path(env.get("TMPDIR") or env.get("TEMP") or env.get("TMP") or tempfile.gettempdir()) / cache_dir_name() + ) + try: + text: Final = ( + codex_stop_message(body, env, config_dir, cache_dir, fetch) + if codex + else status_line(body, env, config_dir, cache_dir, fetch) + ) + except Exception: # noqa: BLE001 # a status line must never break the agent session + stdout.write("" if codex else printable(as_mapping(body.get("model")).get("display_name")) or "claude") + return + stdout.write(text) + + +if __name__ == "__main__": + run(sys.stdin, sys.stdout, os.environ) diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index ffece87ab83..f2624797a5f 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -23,11 +23,12 @@ from .auth import CliContextObj, context_secret_vault, get_stored_api_key, load_ from .claude_settings import ( BACKUP_PATH, CLAUDE_SETTINGS_PATH, - ApiKeyHelper, ClaudeSettingsError, + StaticToken, + install_statusline_script, load_json_or_empty, merge_claude_settings, - resolve_api_key_helper, + write_claude_settings, ) @@ -98,8 +99,7 @@ def restore_claude_settings(settings_path: Path | None = None, backup_path: Path return None if record.existed and record.content is not None: resolved_settings_path.parent.mkdir(parents=True, exist_ok=True) - with open(resolved_settings_path, "w") as f: - json.dump(record.content, f, indent=2) + write_claude_settings(resolved_settings_path, record.content) elif resolved_settings_path.exists(): resolved_settings_path.unlink() resolved_backup_path.unlink() @@ -134,13 +134,10 @@ def ensure_fresh_login(ctx: click.Context) -> None: pkce: Final = _stored_login_is_pkce(vault) login_command: Final = "lite login --pkce" if pkce else "lite login" if not sys.stdin.isatty(): - raise UpError( - f"No fresh LiteLLM login found for this proxy. Run `{login_command}` first (apiKeyHelper " - "reads this token on every Claude Code request)." - ) + raise UpError(f"No fresh LiteLLM login found for this proxy. Run `{login_command}` first.") click.echo("No fresh LiteLLM login found for this proxy; starting login...") - ctx.invoke(login, pkce=pkce) + ctx.invoke(login, config_claude=False, pkce=pkce) if not _usable_login(get_stored_api_key(expected_base_url=base_url, vault=vault), vault): raise UpError("Login did not produce a usable token.") @@ -162,7 +159,9 @@ def up(ctx: click.Context) -> None: """Route every Claude Code session through your LiteLLM proxy until stopped. Patches ~/.claude/settings.json so Claude Code picks up the proxy on its own - next startup, from any terminal -- no need to launch it through `lite`. + next startup, from any terminal -- no need to launch it through `lite`. The + key written is the one this command resolved (your fresh `lite login`, or an + explicit --api-key), copied in as a static token for as long as `up` runs. Press Ctrl-C to stop and restore your original settings. Assumes the proxy is already running (this does not start one for you). Cursor is not supported: it has no equivalent file-based config to patch. @@ -180,7 +179,7 @@ def up(ctx: click.Context) -> None: "running (or crashed without cleanup). Run `lite down` first." ) - api_key_helper: Final = resolve_api_key_helper(base_url) + status_line: Final = install_statusline_script() original_existed: Final = CLAUDE_SETTINGS_PATH.exists() original_settings: Final = load_json_or_empty(CLAUDE_SETTINGS_PATH) write_backup( @@ -191,9 +190,10 @@ def up(ctx: click.Context) -> None: ) CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True) - merged: Final = merge_claude_settings(original_settings, base_url, ApiKeyHelper(api_key_helper)) - with open(CLAUDE_SETTINGS_PATH, "w") as f: - json.dump(merged, f, indent=2) + merged: Final = merge_claude_settings( + original_settings, base_url, StaticToken(api_key), status_line=status_line + ) + write_claude_settings(CLAUDE_SETTINGS_PATH, merged) except (AgentRunError, ClaudeSettingsError) as e: raise click.ClickException(str(e)) @@ -248,7 +248,6 @@ __all__ = [ "load_json_or_empty", "merge_claude_settings", "read_backup", - "resolve_api_key_helper", "restore_claude_settings", "up", "write_backup", diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index b866ecc741f..3a61da164d0 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -103,6 +103,7 @@ class AutoRouterTurnTransaction: cache_ttl_seconds: int | None cache_touched: bool tier: str | None = None + baseline_model: str | None = None class TurnCacheFacts(NamedTuple): @@ -168,7 +169,7 @@ def _write_ttl_seconds(usage_object: Mapping[str, object] | None) -> int | None: SESSION_ID_MAX_CHARS: Final = 256 -def _bounded_session_id(session_id: str) -> str: +def bounded_session_id(session_id: str) -> str: """The session id as stored, bounded so a caller-chosen identifier cannot exceed Postgres's B-tree index entry limit through the composite primary key. Oversized ids map to a stable digest, so their turns still aggregate into one session.""" @@ -194,6 +195,9 @@ def build_autorouter_turn_transaction( classifier_cost folded into this turn's spend: the excluded classifier row is how it was billed, the decision is how it is attributed. Cache facts are derived from the payload's own usage record through the savings owner, never handed in beside it. + The baseline the turn's saved_spend was priced against travels with the turn, so the + row can name the counterfactual for the money it holds even after the router is + reconfigured or removed. """ if payload.get("status") != "success": return None @@ -216,13 +220,15 @@ def build_autorouter_turn_transaction( usage_object_raw: Final = metadata.get("usage_object") cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None) tier_raw: Final = routing_decision.get("tier") + baseline_raw: Final = routing_decision.get("savings_baseline_model") classifier_cost: Final = classifier_cost_from_decision(routing_decision) return AutoRouterTurnTransaction( api_key=api_key, - session_id=_bounded_session_id(session_id), + session_id=bounded_session_id(session_id), router_name=router_name, router_type=str(routing_decision.get("router_type") or "unknown"), tier=tier_raw if isinstance(tier_raw, str) and tier_raw else None, + baseline_model=baseline_raw if isinstance(baseline_raw, str) and baseline_raw else None, model=model, turn_at=turn_at, total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0), @@ -253,6 +259,10 @@ _CACHE_TTL: Final = _p("cache_ttl_seconds") _TOUCHED: Final = _p("cache_touched") _TIER: Final = f"{_p('tier')}::text" _TIER_DELTA: Final = f"(CASE WHEN {_TIER} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_TIER}, 1) END)" +_BASELINE: Final = f"{_p('baseline_model')}::text" +_BASELINE_DELTA: Final = ( + f"(CASE WHEN {_BASELINE} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_BASELINE}, 1) END)" +) _IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at" _SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}" @@ -270,7 +280,8 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t ( last_model, models, turns, unordered_turns, covered_turns, cache_hits, same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, return_turns, return_hits, return_expired_misses, return_within_ttl_misses, - ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns + ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns, + baseline_models ) VALUES ( {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, @@ -281,7 +292,7 @@ VALUES ( (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END), (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END), {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8, - {_p("classifier_cost")}::float8, 1, {_TIER_DELTA} + {_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA} ) ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET turns = t.turns + 1, @@ -317,6 +328,9 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET tier_turns = (CASE WHEN {_TIER} IS NOT NULL AND t.router_type = {_p("router_type")} THEN t.tier_turns || jsonb_build_object({_TIER}, COALESCE((t.tier_turns ->> {_TIER})::int, 0) + 1) ELSE t.tier_turns END), + baseline_models = (CASE WHEN {_BASELINE} IS NOT NULL + THEN t.baseline_models || jsonb_build_object({_BASELINE}, COALESCE((t.baseline_models ->> {_BASELINE})::int, 0) + 1) + ELSE t.baseline_models END), first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at), last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at) """ diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index bbc914a772a..50716e5d474 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -32,11 +32,15 @@ from litellm.proxy.auth.auth_checks import ( can_key_call_resolved_model, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.db.autorouter_session_rollup import AUTOROUTER_BENCHMARKS_SQL +from litellm.proxy.db.autorouter_session_rollup import ( + AUTOROUTER_BENCHMARKS_SQL, + bounded_session_id, +) from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, refresh_proxy_server_request_body_snapshot, ) +from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.team_repository import TeamRepository from litellm.router_strategy.complexity_router import ComplexityRouter @@ -54,6 +58,7 @@ from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterCacheStats, AutoRouterRoutingTestRequest, AutoRouterRoutingTestResponse, + AutoRouterSessionResponse, ComplexityRouterConfigValidationRequest, ComplexityRouterConfigValidationResponse, RequestComplexityRouterConfig, @@ -704,6 +709,51 @@ async def get_auto_router_benchmarks( ) +@router.get( + "/auto_router/session", + tags=("auto router",), + dependencies=(Depends(user_api_key_auth),), + response_model=AutoRouterSessionResponse, +) +async def get_auto_router_session( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + session_id: Annotated[ + str, Query(description="The client session id (x-*-session-id header) the turns were sent under") + ], +) -> AutoRouterSessionResponse: + """ + One auto-routed session, for the key that ran it: the model its last turn was routed to and the + session's spend against the router's savings baseline. Built for a coding agent's status line + or stop hook, so any virtual key may call it and only ever sees rows written under its own + key hash. Reads the LiteLLM_AutoRouterSession rollup, which the asynchronous spend flush + fills a moment after each turn; a session with no flushed auto-routed turn yet is a 404. The + id is bounded the way the writer bounded it, so an oversized client id still finds its row. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) + row: Final = await AutoRouterSessionRepository(prisma_client).find_latest_for_key( + user_api_key_dict.api_key, bounded_session_id(session_id) + ) + if row is None: + raise HTTPException( + status_code=404, detail=f"No auto-routed turns recorded for session {session_id!r} under this key" + ) + return AutoRouterSessionResponse( + session_id=session_id, + router_name=row.router_name, + router_type=row.router_type, + turns=row.turns, + last_model=row.last_model, + spend=row.spend, + saved_spend=row.saved_spend, + baseline_spend=row.spend + row.saved_spend, + baseline_model=row.baseline_model, + baseline_models=row.baseline_models, + ) + + # --------------------------------------------------------------------------- # Shadow eval: pre-adoption evaluation of an auto-router against live traffic. # The job row is immutable config plus stopped_at; status, counts, spend, and errors diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 817df082d8c..7d521d54791 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1514,6 +1514,7 @@ model LiteLLM_AutoRouterSession { classifier_cost Float @default(0) classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") + baseline_models Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py index 881f7a66cea..dcf9ddfc32a 100644 --- a/litellm/repositories/__init__.py +++ b/litellm/repositories/__init__.py @@ -2,6 +2,7 @@ Repository classes for database operations. """ +from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.credentials_repository import CredentialsRepository @@ -92,6 +93,7 @@ __all__ = [ "AdaptiveRouterStateRepository", "AgentsRepository", "AuditLogRepository", + "AutoRouterSessionRepository", "BatchTable", "BudgetCascadeUnitOfWork", "BudgetRepository", diff --git a/litellm/repositories/autorouter_session_repository.py b/litellm/repositories/autorouter_session_repository.py new file mode 100644 index 00000000000..d05ef9421ca --- /dev/null +++ b/litellm/repositories/autorouter_session_repository.py @@ -0,0 +1,34 @@ +""" +Repository for the auto-router per-session rollup (LiteLLM_AutoRouterSession). +""" + +from typing import TYPE_CHECKING, Final + +from litellm.models.autorouter_session import LiteLLM_AutoRouterSession +from litellm.repositories.base_repository import BaseRepository +from litellm.repositories.prisma_protocols import TableActions + +if TYPE_CHECKING: + from prisma import models as prisma_models + + +class AutoRouterSessionRepository(BaseRepository[LiteLLM_AutoRouterSession]): + @property + def table(self) -> TableActions["prisma_models.LiteLLM_AutoRouterSession"]: + return self.prisma_client.db.litellm_autoroutersession + + @property + def model_class(self) -> type[LiteLLM_AutoRouterSession]: + return LiteLLM_AutoRouterSession + + async def find_latest_for_key(self, api_key: str, session_id: str) -> LiteLLM_AutoRouterSession | None: + """The session's most recently active router row under exactly this key hash, or None. + + The key is the row's own partition, not a filter over a wider read: the spend writer keyed the + row under the caller's api_key, so a key can only ever see what it wrote itself. + """ + record: Final = await self.table.find_first( + where={"api_key": api_key, "session_id": session_id}, # mutable-ok: Prisma where filter must be a dict + order={"last_turn_at": "desc"}, # mutable-ok: Prisma order clause must be a dict + ) + return self._to_model(record) diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 6306658ad0b..9f29f27e41d 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -226,6 +226,30 @@ class AutoRouterBenchmarkGroup(AutoRouterBenchmarkTotals): ) +class AutoRouterSessionResponse(BaseModel): + """One auto-routed session as its own key sees it: what the last turn ran on, and what the session cost + against the router's savings baseline (the priciest model in its hardest tier).""" + + session_id: str + router_name: str = Field(description="The auto-router alias the session's requests were sent to") + router_type: str = Field(description="complexity, adaptive or quality") + turns: int = Field(description="Auto-routed turns the rollup has recorded for this session so far") + last_model: str = Field(description="The deployment model the most recent turn was routed to") + spend: float = Field(description="What the session's routed traffic actually cost, classifier calls included") + saved_spend: float = Field(description="Estimated savings against the baseline, net of classifier cost") + baseline_spend: float = Field(description="spend plus saved_spend: the estimated single-model cost") + baseline_model: str | None = Field( + description="The savings baseline most of this session's turns were priced against, recorded turn by " + "turn, so it still names the counterfactual after the router is reconfigured or removed. None when no " + "turn recorded one: rows from before the baseline was recorded, and adaptive and quality routers, " + "which derive no baseline and so report no savings" + ) + baseline_models: Mapping[str, int] = Field( + description="Turns priced against each baseline model; more than one entry means the router's " + "baseline changed mid-session and baseline_spend mixes both" + ) + + class AutoRouterBenchmarksResponse(BaseModel): """Benchmarks for the auto-router dashboard, aggregated from the per-session rollup.""" diff --git a/schema.prisma b/schema.prisma index 817df082d8c..7d521d54791 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1514,6 +1514,7 @@ model LiteLLM_AutoRouterSession { classifier_cost Float @default(0) classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") + baseline_models Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index 9ac6476a03c..e8caa241a53 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -42,6 +42,7 @@ async def _turn( saved: float = 0.02, classifier_cost: float = 0.0, tier: "str | None" = None, + baseline: "str | None" = None, ) -> None: touched: Final = 1 if (hit or ttl is not None or not covered) else 0 await db.execute_raw( @@ -61,6 +62,7 @@ async def _turn( ttl, touched, tier, + baseline, ) @@ -338,6 +340,27 @@ async def test_a_mid_session_router_type_change_keeps_foreign_tier_names_out_of_ assert row["turns"] == 3 +async def test_baseline_models_count_the_turns_priced_against_each_baseline(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, baseline="opus") + await _turn(db, key, "B", T0 + timedelta(seconds=10), baseline="opus") + await _turn(db, key, "A", T0 + timedelta(seconds=20), baseline="sonnet") + + assert (await _row(db, key))["baseline_models"] == {"opus": 2, "sonnet": 1} + + +async def test_a_turn_priced_against_no_baseline_leaves_the_map_alone(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, baseline=None) + assert (await _row(db, key))["baseline_models"] == {} + + await _turn(db, key, "A", T0 + timedelta(seconds=10), baseline="opus") + await _turn(db, key, "A", T0 + timedelta(seconds=20), baseline=None) + row = await _row(db, key) + assert row["baseline_models"] == {"opus": 1} + assert row["turns"] == 3 + + async def test_an_out_of_order_turn_still_counts_toward_its_tier(db): key = f"k-{uuid.uuid4()}" await _turn(db, key, "A", T0 + timedelta(seconds=60), tier="simple") diff --git a/tests/test_litellm/litellm_core_utils/test_private_json.py b/tests/test_litellm/litellm_core_utils/test_private_json.py index cedff61959f..3c9f49607f9 100644 --- a/tests/test_litellm/litellm_core_utils/test_private_json.py +++ b/tests/test_litellm/litellm_core_utils/test_private_json.py @@ -4,7 +4,11 @@ import stat import pytest -from litellm.litellm_core_utils.private_json import overwrite_private_json, write_private_json +from litellm.litellm_core_utils.private_json import ( + overwrite_private_json, + write_private_bytes, + write_private_json, +) class TestOverwritePrivateJson: @@ -35,3 +39,32 @@ class TestOverwritePrivateJson: overwrite_private_json(str(path), {"user_id": "u-1"}) assert stat.S_IMODE(path.stat().st_mode) == 0o600 + + +class TestWritePrivateBytes: + def test_replaces_the_file_in_one_step_so_a_reader_holding_the_old_one_keeps_it_whole(self, tmp_path): + path = tmp_path / "script.py" + write_private_bytes(str(path), b"print('one')\n" * 200) + before = path.stat().st_ino + + with path.open("rb") as reader: + write_private_bytes(str(path), b"print('two')\n") + assert reader.read() == b"print('one')\n" * 200 + + assert path.read_bytes() == b"print('two')\n" + assert path.stat().st_ino != before + assert [child.name for child in tmp_path.iterdir()] == ["script.py"] + + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") + def test_lands_owner_only_and_a_refused_stage_leaves_the_previous_file_untouched(self, tmp_path): + path = tmp_path / "script.py" + write_private_bytes(str(path), b"first") + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + + tmp_path.chmod(0o500) + try: + with pytest.raises(PermissionError): + write_private_bytes(str(path), b"second") + finally: + tmp_path.chmod(0o700) + assert path.read_bytes() == b"first" diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index 9ae9b732066..aa6449c98dd 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -8,6 +8,7 @@ import pytest from pydantic import BaseModel, TypeAdapter from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.models.autorouter_session import LiteLLM_AutoRouterSession from litellm.models.budget import ( LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, @@ -588,3 +589,35 @@ class TestManagedTables: ) assert table.vector_store_id == "vs1" assert table.custom_llm_provider == "openai" + + +class TestAutoRouterSession: + @staticmethod + def _row(baseline_models: dict) -> LiteLLM_AutoRouterSession: + return LiteLLM_AutoRouterSession( + api_key="k", + session_id="s", + router_name="auto", + router_type="complexity", + first_turn_at=datetime(2026, 9, 1, 12, 0, 0), + last_turn_at=datetime(2026, 9, 1, 12, 5, 0), + last_model="anthropic/claude-sonnet-5", + turns=3, + spend=0.14, + saved_spend=0.24, + classifier_cost=0.0, + tier_turns={}, + baseline_models=baseline_models, + ) + + def test_the_baseline_label_is_the_one_most_turns_were_priced_against(self): + assert self._row({"anthropic/claude-opus-5": 2, "anthropic/claude-sonnet-5": 1}).baseline_model == ( + "anthropic/claude-opus-5" + ) + + def test_a_tie_between_baselines_is_broken_deterministically(self): + assert self._row({"b-model": 1, "a-model": 1}).baseline_model == "b-model" + assert self._row({"a-model": 1, "b-model": 1}).baseline_model == "b-model" + + def test_a_row_whose_turns_recorded_no_baseline_has_no_label(self): + assert self._row({}).baseline_model is None diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 38a0e85c1ea..c8b3d789665 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3916,3 +3916,28 @@ def test_claude_code_marketplace_routes_open_to_internal_users(route): """Per-skill visibility is enforced inside the handler, so the route gate must let non-admins through.""" assert RouteChecks.is_llm_api_route(route) is True assert _gate(route, LitellmUserRoles.INTERNAL_USER.value) == "allowed" + + +@pytest.mark.parametrize("user_role", [None, LitellmUserRoles.INTERNAL_USER.value, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value]) +def test_auto_router_session_is_reachable_by_any_key_but_benchmarks_stays_admin_only(user_role): + valid_token = UserAPIKeyAuth(api_key="hash-of-caller", user_role=user_role) + request = MagicMock(spec=Request) + request.query_params = {"session_id": "sess-1"} + + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=user_role, + route="/auto_router/session", + request=request, + valid_token=valid_token, + request_data={}, + ) + with pytest.raises(Exception, match="Only proxy admin"): + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=user_role, + route="/auto_router/benchmarks", + request=request, + valid_token=valid_token, + request_data={}, + ) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py index 77742ea9f9f..028ab58843f 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py @@ -6,6 +6,7 @@ from typing import Optional import yaml from click.testing import CliRunner +from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError from litellm.proxy.client.cli.commands.autoroute import commands as commands_module from litellm.proxy.client.cli.commands.autoroute import process as process_module from litellm.proxy.client.cli.commands.autoroute.commands import down, up @@ -163,6 +164,7 @@ class TestUpCommand: # `lite configure claude --model` or a user pin would 400 on the first message. assert captured["settings"]["model"] == "autorouter" assert captured["settings"]["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter" + assert captured["settings"]["statusLine"]["command"].endswith("statusline.py") assert captured["settings_mode"] == 0o600 assert terminate_calls == [99999] @@ -253,6 +255,33 @@ class TestUpCommand: assert not pid_record_path.exists() assert not backup_path.exists() + def test_a_status_line_install_failure_leaves_no_backup_behind(self, monkeypatch, tmp_path): + # The install runs before the backup is written, so a failure cannot strand a backup that + # would make every later `lite configure` / `lite autoroute up` think a session still owns settings.json + config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + claude_settings_path.write_text(json.dumps({"theme": "dark"})) + + def boom(): + raise ClaudeSettingsError("disk full") + + fake_process = FakeProcess(pid=778) + terminate_calls = [] + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "is_port_available", lambda port: True) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + monkeypatch.setattr(commands_module, "install_statusline_script", boom) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + result = self.runner.invoke(up) + + assert result.exit_code != 0 and "disk full" in result.output + assert terminate_calls == [778] + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == {"theme": "dark"} + def test_up_uses_the_same_port_and_master_key_across_runs(self, monkeypatch, tmp_path): """The LIT-4607/LIT-4608 regression: a client configured against one session must keep working in the next, so consecutive runs must patch settings with an identical base URL diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py index f399b6f957a..47b5459d489 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py @@ -161,7 +161,13 @@ class TestBuildGeneratedModelList: config = _base_config(classifier=HeuristicClassifier(), semantic_matching=NoSemanticMatching()) autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") router_config = autorouter["litellm_params"]["complexity_router_config"] - assert set(router_config.keys()) == {"tiers", "default_model"} + assert set(router_config.keys()) == {"tiers", "default_model", "return_raw_model_name"} + + def test_the_generated_router_reports_the_tier_model_it_routed_to(self): + # The status line reads the routed model from the response body, which the proxy restamps to the + # requested alias unless the deployment opts out; "autorouter" on every line would tell nothing. + autorouter = next(m for m in build_generated_model_list(_base_config()) if m["model_name"] == "autorouter") + assert autorouter["litellm_params"]["complexity_router_config"]["return_raw_model_name"] is True class TestBuildGeneratedProxyConfig: diff --git a/tests/test_litellm/proxy/client/cli/conftest.py b/tests/test_litellm/proxy/client/cli/conftest.py index 50c76d3f125..c77f516a768 100644 --- a/tests/test_litellm/proxy/client/cli/conftest.py +++ b/tests/test_litellm/proxy/client/cli/conftest.py @@ -5,6 +5,8 @@ from typing import Final import pytest +from litellm.proxy.client.cli.commands import claude_settings + REAL_CLAUDE_SETTINGS: Final = Path(os.path.expanduser("~")) / ".claude" / "settings.json" @@ -12,6 +14,11 @@ def _current_bytes() -> bytes | None: return REAL_CLAUDE_SETTINGS.read_bytes() if REAL_CLAUDE_SETTINGS.exists() else None +@pytest.fixture(autouse=True) +def _statusline_script_under_tmp(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setattr(claude_settings, "STATUSLINE_SCRIPT_PATH", tmp_path / "litellm-home" / "statusline.py") + + @pytest.fixture(autouse=True) def isolated_claude_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[Path]: before: Final = _current_bytes() diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index bebea285edc..7804435a60d 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -164,27 +164,6 @@ class TestBuildAgentEnv: assert env["PATH"] == "/usr/bin" assert base == {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} - def test_anthropic_profile_leaves_the_bearer_to_the_api_key_helper(self): - env = build_agent_env( - {"ANTHROPIC_AUTH_TOKEN": "stale-token", "ANTHROPIC_API_KEY": "real-key"}, - "http://localhost:4000/", - "sk-key", - frozenset({"anthropic"}), - export_anthropic_token=False, - ) - assert "ANTHROPIC_AUTH_TOKEN" not in env - assert "ANTHROPIC_API_KEY" not in env - assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" - assert env["ENABLE_TOOL_SEARCH"] == "true" - assert env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" - - def test_helper_mode_still_exports_the_openai_key(self): - env = build_agent_env( - {}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"}), export_anthropic_token=False - ) - assert "ANTHROPIC_AUTH_TOKEN" not in env - assert env["OPENAI_API_KEY"] == "sk-key" - class TestAgentLaunchArgs: def test_claude_and_opencode_get_no_extra_args(self): @@ -531,25 +510,6 @@ class TestRunAgent: assert "ANTHROPIC_API_KEY" not in env assert "OPENAI_BASE_URL" not in env - def test_helper_supplied_token_never_reaches_the_launch_env(self): - calls = {} - verified = [] - - run_agent( - "http://localhost:4000", - "sk-key", - ["claude"], - base_env={"PATH": "/usr/bin", "ANTHROPIC_AUTH_TOKEN": "stale-token"}, - which=lambda name: "/usr/local/bin/claude", - verify=lambda base_url, api_key: verified.append(api_key), - launcher=lambda p, a, e: calls.update(env=dict(e)), - export_anthropic_token=False, - ) - - assert verified == ["sk-key"] - assert "ANTHROPIC_AUTH_TOKEN" not in calls["env"] - assert calls["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" - def test_codex_gets_openai_env(self): calls = {} run_agent( @@ -1093,101 +1053,6 @@ class TestAgentCommands: in result.output ) - def _invoke_claude_with_settings(self, tmp_path, settings, obj, *, default_settings=None): - config_dir = tmp_path / "claude-config" - config_dir.mkdir() - if settings is not None: - (config_dir / "settings.json").write_text(json.dumps(settings)) - default_path = tmp_path / "home-claude" / "settings.json" - default_path.parent.mkdir() - if default_settings is not None: - default_path.write_text(json.dumps(default_settings)) - captured = {} - with ( - patch(f"{CLAUDE_SETTINGS_MODULE}.CLAUDE_SETTINGS_PATH", default_path), - patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"), - patch(f"{AGENTS_MODULE}.run_agent", side_effect=lambda b, k, c, **kw: captured.update(kw)), - ): - result = self.runner.invoke( - _agent_command("claude"), [], obj=obj, env={"CLAUDE_CONFIG_DIR": str(config_dir)} - ) - assert result.exit_code == 0, result.output - return captured, result.output - - def test_helper_is_read_from_the_config_dir_claude_code_uses(self, tmp_path): - captured, output = self._invoke_claude_with_settings( - tmp_path, - {"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, - {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, - ) - - assert captured["export_anthropic_token"] is False - assert str(tmp_path / "claude-config" / "settings.json") in output - - def test_helper_only_in_the_default_file_keeps_the_env_token_when_config_dir_points_elsewhere(self, tmp_path): - captured, output = self._invoke_claude_with_settings( - tmp_path, - None, - {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, - default_settings={"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, - ) - - assert captured["export_anthropic_token"] is True - assert "apiKeyHelper" not in output - - def test_stored_login_with_a_matching_helper_leaves_the_token_to_the_helper(self, tmp_path): - captured, output = self._invoke_claude_with_settings( - tmp_path, - {"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, - {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, - ) - - assert captured["export_anthropic_token"] is False - assert "reads its key from the apiKeyHelper" in output - - def test_explicit_key_is_exported_even_when_a_helper_matches(self, tmp_path): - captured, output = self._invoke_claude_with_settings( - tmp_path, - {"apiKeyHelper": "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"}, - {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": False}, - ) - - assert captured["export_anthropic_token"] is True - assert "apiKeyHelper" not in output - - def test_helper_for_another_proxy_keeps_the_env_token(self, tmp_path): - captured, _ = self._invoke_claude_with_settings( - tmp_path, - {"apiKeyHelper": "/usr/local/bin/lite --base-url https://other.example.com auth print-token"}, - {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, - ) - - assert captured["export_anthropic_token"] is True - - def test_no_claude_settings_keeps_the_env_token(self, tmp_path): - captured, _ = self._invoke_claude_with_settings( - tmp_path, - None, - {"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, - ) - - assert captured["export_anthropic_token"] is True - - def test_codex_never_consults_claude_settings(self): - captured = {} - with ( - patch(f"{AGENTS_MODULE}.lite_api_key_helper_configured", side_effect=AssertionError("consulted")), - patch(f"{AGENTS_MODULE}.run_agent", side_effect=lambda b, k, c, **kw: captured.update(kw)), - ): - result = self.runner.invoke( - _agent_command("codex"), - [], - obj={"base_url": "http://localhost:4000", "api_key": "sk-key", "api_key_from_token_file": True}, - ) - - assert result.exit_code == 0, result.output - assert captured["export_anthropic_token"] is True - def test_codex_shows_friendly_name(self): captured = {} with patch( @@ -1338,3 +1203,53 @@ class TestAgentCommands: ) assert result.exit_code == 0, result.output assert captured["reattach_terminal"] is None + + +class TestPrepareCodex: + def test_registers_the_installed_script_as_a_session_scoped_stop_hook(self): + from litellm.proxy.client.cli.commands.agents import prepare_codex + + args = prepare_codex("http://localhost:4000", "sk-key", {}, install=lambda: "/py /home/me/.litellm/statusline.py") + assert args == ( + "-c", + 'hooks.Stop=[{hooks=[{type="command",command="/py /home/me/.litellm/statusline.py"}]}]', + ) + + def test_a_failed_install_is_an_agent_error_not_a_crash(self): + from litellm.proxy.client.cli.commands.agents import AgentRunError, prepare_codex + from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError + + def boom(): + raise ClaudeSettingsError("disk full") + + with pytest.raises(AgentRunError, match="disk full"): + prepare_codex("http://localhost:4000", "sk-key", {}, install=boom) + + def test_a_config_that_already_declares_hooks_keeps_them_and_skips_ours(self, tmp_path): + from litellm.proxy.client.cli.commands.agents import prepare_codex + + warnings = [] + env = {"CODEX_HOME": str(tmp_path)} + for body in ('[[hooks.Stop]]\nhooks = [{ type = "command", command = "mine" }]\n', 'hooks.Stop = []\n', "[hooks]\n"): + (tmp_path / "config.toml").write_text(body) + assert prepare_codex("http://localhost:4000", "sk", env, install=lambda: "/py /s.py", warn=warnings.append) == () + (tmp_path / "config.toml").write_text('model = "gpt-5.6-sol"\n[projects."/x"]\ntrust_level = "trusted"\n') + assert prepare_codex("http://localhost:4000", "sk", env, install=lambda: "/py /s.py", warn=warnings.append) != () + assert len(warnings) == 3 and "already declares hooks" in warnings[0] + + def test_a_config_that_cannot_be_read_or_decoded_still_lets_codex_launch(self, tmp_path): + # A UTF-16 config.toml (a Windows Notepad save) is Codex's problem to report at launch, not a reason + # for the hook pre-check to abort `lite codex` with a traceback before Codex ever starts. + from litellm.proxy.client.cli.commands.agents import codex_declares_stop_hooks, prepare_codex + + config = tmp_path / "config.toml" + config.write_bytes('[[hooks.Stop]]\nhooks = [{ type = "command", command = "mine" }]\n'.encode("utf-16")) + assert codex_declares_stop_hooks(config) is False + assert codex_declares_stop_hooks(tmp_path / "absent.toml") is False + args = prepare_codex("http://localhost:4000", "sk", {"CODEX_HOME": str(tmp_path)}, install=lambda: "/py /s.py") + assert args[0] == "-c" and "hooks.Stop=" in args[1] + + def test_codex_is_wired_through_the_preparer_registry(self): + from litellm.proxy.client.cli.commands.agents import _PREPARERS, prepare_codex + + assert _PREPARERS["codex"] is prepare_codex diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 1f314e0c9d8..3a7792db1fe 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -17,8 +17,15 @@ from litellm.litellm_core_utils.cli_keyring import ( SecretErased, SecretStored, ) -from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord, save_cli_token +from litellm.litellm_core_utils.cli_token_utils import ( + CliTokenRecord, + CredentialNotRecorded, + CredentialNotSaved, + save_cli_token, +) from litellm.proxy.client.cli import cli +from litellm.proxy.client.cli.commands import claude_settings as claude_settings_module +from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner from litellm.proxy.client.cli.commands.auth import ( get_stored_api_key, login, @@ -26,8 +33,6 @@ from litellm.proxy.client.cli.commands.auth import ( print_token, whoami, ) -from litellm.proxy.client.cli.commands import claude_settings as claude_settings_module -from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner @pytest.fixture @@ -1394,7 +1399,9 @@ class TestLoginConfigClaude: monkeypatch.setattr(claude_settings_module, "CONFIGURE_STATE_PATH", tmp_path / "claude_configure_state.json") return backup_path - def _run_login(self, tmp_path, monkeypatch, args, base_url="https://test.example.com", *, config_dir_env=None): + def _run_login( + self, tmp_path, monkeypatch, args, base_url="https://test.example.com", *, config_dir_env=None, stored=None + ): settings_path = tmp_path / "claude" / "settings.json" backup_path = self._isolate_default_settings(tmp_path, monkeypatch) env = {"CLAUDE_CONFIG_DIR": str(settings_path.parent)} if config_dir_env is None else config_dir_env @@ -1411,12 +1418,8 @@ class TestLoginConfigClaude: patch("webbrowser.open"), patch("requests.post", return_value=_mock_cli_sso_start_response()), patch("requests.get", return_value=poll_response), - patch("litellm.proxy.client.cli.commands.auth.save_cli_token"), + patch("litellm.proxy.client.cli.commands.auth.save_cli_token", return_value=stored or SecretStored()), patch("litellm.proxy.client.cli.interface.show_commands"), - patch( - "litellm.proxy.client.cli.commands.claude_settings.shutil.which", - return_value="/usr/local/bin/lite", - ), ): result = self.runner.invoke(login, args, obj={"base_url": base_url}, env=env) return result, settings_path, backup_path @@ -1436,10 +1439,13 @@ class TestLoginConfigClaude: written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://test.example.com" assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" - assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" + # The minted key goes in as a static token: an apiKeyHelper would make Claude Code spawn `lite` (and + # its keychain probe) on every credential refresh, which is what this flag used to write. + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" + assert "apiKeyHelper" not in written assert f"Configured Claude Code: {settings_path} now routes through https://test.example.com." in result.output - assert "pins a proxy model for every tier" not in result.output - assert "the model Claude Code starts on" in result.output + assert "run `lite login --config-claude` again after it expires" in result.output + assert "the model Claude Code starts and resumes on" in result.output def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path, monkeypatch): settings_path = tmp_path / "claude" / "settings.json" @@ -1490,7 +1496,7 @@ class TestLoginConfigClaude: assert result.exit_code == 0, result.output written = json.loads(settings_path.read_text()) - assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" assert f"Configured Claude Code: {settings_path} now routes through https://test.example.com." in result.output def test_flag_keeps_a_config_dir_receipt_apart_from_the_default_file_receipt(self, tmp_path, monkeypatch): @@ -1503,6 +1509,42 @@ class TestLoginConfigClaude: assert len(receipts) == 1 assert json.loads(receipts[0].read_text())["file_existed"] is False + def test_a_second_login_replaces_the_key_and_unconfigure_still_restores_the_original(self, tmp_path, monkeypatch): + # The stored key expires daily, so the flag is re-run per login; the receipt must keep owning the + # slot across re-logins and hand back what was there before the first one. + from litellm.proxy.client.cli.commands.configure import unconfigure_claude + + settings_path = tmp_path / "claude" / "settings.json" + settings_path.parent.mkdir(parents=True) + settings_path.write_text(json.dumps({"theme": "dark", "env": {"ANTHROPIC_AUTH_TOKEN": "sk-theirs"}})) + self._run_login(tmp_path, monkeypatch, ["--config-claude"]) + first = json.loads(settings_path.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] + self._run_login(tmp_path, monkeypatch, ["--config-claude"]) + assert json.loads(settings_path.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == first != "sk-theirs" + + result = self.runner.invoke(unconfigure_claude, [], env={"CLAUDE_CONFIG_DIR": str(settings_path.parent)}) + assert result.exit_code == 0, result.output + assert json.loads(settings_path.read_text()) == {"theme": "dark", "env": {"ANTHROPIC_AUTH_TOKEN": "sk-theirs"}} + + @pytest.mark.parametrize( + "stored", + [CredentialNotSaved("read-only ~/.litellm"), CredentialNotRecorded()], + ids=["nothing-kept-it", "keychain-took-it-file-refused"], + ) + def test_claude_code_is_configured_even_when_the_cli_could_not_keep_the_credential( + self, tmp_path, monkeypatch, stored + ): + # The key is in hand either way, and --config-claude asked for exactly that key to be written into + # settings.json; whether the CLI's own token file or keychain kept a copy is a separate outcome. + result, settings_path, _backup_path = self._run_login(tmp_path, monkeypatch, ["--config-claude"], stored=stored) + + assert result.exit_code == 0, result.output + written = json.loads(settings_path.read_text()) + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" + assert f"Configured Claude Code: {settings_path}" in result.output + assert "even though the CLI itself could not keep it" in result.output + assert "You can now use the CLI without specifying --api-key" not in result.output + def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path, monkeypatch): settings_path = tmp_path / "claude" / "settings.json" settings_path.parent.mkdir(parents=True) diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index fc2d98f2264..cf52d41e963 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -1,7 +1,9 @@ import json import os +import pathlib import shlex import stat +import sys import time from pathlib import Path from unittest.mock import patch @@ -9,8 +11,6 @@ from unittest.mock import patch import pytest from click.testing import CliRunner -from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord -from litellm.proxy.client.cli import cli from litellm.litellm_core_utils.private_json import commit_staged_json from litellm.proxy.client.cli.commands.claude_settings import ( ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, @@ -21,7 +21,6 @@ from litellm.proxy.client.cli.commands.claude_settings import ( OWNED_ENV_KEYS, OWNED_TOP_LEVEL_KEYS, SETTINGS_FILE_OWNERS, - ApiKeyHelper, ClaudeSettingsError, KeepModel, SettingsFileOwner, @@ -30,11 +29,12 @@ from litellm.proxy.client.cli.commands.claude_settings import ( UnpinModel, claude_settings_path, configure_claude_settings, + install_statusline_script, configure_state_path, - lite_api_key_helper_configured, merge_claude_settings, - resolve_api_key_helper, + statusline_command, unconfigure_claude_settings, + with_status_line, ) @@ -45,101 +45,33 @@ def _owners(*backup_paths): CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" AUTH_MODULE = "litellm.proxy.client.cli.commands.auth" -WINDOWS_LITE_EXE = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE" - -CMD_METACHARACTERS = frozenset("&|<>^()") -CMD_PERCENT_GUARD = "%%cd:~,%" - - -def _through_cmd_exe(command): - """The line cmd.exe hands to CreateProcess after reading the apiKeyHelper. - - A `"` toggles cmd's quote state and the metacharacters only act outside it. cmd expands - `%VAR%` even inside quotes, so every `%` has to arrive as the `%%cd:~,%` guard: the first - `%` has no variable name and stays literal, and `%cd:~,%` is a zero length substring of `cd`. - """ - assert not any(CMD_METACHARACTERS & set(run) for run in command.split('"')[::2]), command - assert command.count("%") == 3 * command.count(CMD_PERCENT_GUARD), command - return command.replace(CMD_PERCENT_GUARD, "%") - - -def _through_c_runtime(command_line): - """argv as the Microsoft C runtime builds it for the `lite` executable. - - Outside quotes whitespace ends an argument. A `"` toggles quoting, and inside quotes `""` - is a literal quote. Backslashes are literal unless they run up to a `"`, where each pair - is one backslash and an odd one left over makes the quote literal. - """ - argv = [] - current = None - quoted = False - i = 0 - while i < len(command_line): - ch = command_line[i] - if ch in " \t" and not quoted: - if current is not None: - argv.append(current) - current = None - i += 1 - continue - if current is None: - current = "" - if ch == "\\": - run = len(command_line[i:]) - len(command_line[i:].lstrip("\\")) - before_quote = command_line[i + run : i + run + 1] == '"' - current += "\\" * (run // 2 if before_quote else run) - if before_quote and run % 2: - current += '"' - i += 1 - i += run - elif ch == '"': - if quoted and command_line[i + 1 : i + 2] == '"': - current += '"' - i += 1 - else: - quoted = not quoted - i += 1 - else: - current += ch - i += 1 - return argv if current is None else [*argv, current] - - @pytest.fixture def paths(tmp_path): return tmp_path / "claude" / "settings.json", tmp_path / "backup.json" -@pytest.fixture -def lite_on_path(): - with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"): - yield - - -def _helper_configure(base_url, settings_path, owners, state_path=None): - """`lite login --config-claude`'s shape: the login credential behind apiKeyHelper, no pinned model.""" +def _static_configure(base_url, settings_path, owners, state_path=None): + """`lite configure claude --api-key`'s shape: a virtual key as a static token, no pinned model.""" state = state_path if state_path is not None else settings_path.parent.parent / "state.json" - root = base_url.rstrip("/") - configure_claude_settings( - root, ApiKeyHelper(resolve_api_key_helper(root)), KeepModel(), settings_path, state, owners - ) + configure_claude_settings(base_url.rstrip("/"), StaticToken("sk-virtual-key"), KeepModel(), settings_path, state, owners) -class TestConfigureWithTheLoginHelper: - def test_creates_the_file_and_its_parent_when_missing(self, paths, lite_on_path): +class TestConfigureClaudeSettings: + def test_creates_the_file_and_its_parent_when_missing(self, paths): settings_path, backup_path = paths assert not settings_path.parent.exists() - _helper_configure("https://proxy.example.com/", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com/", settings_path, _owners(backup_path)) written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key" assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" assert written["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" - assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token" - assert "model" not in written + assert "apiKeyHelper" not in written + assert "model" not in written and "ANTHROPIC_MODEL" not in written["env"] - def test_updates_an_existing_file_preserving_unrelated_settings(self, paths, lite_on_path): + def test_updates_an_existing_file_preserving_unrelated_settings(self, paths): settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.write_text( @@ -148,183 +80,87 @@ class TestConfigureWithTheLoginHelper: "theme": "dark", "permissions": {"allow": ["Bash"]}, "env": {"SOME_OTHER_VAR": "keep-me", "ANTHROPIC_BASE_URL": "https://old.example.com"}, - "apiKeyHelper": "old-helper", } ) ) - _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) written = json.loads(settings_path.read_text()) assert written["theme"] == "dark" assert written["permissions"] == {"allow": ["Bash"]} assert written["env"]["SOME_OTHER_VAR"] == "keep-me" assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" - assert written["apiKeyHelper"] != "old-helper" - def test_rerunning_against_a_new_proxy_refreshes_both_base_url_and_helper(self, paths, lite_on_path): - settings_path, backup_path = paths - - _helper_configure("https://first.example.com", settings_path, _owners(backup_path)) - _helper_configure("https://second.example.com", settings_path, _owners(backup_path)) - - written = json.loads(settings_path.read_text()) - assert written["env"]["ANTHROPIC_BASE_URL"] == "https://second.example.com" - assert "second.example.com" in written["apiKeyHelper"] - assert "first.example.com" not in written["apiKeyHelper"] - - def test_drops_stray_static_credentials_so_the_helper_token_wins(self, paths, lite_on_path): - # Claude Code prefers ANTHROPIC_AUTH_TOKEN over apiKeyHelper, so a virtual key left behind - # by an earlier `lite configure claude --api-key` would silently keep winning. + def test_a_helper_left_by_an_older_lite_is_stripped_so_only_the_static_token_is_sent(self, paths): + # Older `lite` versions wrote `apiKeyHelper: lite auth print-token`; Claude Code would keep spawning + # `lite` (and its keychain probe) on every credential refresh, so configure takes the slot over. settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.write_text( - json.dumps({"env": {"ANTHROPIC_API_KEY": "sk-leaked", "ANTHROPIC_AUTH_TOKEN": "sk-old"}}) + json.dumps({"apiKeyHelper": "/usr/local/bin/lite auth print-token", "env": {"ANTHROPIC_API_KEY": "sk-leaked"}}) ) - _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) - env = json.loads(settings_path.read_text())["env"] - assert "ANTHROPIC_API_KEY" not in env and "ANTHROPIC_AUTH_TOKEN" not in env + written = json.loads(settings_path.read_text()) + assert "apiKeyHelper" not in written + assert "ANTHROPIC_API_KEY" not in written["env"] + assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-virtual-key" - def test_written_file_is_owner_only(self, paths, lite_on_path): + def test_written_file_is_owner_only(self, paths): settings_path, backup_path = paths - _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert stat.S_IMODE(settings_path.stat().st_mode) == 0o600 - def test_refuses_while_lite_up_holds_a_backup(self, paths, lite_on_path): + def test_refuses_while_lite_up_holds_a_backup(self, paths): settings_path, backup_path = paths backup_path.write_text("{}") with pytest.raises(ClaudeSettingsError, match="lite down"): - _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert not settings_path.exists() - def test_refuses_on_corrupt_existing_settings_without_touching_the_file(self, paths, lite_on_path): + def test_refuses_on_corrupt_existing_settings_without_touching_the_file(self, paths): settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.write_text("not json at all {{{") with pytest.raises(ClaudeSettingsError, match="invalid JSON"): - _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert settings_path.read_text() == "not json at all {{{" - def test_reports_an_actionable_error_when_lite_is_not_on_path(self, paths): - settings_path, backup_path = paths - with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=None): - with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"): - _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) - - assert not settings_path.exists() - - def test_reports_an_actionable_error_on_a_non_utf8_file(self, paths, lite_on_path): - """Bytes that are not valid UTF-8 must not escape as UnicodeDecodeError. - - UnicodeDecodeError is a ValueError, not an OSError, so a decode-side catch - is easy to miss; login's broad `except Exception` would then relabel it as - an authentication failure and exit 0. - """ + def test_reports_an_actionable_error_on_a_non_utf8_file(self, paths): + # UnicodeDecodeError is a ValueError, not an OSError, so a decode-side catch is easy to miss. settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.write_bytes(b'{"theme": "\xff\xfe"}') with pytest.raises(ClaudeSettingsError, match="invalid JSON"): - _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) - def test_reports_an_actionable_error_when_the_file_cannot_be_read(self, paths, lite_on_path): - """An unreadable settings file must not surface as "Authentication failed". - - login wraps the whole flow in a broad `except Exception`, so any OSError - escaping this function gets relabelled as an auth failure and sends the - user looking at their SSO config instead of at file permissions. - """ + def test_reports_an_actionable_error_when_the_file_cannot_be_read(self, paths): settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.mkdir() with pytest.raises(ClaudeSettingsError, match="Could not read"): - _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) - def test_reports_an_actionable_error_when_the_file_cannot_be_written(self, paths, lite_on_path): + def test_reports_an_actionable_error_when_the_file_cannot_be_written(self, paths): settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.parent.chmod(0o500) try: with pytest.raises(ClaudeSettingsError, match="Could not write"): - _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) finally: settings_path.parent.chmod(0o700) assert not settings_path.exists() -class TestApiKeyHelperIsActuallyInvocable: - """The helper string is executed verbatim by Claude Code, so it has to parse. - - Asserting only on its text is what let a malformed command (`--base-url`, a - top-level group option, placed after the `print-token` subcommand) ship: click - rejects it with "No such option" and every Claude Code request loses its token. - """ - - def _helper_args(self, base_url): - with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"): - return shlex.split(resolve_api_key_helper(base_url))[1:] - - def test_the_generated_command_parses(self): - result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) - - assert "No such option" not in result.output - assert result.exit_code != 2 - - def test_the_generated_command_reaches_print_token(self): - with patch(f"{AUTH_MODULE}.load_cli_token", return_value=None): - result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) - - assert "Not authenticated" in result.output - - def test_the_generated_command_carries_the_base_url_through(self): - stale = CliTokenRecord( - base_url="http://other-proxy.example.com", - key="sk-stale", - timestamp=time.time(), - ) - with patch(f"{AUTH_MODULE}.load_cli_token", return_value=stale): - result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) - - assert "Not authenticated for this server" in result.output - - def _windows_argv(self, lite_exe, base_url): - with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=lite_exe): - helper = resolve_api_key_helper(base_url, platform="win32") - return _through_c_runtime(_through_cmd_exe(helper)) - - @pytest.mark.parametrize( - ("lite_exe", "base_url"), - [ - (WINDOWS_LITE_EXE, "http://localhost:4000"), - ("C:\\Program Files\\LiteLLM\\lite.EXE", "https://gateway.example.com/?a=1&b=2"), - ("C:\\Users\\u\\Scripts\\lite.EXE", "https://gateway.example.com/team%20a/%7Eproxy"), - ('C:\\odd "dir"\\lite.EXE', "http://localhost:4000/x\\"), - ], - ) - def test_the_windows_command_survives_cmd_exe_and_the_c_runtime(self, lite_exe, base_url): - assert self._windows_argv(lite_exe, base_url) == [lite_exe, "--base-url", base_url, "auth", "print-token"] - - def test_the_windows_command_carries_the_base_url_through_cmd_quoting(self): - stale = CliTokenRecord( - base_url="http://other-proxy.example.com", - key="sk-stale", - timestamp=time.time(), - ) - argv = self._windows_argv(WINDOWS_LITE_EXE, "http://localhost:4000") - with patch(f"{AUTH_MODULE}.load_cli_token", return_value=stale): - result = CliRunner().invoke(cli, argv[1:]) - - assert argv[0] == WINDOWS_LITE_EXE - assert "Not authenticated for this server" in result.output - - class TestConflictingOwnersOfTheSettingsFile: """Both `lite up` and `lite autoroute up` restore a backup when they stop. @@ -332,7 +168,7 @@ class TestConflictingOwnersOfTheSettingsFile: write, which is the exact hazard the guard exists to prevent. """ - def test_any_owner_holding_a_backup_blocks_the_write(self, tmp_path, lite_on_path): + def test_any_owner_holding_a_backup_blocks_the_write(self, tmp_path): settings_path = tmp_path / "claude" / "settings.json" for index, owner in enumerate(SETTINGS_FILE_OWNERS): @@ -340,20 +176,20 @@ class TestConflictingOwnersOfTheSettingsFile: backup.write_text("{}") stand_in = SettingsFileOwner(backup, owner.start_command, owner.stop_command) with pytest.raises(ClaudeSettingsError, match="currently managing"): - _helper_configure("https://proxy.example.com", settings_path, (stand_in,)) + _static_configure("https://proxy.example.com", settings_path, (stand_in,)) backup.unlink() assert not settings_path.exists() - def test_the_error_names_the_owner_that_actually_holds_the_file(self, tmp_path, lite_on_path): + def test_the_error_names_the_owner_that_actually_holds_the_file(self, tmp_path): settings_path = tmp_path / "claude" / "settings.json" backup = tmp_path / "auto.json" backup.write_text("{}") autoroute = SettingsFileOwner(backup, "lite autoroute up", "lite autoroute down") with pytest.raises(ClaudeSettingsError, match="`lite autoroute up` is currently managing"): - _helper_configure("https://proxy.example.com", settings_path, (autoroute,)) + _static_configure("https://proxy.example.com", settings_path, (autoroute,)) with pytest.raises(ClaudeSettingsError, match="Run `lite autoroute down` first"): - _helper_configure("https://proxy.example.com", settings_path, (autoroute,)) + _static_configure("https://proxy.example.com", settings_path, (autoroute,)) def test_the_registry_matches_the_paths_the_commands_actually_use(self): """A second definition of the autoroute dir must not drift from this one.""" @@ -365,7 +201,7 @@ class TestConflictingOwnersOfTheSettingsFile: class TestDoesNotDestroyUserOwnedStructure: - def test_writes_through_a_symlinked_settings_file(self, tmp_path, lite_on_path): + def test_writes_through_a_symlinked_settings_file(self, tmp_path): """os.replace() swaps the symlink for a regular file, detaching a dotfiles repo. There is no backup here to undo that, so the link must survive and its @@ -378,20 +214,20 @@ class TestDoesNotDestroyUserOwnedStructure: link.parent.mkdir() link.symlink_to(real) - _helper_configure("https://proxy.example.com", link, ()) + _static_configure("https://proxy.example.com", link, ()) assert link.is_symlink() assert json.loads(real.read_text())["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" assert json.loads(real.read_text())["theme"] == "dark" - def test_refuses_rather_than_discarding_a_non_object_env(self, paths, lite_on_path): + def test_refuses_rather_than_discarding_a_non_object_env(self, paths): """merge coerces a non-dict env to {}; that is silent data loss on a persistent write.""" settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.write_text(json.dumps({"theme": "dark", "env": "not-an-object"})) with pytest.raises(ClaudeSettingsError, match="non-object"): - _helper_configure("https://proxy.example.com", settings_path, _owners(backup_path)) + _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) assert json.loads(settings_path.read_text())["env"] == "not-an-object" @@ -446,18 +282,13 @@ class TestConfigureStatePath: assert work_state == configure_state_path(tmp_path / "work" / "settings.json") def test_configure_and_unconfigure_under_a_config_dir_leave_the_default_receipt_alone( - self, default_paths, tmp_path, lite_on_path + self, default_paths, tmp_path ): _default_settings, default_state = default_paths work_settings = tmp_path / "work" / "settings.json" work_state = configure_state_path(work_settings) configure_claude_settings( - "https://proxy.example.com", - ApiKeyHelper(resolve_api_key_helper("https://proxy.example.com")), - KeepModel(), - work_settings, - work_state, - (), + "https://proxy.example.com", StaticToken("sk-virtual-key"), KeepModel(), work_settings, work_state, () ) assert work_state.exists() and not default_state.exists() outcome = unconfigure_claude_settings(work_settings, work_state, ()) @@ -465,51 +296,8 @@ class TestConfigureStatePath: assert not work_state.exists() -class TestLiteApiKeyHelperConfigured: - def _settings(self, tmp_path, payload): - settings_path = tmp_path / "settings.json" - settings_path.write_text(payload) - return settings_path - - def test_recognises_the_helper_lite_login_wrote_for_this_proxy(self, tmp_path, lite_on_path): - settings_path = tmp_path / "settings.json" - _helper_configure("https://proxy.example.com/", settings_path, (), tmp_path / "state.json") - - assert lite_api_key_helper_configured("https://proxy.example.com/", settings_path) is True - assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is True - - def test_a_helper_for_another_proxy_does_not_count(self, tmp_path, lite_on_path): - settings_path = tmp_path / "settings.json" - _helper_configure("https://other.example.com", settings_path, (), tmp_path / "state.json") - - assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False - - def test_a_hand_written_helper_does_not_count(self, tmp_path, lite_on_path): - settings_path = self._settings(tmp_path, json.dumps({"apiKeyHelper": "cat ~/.my-proxy-key"})) - - assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False - - def test_missing_or_helperless_settings_do_not_count(self, tmp_path, lite_on_path): - assert lite_api_key_helper_configured("https://proxy.example.com", tmp_path / "absent.json") is False - helperless = json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://proxy.example.com"}}) - settings_path = self._settings(tmp_path, helperless) - assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False - - def test_unreadable_settings_fall_back_to_false(self, tmp_path, lite_on_path): - settings_path = self._settings(tmp_path, "{not json") - - assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False - - def test_lite_missing_from_path_falls_back_to_false(self, tmp_path): - helper = "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token" - settings_path = self._settings(tmp_path, json.dumps({"apiKeyHelper": helper})) - with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=None): - assert lite_api_key_helper_configured("https://proxy.example.com", settings_path) is False - - class TestMergeClaudeSettings: - """One merge for every way Claude Code gets wired: `lite up`, `lite login --config-claude`, - `lite configure claude` and `lite autoroute up`.""" + """One merge for every way Claude Code gets wired: `lite up`, `lite configure claude` and `lite autoroute up`.""" def test_a_static_token_lands_in_env_and_the_helper_slot_is_cleared(self): settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token", "env": {"ANTHROPIC_API_KEY": "leaked"}} @@ -523,12 +311,6 @@ class TestMergeClaudeSettings: assert "model" not in merged assert not any(key in merged["env"] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS) - def test_a_helper_lands_top_level_and_the_static_slots_are_cleared(self): - settings = {"env": {"ANTHROPIC_AUTH_TOKEN": "sk-old", "ANTHROPIC_API_KEY": "leaked"}} - merged = merge_claude_settings(settings, "http://127.0.0.1:4000", ApiKeyHelper("lite auth print-token")) - assert merged["apiKeyHelper"] == "lite auth print-token" - assert "ANTHROPIC_AUTH_TOKEN" not in merged["env"] and "ANTHROPIC_API_KEY" not in merged["env"] - def test_keeps_existing_switch_values_and_unrelated_keys_without_mutating_the_input(self): settings = {"theme": "dark", "env": {"SOME_OTHER_VAR": "value", "ENABLE_TOOL_SEARCH": "false"}} merged = merge_claude_settings(settings, "http://127.0.0.1:4000", StaticToken("token-abc")) @@ -537,13 +319,22 @@ class TestMergeClaudeSettings: assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" assert settings == {"theme": "dark", "env": {"SOME_OTHER_VAR": "value", "ENABLE_TOOL_SEARCH": "false"}} - def test_a_default_model_sets_only_the_row_claude_code_starts_on(self): + def test_a_default_model_pins_the_starting_row_and_the_model_a_resumed_session_keeps(self): + # `model` is only the row Claude Code starts on: a resumed session re-sends the model its transcript + # recorded, which behind a raw-model auto-router is the tier model (403 for a key scoped to the + # router). ANTHROPIC_MODEL outranks the transcript on resume, so the pin has to land there too. merged = merge_claude_settings( {}, "http://127.0.0.1:4000", StaticToken("token-abc"), default_model="claude-auto" ) assert merged["model"] == "claude-auto" + assert merged["env"]["ANTHROPIC_MODEL"] == "claude-auto" assert not any(key in merged["env"] for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS) + def test_without_a_default_model_neither_pin_is_written_and_a_users_own_stays(self): + settings = {"model": "mine", "env": {"ANTHROPIC_MODEL": "mine-too"}} + merged = merge_claude_settings(settings, "http://127.0.0.1:4000", StaticToken("token-abc")) + assert merged["model"] == "mine" and merged["env"]["ANTHROPIC_MODEL"] == "mine-too" + def test_a_tier_model_forces_every_claude_code_tier_as_autoroute_needs(self): # Router's auto-router registry is keyed by the literal requested model string with no # wildcard resolution, so `lite autoroute up` overrides the env var each tier reads. @@ -564,7 +355,7 @@ class TestMergeClaudeSettings: "apiKeyHelper": "old-helper", "model": "old-model", } - for credential in (StaticToken("token-abc"), ApiKeyHelper("helper")): + for credential in (StaticToken("token-abc"), StaticToken("token-rotated")): merged = merge_claude_settings(settings, "http://127.0.0.1:4000", credential, default_model="claude-auto") changed_top_level = {key for key in set(settings) | set(merged) if settings.get(key) != merged.get(key)} assert changed_top_level - {"env"} <= set(OWNED_TOP_LEVEL_KEYS) @@ -580,7 +371,7 @@ class TestMergeClaudeSettings: PROXY = "http://127.0.0.1:4000" ANTHROPIC = "https://api.anthropic.com" -HELPER = ApiKeyHelper("lite auth print-token") +RELOGIN = StaticToken("sk-fresh-login") ORIGINAL = { "theme": "dark", "permissions": {"allow": ["Bash"]}, @@ -664,18 +455,19 @@ UNDO_SCENARIOS = { { "restored": { "env.ANTHROPIC_BASE_URL", + "env.ANTHROPIC_AUTH_TOKEN", "env.ENABLE_TOOL_SEARCH", "env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", - "apiKeyHelper", + "statusLine", }, "kept": (), }, - {"credential": HELPER, "model": KeepModel()}, + {"credential": RELOGIN, "model": KeepModel()}, ), "repeat across credential kinds keeps the first snapshot": ( ORIGINAL, [ - {"credential": HELPER, "model": UnpinModel()}, + {"credential": RELOGIN, "model": UnpinModel()}, {"credential": StaticToken("sk-rotated"), "model": StartOn("claude-sonnet-4-6")}, ], ORIGINAL, @@ -688,13 +480,13 @@ UNDO_SCENARIOS = { {"model": "claude-opus-5"}, {}, ), - "re-login keeps our pin": (None, [{"credential": HELPER, "model": KeepModel()}], None, {"file_removed": True}), + "re-login keeps our pin": (None, [{"credential": RELOGIN, "model": KeepModel()}], None, {"file_removed": True}), "edit between configures survives an unpin repeat": ( ORIGINAL, [ _set("model", "my-favourite"), _set("env.ENABLE_TOOL_SEARCH", "false"), - {"credential": HELPER, "model": UnpinModel()}, + {"credential": RELOGIN, "model": UnpinModel()}, ], {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}}, @@ -704,14 +496,14 @@ UNDO_SCENARIOS = { [ _set("model", "my-favourite"), _set("env.ENABLE_TOOL_SEARCH", "false"), - {"credential": HELPER, "model": KeepModel()}, + {"credential": RELOGIN, "model": KeepModel()}, ], {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, {"kept": {"env.ENABLE_TOOL_SEARCH", "model"}}, ), "edit between configures: a same-model repeat displaces it, so it is what comes back": ( ORIGINAL, - [_set("model", "my-favourite"), _set("env.ENABLE_TOOL_SEARCH", "false"), {"credential": HELPER}], + [_set("model", "my-favourite"), _set("env.ENABLE_TOOL_SEARCH", "false"), {"credential": RELOGIN}], {**ORIGINAL, "env": {**ORIGINAL["env"], "ENABLE_TOOL_SEARCH": "false"}, "model": "my-favourite"}, {"kept": {"env.ENABLE_TOOL_SEARCH"}, "restored_includes": {"model"}}, ), @@ -744,12 +536,12 @@ UNDO_SCENARIOS = { None, [ _set("env.ANTHROPIC_API_KEY", "sk-user"), - {"credential": HELPER, "model": KeepModel()}, + {"credential": RELOGIN, "model": KeepModel()}, _set("env.ANTHROPIC_BASE_URL", _ABSENT), ], None, {"withheld": {("env.ANTHROPIC_API_KEY", PROXY)}, "file_removed": True, "receipt_kept": True}, - {"credential": HELPER, "model": KeepModel()}, + {"credential": RELOGIN, "model": KeepModel()}, ), "a credential the user changed is kept, never also withheld": ( ORIGINAL, @@ -832,11 +624,10 @@ class TestConfigureAndUnconfigure: @pytest.mark.parametrize( ("path", "value", "repeat_credential"), [ - ("env.ANTHROPIC_API_KEY", "sk-user-added-later", HELPER), - ("env.ANTHROPIC_AUTH_TOKEN", "sk-users-own-token", HELPER), + ("env.ANTHROPIC_API_KEY", "sk-user-added-later", RELOGIN), ("apiKeyHelper", "/opt/mine/helper", StaticToken("sk-rotated")), ], - ids=["user-adds-api-key", "user-replaces-our-token", "user-sets-own-helper"], + ids=["user-adds-api-key", "user-sets-own-helper"], ) def test_a_credential_the_user_set_between_two_configures_is_what_comes_back( self, tmp_path, path, value, repeat_credential @@ -844,7 +635,7 @@ class TestConfigureAndUnconfigure: # The repeat's merge clears the slot, so the displaced value is snapshotted and is what returns; # it was set while the file pointed at the proxy, so it returns once the file points there again. rig = _Rig(tmp_path, {"theme": "dark"}) - rig.configure(credential=HELPER, model=KeepModel()) + rig.configure(credential=RELOGIN, model=KeepModel()) rig.edit(_set(path, value)) rig.configure(credential=repeat_credential, model=KeepModel()) assert not _lookup(rig.read(), path) @@ -955,3 +746,95 @@ class TestConfigureAndUnconfigure: def _lookup(settings, path): section, _, key = path.rpartition(".") return (settings.get(section) or {}).get(key) if section else settings.get(key) + + +class TestStatusLine: + """Every configure registers the status line, and only while the slot is empty or already ours.""" + + COMMAND = "/opt/lite/bin/python /Users/me/.litellm/statusline.py" + + def test_an_empty_slot_gets_our_status_line(self): + assert with_status_line({}, self.COMMAND)["statusLine"] == {"type": "command", "command": self.COMMAND} + + def test_a_users_own_status_line_is_never_replaced(self): + theirs = {"type": "command", "command": "~/.claude/my-statusline.sh"} + assert with_status_line({"statusLine": theirs}, self.COMMAND)["statusLine"] == theirs + + def test_ours_under_an_older_interpreter_is_refreshed(self): + stale = {"type": "command", "command": "/old/python /Users/me/.litellm/statusline.py"} + assert with_status_line({"statusLine": stale}, self.COMMAND)["statusLine"]["command"] == self.COMMAND + + def test_the_merge_carries_it(self): + merged = merge_claude_settings({}, PROXY, StaticToken("tok"), status_line=self.COMMAND) + assert merged["statusLine"] == {"type": "command", "command": self.COMMAND} + + def test_the_installed_script_is_the_bundled_one_and_the_command_runs_this_interpreter(self, tmp_path): + from litellm.proxy.client.cli.commands import statusline_script + + script = tmp_path / "lite" / "statusline.py" + command = install_statusline_script(script) + assert script.read_bytes() == pathlib.Path(statusline_script.__file__).read_bytes() + assert shlex.split(command) == [sys.executable, str(script)] + assert command == statusline_command(script) + assert stat.S_IMODE(script.stat().st_mode) == 0o600 + assert stat.S_IMODE(script.parent.stat().st_mode) == 0o700 + assert install_statusline_script(script) == command + + def test_a_reinstall_replaces_the_script_in_one_step_and_a_refused_one_leaves_the_old_script_whole(self, tmp_path): + # Claude Code may be running the script at the moment `lite` reinstalls it; the file it has open + # must stay complete, and a reinstall that cannot land must not leave a truncated script behind. + from litellm.proxy.client.cli.commands import statusline_script + + script = tmp_path / "lite" / "statusline.py" + install_statusline_script(script) + bundled = pathlib.Path(statusline_script.__file__).read_bytes() + with script.open("rb") as running: + install_statusline_script(script) + assert running.read() == bundled + assert [child.name for child in script.parent.iterdir()] == ["statusline.py"] + + if os.geteuid() != 0: + script.parent.chmod(0o500) + try: + with pytest.raises(ClaudeSettingsError, match="Could not install the status line script"): + install_statusline_script(script) + finally: + script.parent.chmod(0o700) + assert script.read_bytes() == bundled + + def test_configure_installs_it_and_unconfigure_removes_only_ours(self, tmp_path): + rig = _Rig(tmp_path, {"theme": "dark"}) + script = tmp_path / "statusline.py" + rig.configure(script_path=script) + assert rig.read()["statusLine"]["command"] == statusline_command(script) + assert script.exists() + + outcome = rig.unconfigure() + assert rig.read() == {"theme": "dark"} + assert "statusLine" in outcome.restored + + def test_a_status_line_the_user_replaced_after_configure_survives_unconfigure(self, tmp_path): + rig = _Rig(tmp_path, None) + rig.configure(model=KeepModel(), script_path=tmp_path / "statusline.py") + theirs = {"type": "command", "command": "~/.claude/my-statusline.sh"} + rig.edit(lambda settings: {**settings, "statusLine": theirs}) + + outcome = rig.unconfigure() + assert rig.read()["statusLine"] == theirs + assert "statusLine" in outcome.kept + + def test_a_receipt_from_before_the_status_line_existed_still_unconfigures(self, tmp_path): + # Older receipts never claimed statusLine; a key no configure wrote is never ours, so it stays. + rig = _Rig(tmp_path, None) + script = tmp_path / "statusline.py" + rig.configure(model=KeepModel(), script_path=script) + receipt = json.loads(rig.state.read_text()) + receipt["written"].pop("statusLine") + receipt["previous"].pop("statusLine") + rig.state.write_text(json.dumps(receipt)) + + outcome = rig.unconfigure() + restored = rig.read() + assert "ANTHROPIC_AUTH_TOKEN" not in restored.get("env", {}) + assert restored["statusLine"]["command"] == statusline_command(script) + assert "statusLine" not in outcome.restored diff --git a/tests/test_litellm/proxy/client/cli/test_configure_commands.py b/tests/test_litellm/proxy/client/cli/test_configure_commands.py index ac7408339fe..8ed188af737 100644 --- a/tests/test_litellm/proxy/client/cli/test_configure_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_configure_commands.py @@ -86,6 +86,7 @@ class TestConfigureClaudeWithAVirtualKey: assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in written["env"] assert state_path.exists() assert VALID_KEY not in result.output + assert written["env"]["ANTHROPIC_MODEL"] == "claude-auto" assert "Starting model: claude-auto" in result.output assert "1 of the proxy's 2 models" in result.output assert "lite unconfigure claude" in result.output @@ -99,7 +100,7 @@ class TestConfigureClaudeWithAVirtualKey: assert result.exit_code == 0, result.output written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY - assert "model" not in written + assert "model" not in written and "ANTHROPIC_MODEL" not in written["env"] assert "Starting model: not pinned" in result.output @responses.activate @@ -159,18 +160,11 @@ class TestConfigureClaudeWithAVirtualKey: assert not settings_path.exists() @responses.activate - @pytest.mark.parametrize("entry", ["virtual-key", "login", "interactive"]) - def test_refuses_while_lite_up_holds_a_backup_before_any_login_or_request( - self, runner, paths, monkeypatch, lite_up_backup, entry - ): + @pytest.mark.parametrize("entry", ["virtual-key", "no-key", "interactive"]) + def test_refuses_while_lite_up_holds_a_backup_before_any_request(self, runner, paths, lite_up_backup, entry): _mock_models() - - def login_must_not_run(ctx): - raise AssertionError("the local precondition must be checked before a login is attempted") - - monkeypatch.setattr(configure_module, "ensure_fresh_login", login_must_not_run) if entry == "interactive": - ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": None}) + ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY}) with pytest.raises(click.ClickException, match="lite down"): interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=lambda listed: None) else: @@ -195,33 +189,27 @@ class TestConfigureClaudeWithAVirtualKey: assert json.loads(target.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY -class TestConfigureClaudeWithTheLogin: - def _stored_login(self, monkeypatch): - monkeypatch.setattr(configure_module, "ensure_fresh_login", lambda ctx: None) - monkeypatch.setattr(configure_module, "get_stored_api_key", lambda expected_base_url, vault: VALID_KEY) - +class TestConfigureClaudeWithoutAKey: @responses.activate - def test_uses_the_login_through_the_helper_and_writes_no_secret(self, runner, paths, monkeypatch, lite_on_path): + def test_refuses_and_names_the_ways_to_pass_a_key_without_writing_or_logging_in(self, runner, paths): + # A `lite login` credential expires within a day; the old fallback wrote an apiKeyHelper that made + # Claude Code spawn `lite` (and its keychain probe) on every credential refresh. _mock_models() - self._stored_login(monkeypatch) - settings_path, _ = paths + settings_path, state_path = paths result = runner.invoke( configure_claude, ["--model", "claude-auto"], - obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": True}, + obj={"base_url": PROXY, "api_key": "sk-login-jwt", "api_key_from_token_file": True}, ) - assert result.exit_code == 0, result.output - written = json.loads(settings_path.read_text()) - assert written["apiKeyHelper"] == f"{lite_on_path} --base-url {PROXY} auth print-token" - assert "ANTHROPIC_AUTH_TOKEN" not in written["env"] - assert written["model"] == "claude-auto" - assert VALID_KEY not in settings_path.read_text() - assert "read through apiKeyHelper" in result.output + assert result.exit_code != 0 + assert "--api-key" in result.output and "LITELLM_PROXY_API_KEY" in result.output + assert "apiKeyHelper" not in result.output + assert not settings_path.exists() and not state_path.exists() + assert len(responses.calls) == 0 @responses.activate - def test_an_explicit_key_still_wins_over_a_stored_login(self, runner, paths, monkeypatch, lite_on_path): + def test_an_explicit_key_still_wins_over_a_stored_login(self, runner, paths): _mock_models() - self._stored_login(monkeypatch) settings_path, _ = paths result = runner.invoke( configure_claude, @@ -302,11 +290,13 @@ class TestUnconfigureClaude: edited = json.loads(settings_path.read_text()) edited["env"] = {key: f"{value}-edited" for key, value in edited["env"].items()} edited["model"] = "mine" + edited["statusLine"] = {"type": "command", "command": "~/.claude/my-statusline.sh"} settings_path.write_text(json.dumps(edited)) result = runner.invoke(cli, ["unconfigure", "claude"]) assert result.exit_code == 0, result.output assert "Nothing in" in result.output and "was still ours to restore" in result.output assert "Left as you changed them since:" in result.output and "model" in result.output + assert "statusLine" in result.output @responses.activate def test_names_the_server_a_withheld_credential_was_captured_with_and_keeps_the_receipt(self, runner, paths): diff --git a/tests/test_litellm/proxy/client/cli/test_statusline_script.py b/tests/test_litellm/proxy/client/cli/test_statusline_script.py new file mode 100644 index 00000000000..691764fbef4 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_statusline_script.py @@ -0,0 +1,404 @@ +"""The status line script is copied verbatim to the user's machine, so these drive it the way Claude Code +and Codex do: the documented stdin payload, a transcript on disk, and the proxy behind an injected fetch.""" + +import io +import json +import os +import re +import subprocess +import sys +from pathlib import Path + +import pytest + +from litellm.proxy.client.cli.commands import statusline_script +from litellm.proxy.client.cli.commands.statusline_script import ( + CACHE_TTL_SECONDS, + Credentials, + Fetched, + Session, + cache_dir_name, + cache_path, + claude_credentials, + codex_credentials, + latest_transcript_model, + load_session, + render, + run, +) + +SESSION_ID = "cf712ab8-4c7c-4d48-ba91-eed54bc2956b" +ANSI = re.compile(r"\x1b\[[0-9;]*m") +RECORDED = Session( + router_name="claude-auto", + last_model="anthropic/claude-sonnet-5", + spend=0.14, + baseline_spend=0.38, + baseline_model="anthropic/claude-opus-5", +) + + +def _assistant_line(model: str, **extra: object) -> str: + return json.dumps({"type": "assistant", "message": {"model": model, "role": "assistant"}, **extra}) + + +@pytest.fixture +def transcript(tmp_path: Path) -> Path: + path = tmp_path / "session.jsonl" + path.write_text( + "\n".join( + ( + json.dumps({"type": "user", "message": {"role": "user", "content": "hi"}}), + _assistant_line("claude-haiku-4-5"), + json.dumps({"type": "user", "message": {"role": "user", "content": "harder"}}), + _assistant_line("claude-sonnet-5"), + _assistant_line("claude-haiku-4-5", isSidechain=True), + _assistant_line("claude-haiku-4-5", agentId="agent-1"), + json.dumps({"type": "progress", "data": {}}), + ) + ) + + "\n" + ) + return path + + +@pytest.fixture +def config_dir(tmp_path: Path) -> Path: + directory = tmp_path / "claude" + (directory / "cache").mkdir(parents=True) + (directory / "cache" / "gateway-models.json").write_text( + json.dumps({"models": [{"id": "claude-opus-5", "display_name": "Claude Opus 5"}]}) + ) + return directory + + +def _payload(transcript: Path, session_id: str = SESSION_ID) -> dict: + return { + "session_id": session_id, + "transcript_path": str(transcript), + "model": {"id": "claude-auto", "display_name": "claude-auto"}, + } + + +def _env(tmp_path: Path, config_dir: Path, **extra: str) -> dict[str, str]: + return { + "TMPDIR": str(tmp_path / "tmp"), + "CLAUDE_CONFIG_DIR": str(config_dir), + "TERM": "dumb", + "ANTHROPIC_BASE_URL": "http://127.0.0.1:4000", + "ANTHROPIC_AUTH_TOKEN": "sk-virtual", + **extra, + } + + +def _run(payload: object, env: dict[str, str], fetch) -> str: + out = io.StringIO() + run(io.StringIO(json.dumps(payload)), out, env, fetch) + return out.getvalue() + + +class TestTranscript: + def test_the_latest_foreground_assistant_line_wins_over_later_sidechain_and_agent_lines(self, transcript): + assert latest_transcript_model(str(transcript)) == "claude-sonnet-5" + + def test_a_synthetic_line_is_not_a_served_model(self, tmp_path): + # Claude Code writes `` for messages it produced locally (an API error on resume, for one); + # showing "Routed to: " would name a model no proxy served. + path = tmp_path / "t.jsonl" + path.write_text(_assistant_line("claude-haiku-4-5") + "\n" + _assistant_line("") + "\n") + assert latest_transcript_model(str(path)) == "claude-haiku-4-5" + + def test_a_missing_or_empty_transcript_yields_nothing(self, tmp_path): + empty = tmp_path / "empty.jsonl" + empty.write_text("") + assert latest_transcript_model(str(tmp_path / "missing.jsonl")) == "" + assert latest_transcript_model(str(empty)) == "" + assert latest_transcript_model("") == "" + + +class TestCredentials: + MIXED = { + "ANTHROPIC_BASE_URL": "http://anthropic-side:4000", + "ANTHROPIC_AUTH_TOKEN": "sk-ant", + "OPENAI_BASE_URL": "http://openai-side:4000/v1/", + "OPENAI_API_KEY": "sk-openai", + } + + def test_each_agent_reads_the_pair_it_dials_itself(self): + # A shell that exports both families must not send Codex's hook to the Anthropic proxy. + assert claude_credentials(self.MIXED) == Credentials("http://anthropic-side:4000", "sk-ant") + assert codex_credentials(self.MIXED) == Credentials("http://openai-side:4000", "sk-openai") + + def test_lites_own_shell_variables_are_not_a_credential_either_agent_sends(self): + # A `lite login` shell exports LITELLM_PROXY_*; Claude Code and Codex never read them, so the + # status line must not query the proxy as that principal while the agent used another. + env = {"LITELLM_PROXY_URL": "http://lite:4000/", "LITELLM_PROXY_API_KEY": "sk-lite", **self.MIXED} + assert claude_credentials(env) == Credentials("http://anthropic-side:4000", "sk-ant") + assert codex_credentials(env) == Credentials("http://openai-side:4000", "sk-openai") + assert not claude_credentials({"LITELLM_PROXY_API_KEY": "sk-lite", "ANTHROPIC_BASE_URL": "http://p"}).usable + assert codex_credentials({}) == Credentials("", "") + + def test_claude_code_prefers_the_auth_token_over_a_stray_api_key(self): + env = {"ANTHROPIC_BASE_URL": "http://p", "ANTHROPIC_API_KEY": "sk-stray", "ANTHROPIC_AUTH_TOKEN": "sk-ours"} + assert claude_credentials(env).api_key == "sk-ours" + + def test_an_api_key_helper_in_settings_is_never_run(self, tmp_path, transcript, config_dir): + # `lite` once wrote `apiKeyHelper: lite auth print-token`; running it from a status line that + # refreshes every 300ms spawned `lite` (and a keychain prompt) on every refresh. Without a key + # in the env the proxy is simply not asked. + (config_dir / "settings.json").write_text(json.dumps({"apiKeyHelper": "printf sk-from-helper"})) + asked = [] + env = {k: v for k, v in _env(tmp_path, config_dir).items() if k != "ANTHROPIC_AUTH_TOKEN"} + text = _run(_payload(transcript), env, lambda c, s: asked.append(c) or Fetched(RECORDED, True)) + assert text == "Routed to: claude-sonnet-5" and asked == [] + + +class TestSessionCache: + def test_a_definite_answer_is_served_from_the_cache_within_the_ttl(self, tmp_path): + calls = [] + + def fetch(credentials, session_id): + calls.append(session_id) + return Fetched(RECORDED, definitive=True) + + clock = [100.0] + credentials = Credentials("http://p", "sk") + first = load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: clock[0]) + clock[0] = 100.0 + CACHE_TTL_SECONDS - 1 + second = load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: clock[0]) + clock[0] = 100.0 + CACHE_TTL_SECONDS + 1 + third = load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: clock[0]) + assert first == second == third == RECORDED + assert calls == [SESSION_ID, SESSION_ID] + + def test_a_404_is_cached_as_absence_but_a_transport_failure_is_retried(self, tmp_path): + outcomes = iter((Fetched(None, definitive=False), Fetched(None, definitive=True), Fetched(RECORDED, True))) + calls = [] + + def fetch(credentials, session_id): + calls.append(session_id) + return next(outcomes) + + credentials = Credentials("http://p", "sk") + assert load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: 1.0) is None + assert load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: 1.0) is None + assert load_session(credentials, SESSION_ID, tmp_path, fetch, now=lambda: 1.0) is None + assert len(calls) == 2 + + def test_a_cache_directory_that_is_not_private_is_never_used(self, tmp_path): + # A shared temp root lets another user pre-create the directory; refuse it rather than write into it. + calls = [] + + def fetch(credentials, session_id): + calls.append(session_id) + return Fetched(RECORDED, definitive=True) + + shared = tmp_path / "litellm-statusline" + shared.mkdir(mode=0o755) + credentials = Credentials("http://p", "sk") + for _ in range(2): + assert load_session(credentials, SESSION_ID, shared, fetch, now=lambda: 1.0) == RECORDED + assert calls == [SESSION_ID, SESSION_ID] + assert list(shared.iterdir()) == [] + + def test_any_client_error_is_a_definite_answer_and_a_server_error_is_not(self, monkeypatch): + import urllib.error + + from litellm.proxy.client.cli.commands.statusline_script import fetch_session + + def fail_with(code): + def opener(request, timeout): + raise urllib.error.HTTPError(request.full_url, code, "x", {}, None) + + return opener + + for code, definitive in ((403, True), (401, True), (404, True), (502, False)): + monkeypatch.setattr("urllib.request.urlopen", fail_with(code)) + assert fetch_session(Credentials("http://127.0.0.1:1", "sk"), SESSION_ID) == Fetched(None, definitive) + + def test_the_cache_file_holds_the_proxy_answer_and_never_the_key(self, tmp_path): + credentials = Credentials("http://p", "sk-secret") + load_session(credentials, SESSION_ID, tmp_path, lambda c, s: Fetched(RECORDED, True)) + path = cache_path(tmp_path, credentials, SESSION_ID) + assert "sk-secret" not in written and SESSION_ID not in written if (written := path.read_text()) else False + assert "sk-secret" not in path.name + assert json.loads(written)["session"]["baseline_model"] == "anthropic/claude-opus-5" + assert (path.stat().st_mode & 0o777) == 0o600 + assert (path.parent.stat().st_mode & 0o777) == 0o700 + + def test_a_refresh_replaces_the_entry_in_one_step_so_a_concurrent_refresh_never_reads_a_torn_one(self, tmp_path): + credentials = Credentials("http://p", "sk") + load_session(credentials, SESSION_ID, tmp_path, lambda c, s: Fetched(RECORDED, True), now=lambda: 1.0) + path = cache_path(tmp_path, credentials, SESSION_ID) + first = path.read_text() + + with path.open() as concurrent_reader: + newer = RECORDED._replace(spend=0.5) + load_session(credentials, SESSION_ID, tmp_path, lambda c, s: Fetched(newer, True), now=lambda: 100.0) + assert concurrent_reader.read() == first + assert json.loads(path.read_text())["session"]["spend"] == 0.5 + assert (path.stat().st_mode & 0o777) == 0o600 + assert [child.name for child in tmp_path.iterdir()] == [path.name] + + def test_the_same_session_id_against_another_proxy_or_key_is_not_served_from_the_cache(self, tmp_path): + answers = iter((Fetched(RECORDED, True), Fetched(RECORDED._replace(spend=9.0), True))) + first = load_session(Credentials("http://p", "sk-a"), SESSION_ID, tmp_path, lambda c, s: next(answers)) + second = load_session(Credentials("http://p", "sk-b"), SESSION_ID, tmp_path, lambda c, s: next(answers)) + assert first == RECORDED and second is not None and second.spend == 9.0 + + +class TestRender: + def test_savings_header_and_bars_against_the_routers_baseline(self, config_dir): + text = render("claude-sonnet-5", RECORDED, config_dir, use_color=False, bar_width=10) + assert text.splitlines() == [ + "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5", + "LiteLLM ████░░░░░░ $0.14", + "Claude Opus 5 ██████████ $0.38", + ] + + def test_control_characters_in_any_externally_sourced_label_never_reach_the_terminal(self, tmp_path, config_dir): + # The transcript, the proxy payload and Claude Code's model cache all feed labels straight into a + # terminal, and none is under this script's control. Only the control bytes are dropped (ESC, BEL, + # C1), which is what disarms an OSC-52 clipboard write or a screen clear; the printable remainder of + # such a sequence is inert text and is kept as is. + from litellm.proxy.client.cli.commands.statusline_script import _session_from_payload, model_label + + hostile = "claude-\x1b\x07\x9bsonnet" + path = tmp_path / "t.jsonl" + path.write_text(_assistant_line(hostile) + "\n") + assert latest_transcript_model(str(path)) == "claude-sonnet" + (config_dir / "cache" / "gateway-models.json").write_text( + json.dumps({"models": [{"id": "claude-sonnet", "display_name": "Son\x1b\x07net"}]}) + ) + assert model_label("claude-sonnet", config_dir) == "Sonnet" + session = _session_from_payload( + {"router_name": "auto\x07", "last_model": hostile, "spend": 0.1, "baseline_spend": 0.2, "baseline_model": "op\x1bus"} + ) + assert session == Session("auto", "claude-sonnet", 0.1, 0.2, "opus") + assert latest_transcript_model(str(path)) == "claude-sonnet" + assert "\x1b]52;c;ZXZpbA==" not in render( + latest_transcript_model(str(path)), + _session_from_payload({"router_name": "a", "last_model": "m", "spend": 0.1, "baseline_spend": 0.2, "baseline_model": "\x1b]52;c;ZXZpbA==\x07"}), + config_dir, + use_color=False, + ) + + def test_a_session_that_cost_more_than_its_baseline_reads_as_a_plus(self, config_dir): + dearer = RECORDED._replace(spend=0.50, baseline_spend=0.40) + assert "+25% vs Claude Opus 5" in render("m", dearer, config_dir, use_color=False) + + def test_without_a_baseline_only_the_routed_line_shows(self, config_dir): + assert render("m", RECORDED._replace(baseline_model=None), config_dir, False) == "claude-auto · Routed to: m" + assert render("m", None, config_dir, False) == "Routed to: m" + + def test_color_wraps_the_same_text(self, config_dir): + colored = render("claude-sonnet-5", RECORDED, config_dir, use_color=True, bar_width=10) + assert ANSI.sub("", colored) == render("claude-sonnet-5", RECORDED, config_dir, use_color=False, bar_width=10) + + +class TestClaudeCodeMode: + def test_the_transcript_names_the_routed_model_and_the_proxy_adds_the_savings(self, tmp_path, transcript, config_dir): + seen = [] + + def fetch(credentials, session_id): + seen.append((credentials, session_id)) + return Fetched(RECORDED, definitive=True) + + text = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) + assert text.startswith("claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5\n") + assert seen == [(Credentials("http://127.0.0.1:4000", "sk-virtual"), SESSION_ID)] + + def test_an_unrecorded_session_degrades_to_the_routed_line(self, tmp_path, transcript, config_dir): + assert _run(_payload(transcript), _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) == ( + "Routed to: claude-sonnet-5" + ) + + def test_without_credentials_the_proxy_is_never_asked(self, tmp_path, transcript, config_dir): + def fetch(credentials, session_id): + raise AssertionError("must not fetch") + + env = {k: v for k, v in _env(tmp_path, config_dir).items() if k != "ANTHROPIC_AUTH_TOKEN"} + assert _run(_payload(transcript), env, fetch) == "Routed to: claude-sonnet-5" + + def test_before_the_first_response_the_payloads_display_name_shows(self, tmp_path, config_dir): + payload = _payload(tmp_path / "missing.jsonl") + assert _run(payload, _env(tmp_path, config_dir), lambda c, s: Fetched(RECORDED, True)) == "claude-auto" + + def test_a_discovered_display_name_labels_the_routed_model(self, tmp_path, config_dir): + path = tmp_path / "t.jsonl" + path.write_text(_assistant_line("anthropic/claude-opus-5") + "\n") + assert _run(_payload(path), _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) == ( + "Routed to: Claude Opus 5" + ) + + def test_the_cache_lands_under_the_platforms_temp_dir(self, tmp_path, transcript, config_dir): + env = {k: v for k, v in _env(tmp_path, config_dir).items() if k != "TMPDIR"} + env["TEMP"] = str(tmp_path / "wintemp") + _run(_payload(transcript), env, lambda c, s: Fetched(RECORDED, True)) + assert (tmp_path / "wintemp" / cache_dir_name()).is_dir() + assert cache_dir_name().endswith(str(os.getuid())) + + def test_a_crash_falls_back_to_the_model_label_claude_code_already_knows(self, tmp_path, transcript, config_dir): + def fetch(credentials, session_id): + raise RuntimeError("boom") + + assert _run(_payload(transcript), _env(tmp_path, config_dir), fetch) == "claude-auto" + + def test_garbage_on_stdin_still_prints_something(self, tmp_path, config_dir): + out = io.StringIO() + run(io.StringIO("not json"), out, _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) + assert out.getvalue() == "claude" + + +class TestCodexMode: + def test_the_stop_hook_prints_a_system_message_from_the_proxys_record(self, tmp_path, config_dir): + env = _env(tmp_path, config_dir, OPENAI_BASE_URL="http://127.0.0.1:4000/v1", OPENAI_API_KEY="sk-codex") + env = {k: v for k, v in env.items() if not k.startswith("ANTHROPIC_")} + seen = [] + + def fetch(credentials, session_id): + seen.append(credentials) + return Fetched(RECORDED, definitive=True) + + out = _run({"hook_event_name": "Stop", "session_id": SESSION_ID, "transcript_path": "/nope"}, env, fetch) + message = json.loads(out)["systemMessage"] + assert message.splitlines()[1] == "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5" + assert message.startswith("\n") + assert seen == [Credentials("http://127.0.0.1:4000", "sk-codex")] + + def test_an_unrecorded_session_prints_nothing_so_codex_shows_no_message(self, tmp_path, config_dir): + payload = {"hook_event_name": "Stop", "session_id": SESSION_ID} + assert _run(payload, _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) == "" + + def test_a_crash_prints_nothing_rather_than_text_codex_would_reject(self, tmp_path, config_dir): + def fetch(credentials, session_id): + raise RuntimeError("boom") + + env = _env(tmp_path, config_dir, OPENAI_BASE_URL="http://127.0.0.1:4000/v1", OPENAI_API_KEY="sk-codex") + assert _run({"hook_event_name": "Stop", "session_id": SESSION_ID}, env, fetch) == "" + + def test_a_turn_right_after_an_unrecorded_one_still_asks_the_proxy(self, tmp_path, config_dir): + # One hook run per turn: a miss on turn one must not be cached across turn two's fetch. + answers = iter((Fetched(None, definitive=True), Fetched(RECORDED, definitive=True))) + payload = {"hook_event_name": "Stop", "session_id": SESSION_ID} + env = _env(tmp_path, config_dir, OPENAI_BASE_URL="http://127.0.0.1:4000/v1", OPENAI_API_KEY="sk-codex") + assert _run(payload, env, lambda c, s: next(answers)) == "" + assert "Routed to: claude-sonnet-5" in json.loads(_run(payload, env, lambda c, s: next(answers)))["systemMessage"] + + +class TestStandalone: + def test_the_file_runs_under_a_bare_interpreter_with_no_litellm_on_the_path(self, tmp_path, transcript, config_dir): + # It is copied verbatim to ~/.litellm/statusline.py, so it must be self-contained. + script = tmp_path / "statusline.py" + script.write_bytes(Path(statusline_script.__file__).read_bytes()) + env = {k: v for k, v in _env(tmp_path, config_dir).items() if k != "ANTHROPIC_AUTH_TOKEN"} + completed = subprocess.run( + [sys.executable, "-I", str(script)], + input=json.dumps(_payload(transcript)), + capture_output=True, + text=True, + env=env, + check=True, + timeout=30, + ) + assert completed.stdout == "Routed to: claude-sonnet-5" diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index 111d2f3d682..ddd54dd1374 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -11,7 +11,7 @@ from click.testing import CliRunner from litellm.proxy.client.cli.commands import up as up_module from litellm.proxy.client.cli.commands.agents import AgentRunError -from litellm.proxy.client.cli.commands.claude_settings import ApiKeyHelper, ClaudeSettingsError +from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError, StaticToken from litellm.proxy.client.cli.commands.up import ( BackupRecord, UpError, @@ -20,7 +20,6 @@ from litellm.proxy.client.cli.commands.up import ( load_json_or_empty, merge_claude_settings, read_backup, - resolve_api_key_helper, restore_claude_settings, up, write_backup, @@ -40,52 +39,54 @@ def _patch_paths(monkeypatch, tmp_path): class TestMergeClaudeSettings: def test_preserves_unrelated_top_level_keys(self): - merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", ApiKeyHelper("helper")) + merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["theme"] == "dark" def test_preserves_unrelated_env_keys(self): settings = {"env": {"SOME_OTHER_VAR": "value"}} - merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) + merged = merge_claude_settings(settings, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["env"]["SOME_OTHER_VAR"] == "value" - def test_overrides_base_url_and_helper(self): + def test_overrides_base_url_and_strips_an_old_helper(self): settings = { "env": {"ANTHROPIC_BASE_URL": "https://old.example.com"}, "apiKeyHelper": "old-helper", } - merged = merge_claude_settings(settings, "http://localhost:4000/", ApiKeyHelper("new-helper")) + merged = merge_claude_settings(settings, "http://localhost:4000/", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-fresh" assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" - assert merged["apiKeyHelper"] == "new-helper" + assert "apiKeyHelper" not in merged def test_preserves_existing_gateway_model_discovery(self): settings = {"env": {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "0"}} - merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) + merged = merge_claude_settings(settings, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "0" def test_preserves_existing_tool_search(self): settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} - merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) + merged = merge_claude_settings(settings, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" def test_drops_stray_api_key(self): settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} - merged = merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) + merged = merge_claude_settings(settings, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert "ANTHROPIC_API_KEY" not in merged["env"] def test_works_from_empty_settings(self): - merged = merge_claude_settings({}, "http://localhost:4000", ApiKeyHelper("helper")) + merged = merge_claude_settings({}, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert merged["env"] == { "ANTHROPIC_BASE_URL": "http://localhost:4000", + "ANTHROPIC_AUTH_TOKEN": "sk-fresh", "ENABLE_TOOL_SEARCH": "true", "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "1", } - assert merged["apiKeyHelper"] == "helper" + assert "apiKeyHelper" not in merged def test_does_not_mutate_input(self): settings = {"env": {"FOO": "bar"}} - merge_claude_settings(settings, "http://localhost:4000", ApiKeyHelper("helper")) + merge_claude_settings(settings, "http://localhost:4000", StaticToken("sk-fresh"), status_line="statusline-cmd") assert settings == {"env": {"FOO": "bar"}} @@ -154,6 +155,25 @@ class TestBackupRoundTrip: assert restore_claude_settings() is None assert not settings_path.exists() + def test_restore_writes_owner_only_and_through_a_symlink(self, monkeypatch, tmp_path): + # The backup can hold a token the user had in the file before `up`; a plain open() would put it + # back under the umask, and would replace a dotfiles symlink with a regular file. + target = tmp_path / "dotfiles" / "settings.json" + target.parent.mkdir() + target.write_text("{}") + target.chmod(0o644) + settings_path = tmp_path / "settings.json" + settings_path.symlink_to(target) + monkeypatch.setattr(up_module, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(up_module, "BACKUP_PATH", tmp_path / "backup.json") + write_backup(BackupRecord(existed=True, content={"env": {"ANTHROPIC_AUTH_TOKEN": "sk-theirs"}})) + + restore_claude_settings() + + assert settings_path.is_symlink() + assert json.loads(target.read_text()) == {"env": {"ANTHROPIC_AUTH_TOKEN": "sk-theirs"}} + assert stat.S_IMODE(target.stat().st_mode) == 0o600 + def test_recreates_claude_dir_if_it_was_deleted_while_up_was_running(self, monkeypatch, tmp_path): """If ~/.claude/ is removed while `lite up` holds it open, restoring must recreate the directory rather than crash with FileNotFoundError and strand the backup file, which @@ -216,49 +236,6 @@ class TestBackupRoundTrip: assert not backup_path.exists() -class TestResolveApiKeyHelper: - def test_returns_helper_command_bound_to_the_selected_proxy(self, monkeypatch): - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") - helper = resolve_api_key_helper("http://localhost:4000") - assert helper == "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token" - - def test_quotes_a_base_url_containing_shell_metacharacters(self, monkeypatch): - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") - helper = resolve_api_key_helper("http://example.com/path; rm -rf /") - assert helper == "/usr/local/bin/lite --base-url 'http://example.com/path; rm -rf /' auth print-token" - - def test_raises_when_lite_not_on_path(self, monkeypatch): - monkeypatch.setattr(shutil, "which", lambda name: None) - with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"): - resolve_api_key_helper("http://localhost:4000") - - def test_windows_quotes_for_cmd_exe_instead_of_posix_sh(self, monkeypatch): - """cmd.exe takes a single quote literally, so a POSIX-quoted backslashed path is unrunnable.""" - lite_exe = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE" - monkeypatch.setattr(shutil, "which", lambda name: lite_exe) - - helper = resolve_api_key_helper("https://gateway.example.com", platform="win32") - - assert helper == f'"{lite_exe}" "--base-url" "https://gateway.example.com" "auth" "print-token"' - - def test_windows_keeps_a_spaced_path_and_a_metacharacter_url_as_single_tokens(self, monkeypatch): - monkeypatch.setattr(shutil, "which", lambda name: "C:\\Program Files\\LiteLLM\\lite.EXE") - - helper = resolve_api_key_helper("https://gateway.example.com/?a=1&b=2", platform="win32") - - assert helper == ( - '"C:\\Program Files\\LiteLLM\\lite.EXE" "--base-url" "https://gateway.example.com/?a=1&b=2" ' - '"auth" "print-token"' - ) - - def test_non_windows_platforms_keep_posix_quoting(self, monkeypatch): - monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") - - helper = resolve_api_key_helper("http://example.com/path; rm -rf /", platform="darwin") - - assert helper == "/usr/local/bin/lite --base-url 'http://example.com/path; rm -rf /' auth print-token" - - def _make_ctx(base_url): return click.Context(click.Command("test"), obj={"base_url": base_url}) @@ -307,7 +284,8 @@ def _capture_login(monkeypatch, on_login=lambda: None): login_calls = [] @click.pass_context - def fake_login(ctx, pkce=False): + def fake_login(ctx, config_claude=False, pkce=False): + assert config_claude is False, "`lite up` patches settings itself; the login it starts must not also configure" login_calls.append((ctx.obj["base_url"], pkce)) on_login() @@ -317,8 +295,8 @@ def _capture_login(monkeypatch, on_login=lambda: None): class TestEnsureFreshLogin: """A token that is fresh but was issued for a *different* proxy must not be trusted: without - this check, a user logged into proxy A who runs `up --base-url proxy-b` would silently get an - apiKeyHelper wired up around proxy A's real token, which print-token would then hand to proxy B.""" + this check, a user logged into proxy A who runs `up --base-url proxy-b` would silently get proxy A's + real token written into settings pointed at proxy B.""" def test_reuses_a_fresh_token_issued_for_the_same_proxy(self, monkeypatch): _FakeTokenStore( @@ -500,11 +478,13 @@ class TestUpCommand: settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) original = {"theme": "dark"} settings_path.write_text(json.dumps(original)) + settings_path.chmod(0o644) captured = {} def fake_wait(self, timeout=None): captured["settings"] = json.loads(settings_path.read_text()) + captured["settings_mode"] = stat.S_IMODE(settings_path.stat().st_mode) captured["backup_existed"] = backup_path.exists() return True @@ -514,10 +494,6 @@ class TestUpCommand: patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), patch(f"{UP_MODULE}.verify_proxy_key"), - patch( - f"{UP_MODULE}.resolve_api_key_helper", - return_value="/usr/local/bin/lite auth print-token", - ), patch(f"{UP_MODULE}.signal.signal"), patch(f"{UP_MODULE}.atexit.register"), patch("threading.Event.wait", new=fake_wait), @@ -529,7 +505,11 @@ class TestUpCommand: assert captured["settings"]["theme"] == "dark" assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" assert captured["settings"]["env"]["ENABLE_TOOL_SEARCH"] == "true" - assert captured["settings"]["apiKeyHelper"] == "/usr/local/bin/lite auth print-token" + assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-fresh" + # The file now carries the key, so the umask (and the file's earlier 0644) must not decide who reads it. + assert captured["settings_mode"] == 0o600 + assert "apiKeyHelper" not in captured["settings"] + assert captured["settings"]["statusLine"]["command"].endswith("statusline.py") assert json.loads(settings_path.read_text()) == original assert not backup_path.exists() @@ -598,7 +578,7 @@ class TestUpCanInvokeTheRealLoginCommand: @click.pass_context def driver(ctx): ctx.obj = {"base_url": "http://127.0.0.1:9"} - ctx.invoke(real_login, pkce=False) + ctx.invoke(real_login, config_claude=False, pkce=False) with patch( f"{AUTH_MODULE}._start_cli_sso_flow", @@ -618,7 +598,7 @@ class TestUpCanInvokeTheRealLoginCommand: @click.pass_context def driver(ctx): ctx.obj = {"base_url": "http://127.0.0.1:9"} - ctx.invoke(real_login, pkce=False) + ctx.invoke(real_login, config_claude=False, pkce=False) with patch(f"{AUTH_MODULE}._start_cli_sso_flow", side_effect=RuntimeError("stop")): CliRunner().invoke(driver, [], standalone_mode=False, env={"CLAUDE_CONFIG_DIR": str(tmp_path)}) diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index 4507892bd0f..271751a3ff8 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -59,7 +59,8 @@ class TestBuildTransaction: def test_successful_auto_routed_turn_builds_every_field(self): transaction = _build( metadata=_metadata( - usage_object={"prompt_tokens": 90, "cache_read_input_tokens": 5, "cache_creation_input_tokens": 7} + routing_decision={**ROUTING_DECISION, "savings_baseline_model": "anthropic/claude-opus-5"}, + usage_object={"prompt_tokens": 90, "cache_read_input_tokens": 5, "cache_creation_input_tokens": 7}, ) ) assert transaction == AutoRouterTurnTransaction( @@ -77,6 +78,7 @@ class TestBuildTransaction: cache_hit=True, cache_ttl_seconds=300, cache_touched=True, + baseline_model="anthropic/claude-opus-5", ) @pytest.mark.parametrize( @@ -109,6 +111,17 @@ class TestBuildTransaction: transaction = _build() assert transaction is not None and transaction.tier is None + def test_the_baseline_the_turn_was_priced_against_travels_with_the_turn(self): + decision = {**ROUTING_DECISION, "savings_baseline_model": "anthropic/claude-opus-5"} + transaction = _build(metadata=_metadata(routing_decision=decision)) + assert transaction is not None and transaction.baseline_model == "anthropic/claude-opus-5" + + @pytest.mark.parametrize("baseline", [None, "", 3]) + def test_a_decision_without_a_usable_baseline_records_none(self, baseline: object): + decision = {**ROUTING_DECISION, "savings_baseline_model": baseline} + transaction = _build(metadata=_metadata(routing_decision=decision)) + assert transaction is not None and transaction.baseline_model is None + def test_a_priced_classifier_rides_the_turns_spend(self): """The classifier row is excluded from the rollup, so its charge lands here, folded once into the turn that paid for it (GH #38816).""" @@ -215,6 +228,7 @@ def _transaction( session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12, 0, 0), tier: str | None = "medium", + baseline_model: str | None = "anthropic/claude-opus-5", ) -> AutoRouterTurnTransaction: return AutoRouterTurnTransaction( api_key="k1", @@ -232,6 +246,7 @@ def _transaction( cache_ttl_seconds=None, cache_touched=False, tier=tier, + baseline_model=baseline_model, ) @@ -265,6 +280,7 @@ class TestFlush: None, 0, "medium", + "anthropic/claude-opus-5", ) def test_a_connect_error_retries_the_same_statement(self): diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index dc18e0f7d4a..ef843adad98 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -844,6 +844,143 @@ from litellm.proxy.management_endpoints.auto_router_endpoints import ( from litellm.types.management_endpoints.auto_router_endpoints import SHADOW_EVAL_TURN_VALVE, StartShadowEvalRequest VIEWER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, api_key="sk-view", user_id="viewer") + + +class TestAutoRouterSession: + """GET /auto_router/session: a key reads its own session's routed model and savings, nothing else.""" + + ROW = { + "router_name": "claude-auto", + "router_type": "complexity", + "first_turn_at": datetime(2026, 9, 1, 12, 0, 0), + "last_turn_at": datetime(2026, 9, 1, 12, 5, 0), + "turns": 3, + "last_model": "anthropic/claude-sonnet-5", + "spend": 0.14, + "saved_spend": 0.24, + "classifier_cost": 0.0, + "tier_turns": {"simple": 1, "complex": 2}, + "baseline_models": {"anthropic/claude-opus-5": 3}, + } + + @staticmethod + def _rig(monkeypatch: pytest.MonkeyPatch, rows: Sequence[Mapping[str, object]]): + from litellm.proxy import proxy_server + + lookups: list[tuple[Mapping[str, object], Mapping[str, object]]] = [] + + class _Table: + async def find_first(self, where: Mapping[str, object], order: Mapping[str, object]): + lookups.append((where, order)) + matching = [r for r in rows if (r["api_key"], r["session_id"]) == (where["api_key"], where["session_id"])] + return max(matching, key=lambda r: r["last_turn_at"], default=None) + + monkeypatch.setattr( + proxy_server, "prisma_client", type("P", (), {"db": type("D", (), {"litellm_autoroutersession": _Table()})()})() + ) + return lookups + + @pytest.mark.asyncio + async def test_a_key_reads_its_own_session_with_the_baseline_its_turns_were_priced_against( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + caller = UserAPIKeyAuth(api_key="sk-caller") + self._rig(monkeypatch, [{**self.ROW, "api_key": caller.api_key, "session_id": "sess-1"}]) + response = await get_auto_router_session(user_api_key_dict=caller, session_id="sess-1") + assert response.model_dump() == { + "session_id": "sess-1", + "router_name": "claude-auto", + "router_type": "complexity", + "turns": 3, + "last_model": "anthropic/claude-sonnet-5", + "spend": 0.14, + "saved_spend": 0.24, + "baseline_spend": pytest.approx(0.38), + "baseline_model": "anthropic/claude-opus-5", + "baseline_models": {"anthropic/claude-opus-5": 3}, + } + + @pytest.mark.asyncio + async def test_another_keys_session_is_a_404_even_for_an_admin(self, monkeypatch: pytest.MonkeyPatch): + # The scope is the caller's own key hash, exactly what the spend writer keyed the row under; + # an admin wanting every key's sessions has /auto_router/benchmarks. + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + other = UserAPIKeyAuth(api_key="sk-other") + lookups = self._rig(monkeypatch, [{**self.ROW, "api_key": other.api_key, "session_id": "sess-1"}]) + with pytest.raises(HTTPException) as err: + await get_auto_router_session(user_api_key_dict=ADMIN, session_id="sess-1") + assert err.value.status_code == 404 + assert lookups == [({"api_key": ADMIN.api_key, "session_id": "sess-1"}, {"last_turn_at": "desc"})] + assert ADMIN.api_key != "sk-test" + + @pytest.mark.asyncio + async def test_the_sessions_most_recently_active_router_is_the_one_reported(self, monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + older = {**self.ROW, "api_key": ADMIN.api_key, "session_id": "s", "router_name": "old-auto"} + newer = { + **self.ROW, + "api_key": ADMIN.api_key, + "session_id": "s", + "router_name": "new-auto", + "last_turn_at": datetime(2026, 9, 1, 13, 0, 0), + } + self._rig(monkeypatch, [older, newer]) + response = await get_auto_router_session(user_api_key_dict=ADMIN, session_id="s") + assert response.router_name == "new-auto" + + @pytest.mark.asyncio + async def test_a_reconfigured_router_keeps_the_label_the_money_was_priced_against( + self, monkeypatch: pytest.MonkeyPatch + ): + # The proxy's router now prices against a different baseline, but the row's money was priced + # against opus for two of three turns, and the label says so; the full split is on the response. + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + priced = {"anthropic/claude-opus-5": 2, "anthropic/claude-sonnet-5": 1} + self._rig(monkeypatch, [{**self.ROW, "api_key": ADMIN.api_key, "session_id": "s", "baseline_models": priced}]) + response = await get_auto_router_session(user_api_key_dict=ADMIN, session_id="s") + assert response.baseline_model == "anthropic/claude-opus-5" + assert response.baseline_models == priced + + @pytest.mark.asyncio + async def test_a_session_whose_turns_recorded_no_baseline_reports_the_money_without_a_name( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + self._rig(monkeypatch, [{**self.ROW, "api_key": ADMIN.api_key, "session_id": "s", "baseline_models": {}}]) + response = await get_auto_router_session(user_api_key_dict=ADMIN, session_id="s") + assert response.baseline_model is None + assert response.baseline_spend == pytest.approx(0.38) + + @pytest.mark.asyncio + async def test_an_oversized_client_session_id_is_bounded_like_the_writer_bounded_it( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.db.autorouter_session_rollup import bounded_session_id + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + long_id = "s" * 300 + self._rig(monkeypatch, [{**self.ROW, "api_key": ADMIN.api_key, "session_id": bounded_session_id(long_id)}]) + response = await get_auto_router_session(user_api_key_dict=ADMIN, session_id=long_id) + assert response.session_id == long_id + assert response.turns == 3 + + @pytest.mark.asyncio + async def test_without_a_database_the_endpoint_says_so(self, monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_session + + monkeypatch.setattr(proxy_server, "prisma_client", None) + with pytest.raises(HTTPException) as err: + await get_auto_router_session(user_api_key_dict=ADMIN, session_id="s") + assert err.value.status_code == 500 + + NON_ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="user") diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index a890d7ceed0..c7f6b8ff83a 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -2407,3 +2407,60 @@ class TestCountBillableUsers: client.db.litellm_usertable = _RacyTable() repo = UserRepository(client) assert await repo.count_billable_users() == 0 + + +class TestAutoRouterSessionRepository: + ROW: Final = { + "api_key": "hashed-key", + "session_id": "s1", + "router_name": "claude-auto", + "router_type": "complexity", + "first_turn_at": datetime(2026, 9, 1, 12, 0, 0), + "last_turn_at": datetime(2026, 9, 1, 12, 5, 0), + "last_model": "anthropic/claude-sonnet-5", + "models": {"anthropic/claude-sonnet-5": {"at": 1.0, "ttl": None}}, + "turns": 3, + "spend": 0.14, + "saved_spend": 0.24, + "classifier_cost": 0.01, + "tier_turns": {"complex": 3}, + "baseline_models": {"anthropic/claude-opus-5": 3}, + } + + @staticmethod + def _repo(record: Optional[Dict[str, Any]]): + from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository + + lookups: List[Dict[str, Any]] = [] + + class _Table: + async def find_first(self, where: Dict[str, Any], order: Dict[str, str]): + lookups.append({"where": where, "order": order}) + return MockRecord(record) if record is not None else None + + client = MagicMock() + client.db.litellm_autoroutersession = _Table() + return AutoRouterSessionRepository(client), lookups + + @pytest.mark.asyncio + async def test_find_latest_for_key_reads_the_keys_own_partition_newest_router_first(self): + repo, lookups = self._repo(dict(self.ROW)) + row = await repo.find_latest_for_key("hashed-key", "s1") + assert lookups == [{"where": {"api_key": "hashed-key", "session_id": "s1"}, "order": {"last_turn_at": "desc"}}] + assert row is not None + assert (row.router_name, row.turns, row.spend, row.saved_spend) == ("claude-auto", 3, 0.14, 0.24) + assert row.baseline_models == {"anthropic/claude-opus-5": 3} + assert row.baseline_model == "anthropic/claude-opus-5" + + @pytest.mark.asyncio + async def test_find_latest_for_key_is_none_when_the_key_wrote_no_such_session(self): + repo, _ = self._repo(None) + assert await repo.find_latest_for_key("hashed-key", "unknown") is None + + def test_table_is_the_session_rollup_and_needs_a_database(self): + from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository + + client = MagicMock() + assert AutoRouterSessionRepository(client).table is client.db.litellm_autoroutersession + with pytest.raises(RuntimeError, match="No DB Connected"): + _ = AutoRouterSessionRepository(None).table diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0664a9f2fdc..dd455cba44a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1242,6 +1242,31 @@ export interface paths { patch?: never; trace?: never; }; + "/auto_router/session": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Auto Router Session + * @description One auto-routed session, for the key that ran it: the model its last turn was routed to and the + * session's spend against the router's savings baseline. Built for a coding agent's status line + * or stop hook, so any virtual key may call it and only ever sees rows written under its own + * key hash. Reads the LiteLLM_AutoRouterSession rollup, which the asynchronous spend flush + * fills a moment after each turn; a session with no flushed auto-routed turn yet is a 404. The + * id is bounded the way the writer bounded it, so an oversized client id still finds its row. + */ + get: operations["get_auto_router_session_auto_router_session_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/auto_router/shadow_eval": { parameters: { query?: never; @@ -23572,6 +23597,62 @@ export interface components { /** @description The decision record this request would have written to its log row */ routing_decision: components["schemas"]["StandardLoggingRoutingDecision"]; }; + /** + * AutoRouterSessionResponse + * @description One auto-routed session as its own key sees it: what the last turn ran on, and what the session cost + * against the router's savings baseline (the priciest model in its hardest tier). + */ + AutoRouterSessionResponse: { + /** + * Baseline Model + * @description The savings baseline most of this session's turns were priced against, recorded turn by turn, so it still names the counterfactual after the router is reconfigured or removed. None when no turn recorded one: rows from before the baseline was recorded, and adaptive and quality routers, which derive no baseline and so report no savings + */ + baseline_model: string | null; + /** + * Baseline Models + * @description Turns priced against each baseline model; more than one entry means the router's baseline changed mid-session and baseline_spend mixes both + */ + baseline_models: { + [key: string]: number; + }; + /** + * Baseline Spend + * @description spend plus saved_spend: the estimated single-model cost + */ + baseline_spend: number; + /** + * Last Model + * @description The deployment model the most recent turn was routed to + */ + last_model: string; + /** + * Router Name + * @description The auto-router alias the session's requests were sent to + */ + router_name: string; + /** + * Router Type + * @description complexity, adaptive or quality + */ + router_type: string; + /** + * Saved Spend + * @description Estimated savings against the baseline, net of classifier cost + */ + saved_spend: number; + /** Session Id */ + session_id: string; + /** + * Spend + * @description What the session's routed traffic actually cost, classifier calls included + */ + spend: number; + /** + * Turns + * @description Auto-routed turns the rollup has recorded for this session so far + */ + turns: number; + }; /** AwsSessionTag */ AwsSessionTag: { /** Key */ @@ -41687,6 +41768,38 @@ export interface operations { }; }; }; + get_auto_router_session_auto_router_session_get: { + parameters: { + query: { + /** @description The client session id (x-*-session-id header) the turns were sent under */ + session_id: string; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AutoRouterSessionResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; list_shadow_eval_jobs_auto_router_shadow_eval_get: { parameters: { query?: {