diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index ed1447e4e65..f7d9eb7da9a 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -478,6 +478,7 @@ Launch a coding agent with all of its LLM traffic routed through your LiteLLM pr lite claude lite codex lite opencode +lite pi ``` Anything you type after the agent name is forwarded to it untouched, so the usual flags keep working: @@ -491,17 +492,19 @@ Each command resolves your LiteLLM key (logging in via SSO when none is stored a The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, and older versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. +pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/` wins, and inside the TUI the `/model` picker lists every synced litellm model. + Options (these belong to the wrapper, so put them before the agent's own flags): - `--skip-verify`: Skip the pre-launch key check (useful offline or with non-standard auth). -To pin the model, pass the agent's own model flag (for example `lite claude --model my-proxy-model` or `lite codex -m my-proxy-model`), or export the variable the agent reads (`ANTHROPIC_MODEL` / `ANTHROPIC_SMALL_FAST_MODEL` for Claude Code); the wrapper preserves anything you already have set. Whatever model the agent ends up requesting must exist on the proxy, since requests land on the proxy's `/v1/messages` (Anthropic) or `/v1/chat/completions` and `/v1/responses` (OpenAI) endpoints. +To pin the model, pass the agent's own model flag (for example `lite claude --model my-proxy-model`, `lite codex -m my-proxy-model`, or `lite pi --model my-proxy-model`), or export the variable the agent reads (`ANTHROPIC_MODEL` / `ANTHROPIC_SMALL_FAST_MODEL` for Claude Code); the wrapper preserves anything you already have set. Whatever model the agent ends up requesting must exist on the proxy, since requests land on the proxy's `/v1/messages` (Anthropic) or `/v1/chat/completions` and `/v1/responses` (OpenAI) endpoints. #### About the `lite login` credential The token minted by `lite login` is a short-lived, per-session agent credential, not a managed virtual key. It is scoped to the user and team you authenticated as, inherits that user's and team's models and budgets, and is enforced on the proxy exactly like a virtual key on the same team (guardrails, routing, logging, spend). Spend is tracked against the shared team and user budgets, so running several agents (or logging in more than once) does not hand each session its own separate budget; they all draw down the same team/user allowance. There is no separate per-session cap, so sustained agent use is not capped at a small chat-session limit. -The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, and `lite opencode` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the daily re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead. +The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, `lite opencode`, and `lite pi` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the daily re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead. ### Route Every Claude Code Session Through the Proxy diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 9f79240fe55..7a0ae9dc955 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -5,7 +5,7 @@ import sys from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Final +from typing import Final, TypeAlias import click import requests @@ -13,6 +13,15 @@ from pydantic import BaseModel, TypeAdapter, ValidationError from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login from .cmd_quoting import quote_for_cmd +from .pi import ( + LITELLM_PROXY_API_KEY_ENV, + PI_PROVIDER_NAME, + PiSyncError, + fetch_model_ids, + fetch_model_limits, + models_json_path, + sync_models_json, +) ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN" @@ -32,19 +41,24 @@ _SKIP_VERIFY_FLAG: Final = "--skip-verify" PROFILE_ANTHROPIC: Final = "anthropic" PROFILE_OPENAI: Final = "openai" +PROFILE_LITELLM: Final = "litellm" _KNOWN_AGENTS: Final[dict[str, tuple[str, frozenset[str]]]] = { "claude": ("Claude Code", frozenset({PROFILE_ANTHROPIC})), "codex": ("Codex", frozenset({PROFILE_OPENAI})), "opencode": ("OpenCode", frozenset({PROFILE_OPENAI})), + "pi": ("pi", frozenset({PROFILE_LITELLM})), } _INSTALL_DOCS: Final[dict[str, str]] = { "claude": "https://docs.claude.com/en/docs/claude-code/setup", "codex": "https://developers.openai.com/codex/cli", "opencode": "https://opencode.ai/docs", + "pi": "https://pi.dev", } +_HIDDEN_AGENTS: Final = frozenset({"pi"}) + CODEX_PROXY_PROVIDER: Final = "litellm" @@ -81,6 +95,8 @@ def build_agent_env( the environment is left alone. CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY defaults to 1 so Claude Code (v2.1.129+) fills its /model picker from the 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. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") @@ -95,6 +111,8 @@ def build_agent_env( if PROFILE_OPENAI in profiles: env[OPENAI_BASE_URL_ENV] = root + "/v1" env[OPENAI_API_KEY_ENV] = api_key + if PROFILE_LITELLM in profiles: + env[LITELLM_PROXY_API_KEY_ENV] = api_key return env @@ -130,6 +148,40 @@ _PROXY_ARGS: Final[dict[str, Callable[[str], list[str]]]] = { } +def prepare_pi( + base_url: str, + api_key: str, + base_env: Mapping[str, str], + *, + get: Callable[..., requests.Response] = requests.get, +) -> tuple[str, ...]: + """Sync the proxy's model list into pi's models.json before handoff. + + pi has no base-URL env vars, so this file is the only way to point it at the + proxy. Only the litellm provider entry is touched; the synced entry references + the key as $LITELLM_PROXY_API_KEY, which build_agent_env exports. The returned + --model pin is needed because pi ignores a bare --provider when picking the + interactive startup model; a user-supplied --model comes later in argv and wins. + """ + ids: Final = fetch_model_ids(base_url, api_key, get=get) + if isinstance(ids, PiSyncError): + raise AgentRunError(ids.message) + limits: Final = fetch_model_limits(base_url, api_key, get=get) + path: Final = models_json_path(base_env) + error: Final = sync_models_json(path, base_url, ids, limits) + if error is not None: + raise AgentRunError(error.message) + click.echo(f"litellm: synced {len(ids)} proxy models into {path}") + return ("--model", f"{PI_PROVIDER_NAME}/{ids[0]}") + + +_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 +) + + def agent_launch_args(command: str, base_url: str) -> list[str]: """Extra CLI args an agent needs to actually honor the proxy. @@ -407,15 +459,14 @@ def run_agent( warn: Callable[[str], None] = _warn, launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off, reattach_terminal: Callable[[], None] | None = None, + preparers: Mapping[str, _Preparer] = MappingProxyType(_PREPARERS), ) -> None: """Validate, wire the environment, and hand off to the agent. - On success this never returns: POSIX replaces the current process, Windows - waits on the agent and exits with its status. Raises AgentRunError for - missing binaries, an unreachable proxy, or a rejected key. The model list is - synced only once the key check passed, so an unreachable proxy costs one - timeout rather than two, and --skip-verify keeps the launch fully offline. - reattach_terminal, when given, runs just before handoff to restore stdin. + On success this replaces the current process and never returns. Raises + AgentRunError for missing binaries, an unreachable proxy, a rejected key, or + a failed pre-launch config sync (pi). reattach_terminal, when given, runs + just before handoff to restore stdin. """ if not command: raise AgentRunError("Nothing to run.") @@ -435,13 +486,16 @@ def run_agent( if isinstance(synced, ModelSyncSkipped): warn(f"litellm: not syncing {display_name} models from the proxy: {synced.reason}") + prepare: Final = preparers.get(os.path.basename(command[0])) + prepared_args: Final = tuple(prepare(base_url, api_key, env_before_sync)) if prepare is not None else () + env: Final = MappingProxyType( { **build_agent_env(env_before_sync, base_url, api_key, profiles), **(_NO_EXTRA_ENV if isinstance(synced, ModelSyncSkipped) else synced), } ) - extra_args: Final = agent_launch_args(command[0], base_url) + extra_args: Final = (*agent_launch_args(command[0], base_url), *prepared_args) if reattach_terminal is not None: reattach_terminal() launcher(binary, [command[0], *extra_args, *command[1:]], env) @@ -501,6 +555,7 @@ def _make_agent_command(binary: str, display_name: str) -> click.Command: name=binary, context_settings={"ignore_unknown_options": True}, short_help=f"Run {display_name} through your LiteLLM proxy", + hidden=binary in _HIDDEN_AGENTS, ) @click.option("--skip-verify", is_flag=True, default=False, help=_SKIP_VERIFY_HELP) @click.argument("args", nargs=-1, type=click.UNPROCESSED) @@ -533,6 +588,7 @@ __all__ = [ "build_agent_env", "opencode_model_sync_env", "opencode_provider_config", + "prepare_pi", "resolve_api_key", "run_agent", "verify_proxy_key", diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py new file mode 100644 index 00000000000..7b0c1970c4e --- /dev/null +++ b/litellm/proxy/client/cli/commands/pi.py @@ -0,0 +1,210 @@ +"""Sync a LiteLLM provider into pi's models.json. + +pi ignores ANTHROPIC_BASE_URL/OPENAI_BASE_URL, so `lite pi` routes it through the +proxy by writing a provider entry instead. The key is stored as a $-reference so +the short-lived login token never lands on disk. +""" + +import json +import os +import tempfile +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import requests +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError + +PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR" +PI_PROVIDER_NAME: Final = "litellm" +LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY" + + +@dataclass(frozen=True, slots=True) +class PiSyncError: + message: str + + +@dataclass(frozen=True, slots=True) +class ModelLimits: + context_window: int | None + max_tokens: int | None + + +class _Model(BaseModel): + id: str + + +class _ModelList(BaseModel): + data: tuple[_Model, ...] + + +class _ModelGroup(BaseModel): + model_group: str + max_input_tokens: float | None = None + max_output_tokens: float | None = None + + +class _ModelGroupList(BaseModel): + data: tuple[_ModelGroup, ...] + + +def fetch_model_ids( + base_url: str, + api_key: str, + *, + get: Callable[..., requests.Response] = requests.get, +) -> tuple[str, ...] | PiSyncError: + url: Final = base_url.rstrip("/") + "/v1/models" + try: + resp: Final = get( + url, + headers={"Authorization": f"Bearer {api_key}"}, # mutable-ok: requests headers require a dict + timeout=10, + ) + except requests.RequestException as e: + return PiSyncError(f"Could not list models from the proxy: {e}") + if resp.status_code != 200: + return PiSyncError(f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot build pi's model list.") + try: + listing: Final = _ModelList.model_validate(resp.json()) + except (ValueError, ValidationError) as e: + return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}") + ids: Final = tuple(dict.fromkeys(model.id for model in listing.data)) + if not ids: + return PiSyncError("The proxy returned no models for your key, so pi would have nothing to run.") + return ids + + +_NO_LIMITS: Final[Mapping[str, ModelLimits]] = MappingProxyType({}) + + +def fetch_model_limits( + base_url: str, + api_key: str, + *, + get: Callable[..., requests.Response] = requests.get, +) -> Mapping[str, ModelLimits]: + """Best effort: pi falls back to its own defaults for models without limits, + so an unavailable /model_group/info must not block the launch.""" + url: Final = base_url.rstrip("/") + "/model_group/info" + try: + resp: Final = get( + url, + headers={"Authorization": f"Bearer {api_key}"}, # mutable-ok: requests headers require a dict + timeout=10, + ) + if resp.status_code != 200: + return _NO_LIMITS + listing: Final = _ModelGroupList.model_validate(resp.json()) + except (requests.RequestException, ValueError, ValidationError): + return _NO_LIMITS + return MappingProxyType( + { + group.model_group: ModelLimits( + context_window=int(group.max_input_tokens) if group.max_input_tokens else None, + max_tokens=int(group.max_output_tokens) if group.max_output_tokens else None, + ) + for group in listing.data + } + ) + + +def models_json_path(env: Mapping[str, str]) -> Path: + override: Final = env.get(PI_CONFIG_DIR_ENV) + root: Final = Path(override) if override else Path.home() / ".pi" / "agent" + return root / "models.json" + + +def _model_entry( + model_id: str, limits: Mapping[str, ModelLimits] +) -> dict[str, JsonValue]: # mutable-ok: JSON object is serialized + limit: Final = limits.get(model_id) + context: Final[dict[str, JsonValue]] = ( # mutable-ok: JSON field + {"contextWindow": limit.context_window} if limit and limit.context_window else {} # mutable-ok: JSON field + ) + output: Final[dict[str, JsonValue]] = ( # mutable-ok: JSON field + {"maxTokens": limit.max_tokens} if limit and limit.max_tokens else {} + ) # mutable-ok: JSON field + return {"id": model_id, **context, **output} # mutable-ok: JSON serialization requires a mutable object + + +def provider_block( + base_url: str, + model_ids: tuple[str, ...], + limits: Mapping[str, ModelLimits] = _NO_LIMITS, +) -> dict[str, JsonValue]: # mutable-ok: JSON object is serialized + """openai-completions is the one API shape every LiteLLM model serves. + + Real contextWindow/maxTokens matter: pi otherwise assumes 128k/16384, which + breaks compaction thresholds and over-asks models with smaller output caps. + """ + return { # mutable-ok: JSON serialization requires a mutable object + "baseUrl": base_url.rstrip("/") + "/v1", + "api": "openai-completions", + "apiKey": f"${LITELLM_PROXY_API_KEY_ENV}", + "models": [_model_entry(model_id, limits) for model_id in model_ids], # mutable-ok: JSON array + } + + +_MODELS_FILE_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) + + +def sync_models_json( + path: Path, + base_url: str, + model_ids: tuple[str, ...], + limits: Mapping[str, ModelLimits] = _NO_LIMITS, +) -> PiSyncError | None: + """Replace only the litellm provider entry, leaving the rest of the file intact.""" + try: + current: Final = ( # mutable-ok: JSON object default + _MODELS_FILE_ADAPTER.validate_json(path.read_text()) if path.exists() else {} + ) + except (OSError, ValidationError) as e: + return PiSyncError(f"Could not read {path} as a JSON object: {e}. Fix or move the file, then retry.") + existing_providers: Final = current.get("providers", {}) # mutable-ok: JSON object default + if not isinstance(existing_providers, dict): + return PiSyncError(f'"providers" in {path} is not an object; fix or move the file, then retry.') + updated: Final = { # mutable-ok: JSON serialization requires a mutable object + **current, + "providers": { # mutable-ok: JSON serialization requires a mutable object + **existing_providers, + PI_PROVIDER_NAME: provider_block(base_url, model_ids, limits), + }, + } + try: + path.parent.mkdir(parents=True, exist_ok=True) + except OSError as e: + return PiSyncError(f"Could not write {path}: {e}") + try: + fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp") + except OSError as e: + return PiSyncError(f"Could not write {path}: {e}") + try: + with os.fdopen(fd, "w") as file: + file.write(json.dumps(updated, indent=2) + "\n") + os.replace(tmp_name, path) + except OSError as e: + try: + os.unlink(tmp_name) + except FileNotFoundError: + pass + return PiSyncError(f"Could not write {path}: {e}") + return None + + +__all__ = ( + "LITELLM_PROXY_API_KEY_ENV", + "PI_CONFIG_DIR_ENV", + "PI_PROVIDER_NAME", + "ModelLimits", + "PiSyncError", + "fetch_model_ids", + "fetch_model_limits", + "models_json_path", + "provider_block", + "sync_models_json", +) diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index 5b99d368cbb..a8a6659fe9a 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -54,6 +54,15 @@ class _Recorder: return self.returns +class _FakeJsonResponse: + def __init__(self, status_code, payload=None): + self.status_code = status_code + self._payload = payload + + def json(self): + return self._payload + + class TestAgentProfile: def test_claude_is_anthropic(self): name, profiles = agent_profile("claude") @@ -69,6 +78,9 @@ class TestAgentProfile: assert agent_profile("codex") == ("Codex", frozenset({"openai"})) assert agent_profile("opencode") == ("OpenCode", frozenset({"openai"})) + def test_pi_is_litellm(self): + assert agent_profile("pi") == ("pi", frozenset({"litellm"})) + def test_unknown_command_gets_both_profiles(self): name, profiles = agent_profile("mytool") assert name == "mytool" @@ -134,6 +146,15 @@ class TestBuildAgentEnv: assert env["OPENAI_API_KEY"] == "sk-key" assert env["ENABLE_TOOL_SEARCH"] == "true" + def test_litellm_profile_exports_only_the_proxy_key(self): + env = build_agent_env( + {}, "http://localhost:4000/", "sk-key", frozenset({"litellm"}) + ) + assert env["LITELLM_PROXY_API_KEY"] == "sk-key" + assert "ANTHROPIC_BASE_URL" not in env + assert "OPENAI_BASE_URL" not in env + assert "OPENAI_API_KEY" not in env + def test_preserves_unrelated_env_and_does_not_mutate_input(self): base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} env = build_agent_env( @@ -166,6 +187,9 @@ class TestAgentLaunchArgs: agent_launch_args("codex", "http://localhost:4000") ) + def test_pi_gets_no_static_args(self): + assert agent_launch_args("pi", "http://localhost:4000") == [] + class TestVerifyProxyKey: def test_ok_status_passes_and_uses_models_endpoint(self): @@ -520,6 +544,126 @@ class TestRunAgent: # overrides must precede the codex subcommand so codex parses them assert args.index('model_provider="litellm"') < args.index("exec") + def test_pi_preparer_runs_after_verify_and_before_launch(self): + order = [] + captured = {} + + def fake_prepare(base_url, api_key, base_env): + order.append("prepare") + captured["args"] = (base_url, api_key, dict(base_env)) + return [] + + run_agent( + "http://localhost:4000", + "sk-key", + ["pi"], + base_env={"HOME": "/home/u"}, + which=lambda name: "/usr/local/bin/pi", + verify=lambda *a: order.append("verify"), + launcher=lambda *a: order.append("launch"), + preparers={"pi": fake_prepare}, + ) + assert order == ["verify", "prepare", "launch"] + assert captured["args"] == ( + "http://localhost:4000", + "sk-key", + {"HOME": "/home/u"}, + ) + + def test_pi_prepared_args_precede_user_args_and_env_has_proxy_key(self): + calls = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["pi", "-p", "hello"], + base_env={}, + which=lambda name: "/usr/local/bin/pi", + verify=lambda *a: None, + launcher=lambda p, a, e: calls.update(args=tuple(a), env=dict(e)), + preparers={"pi": lambda *a: ["--model", "litellm/m-1"]}, + ) + assert calls["args"] == ("pi", "--model", "litellm/m-1", "-p", "hello") + assert calls["env"]["LITELLM_PROXY_API_KEY"] == "sk-key" + assert "OPENAI_API_KEY" not in calls["env"] + assert "ANTHROPIC_BASE_URL" not in calls["env"] + + def test_failed_preparer_aborts_before_launch(self): + launched = [] + + def boom(*a): + raise AgentRunError("sync failed") + + with pytest.raises(AgentRunError, match="sync failed"): + run_agent( + "http://localhost:4000", + "sk-key", + ["pi"], + base_env={}, + which=lambda name: "/usr/local/bin/pi", + verify=lambda *a: None, + launcher=lambda *a: launched.append(a), + preparers={"pi": boom}, + ) + assert launched == [] + + def test_prepare_pi_syncs_models_json_and_pins_first_model(self, tmp_path): + from litellm.proxy.client.cli.commands.agents import prepare_pi + + def fake_get(url, headers, timeout): + if url.endswith("/model_group/info"): + return _FakeJsonResponse( + 200, + {"data": [{"model_group": "m-first", "max_input_tokens": 131072, "max_output_tokens": 8192}]}, + ) + return _FakeJsonResponse(200, {"data": [{"id": "m-first"}, {"id": "m-second"}]}) + + pin = prepare_pi( + "http://localhost:4000", + "sk-key", + {"PI_CODING_AGENT_DIR": str(tmp_path)}, + get=fake_get, + ) + + assert pin == ("--model", "litellm/m-first") + import json + + written = json.loads((tmp_path / "models.json").read_text()) + assert written["providers"]["litellm"]["apiKey"] == "$LITELLM_PROXY_API_KEY" + assert written["providers"]["litellm"]["models"] == [ + {"id": "m-first", "contextWindow": 131072, "maxTokens": 8192}, + {"id": "m-second"}, + ] + + def test_prepare_pi_surfaces_fetch_failure_as_agent_error(self, tmp_path): + from litellm.proxy.client.cli.commands.agents import prepare_pi + + with pytest.raises(AgentRunError, match="HTTP 500"): + prepare_pi( + "http://localhost:4000", + "sk-key", + {"PI_CODING_AGENT_DIR": str(tmp_path)}, + get=lambda *a, **k: _FakeJsonResponse(500), + ) + + def test_claude_has_no_preparer(self): + prepared = [] + + def fake_prepare(*a): + prepared.append(a) + return [] + + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda *a: None, + launcher=lambda *a: None, + preparers={"pi": fake_prepare}, + ) + assert prepared == [] + def test_claude_launches_without_injected_args(self): calls = {} run_agent( @@ -877,7 +1021,11 @@ class TestAgentCommands: self.runner = CliRunner() def test_one_command_per_known_agent(self): - assert {c.name for c in agent_commands()} == {"claude", "codex", "opencode"} + assert {c.name for c in agent_commands()} == {"claude", "codex", "opencode", "pi"} + + def test_pi_is_hidden_from_help_but_still_registered(self): + hidden_by_name = {c.name: c.hidden for c in agent_commands()} + assert hidden_by_name == {"claude": False, "codex": False, "opencode": False, "pi": True} def test_claude_launches_with_stored_key_and_forwards_args(self): captured = {} diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py new file mode 100644 index 00000000000..68c0ac70064 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -0,0 +1,236 @@ +import json +import os +import stat +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import requests + +from litellm.proxy.client.cli.commands.pi import ( + ModelLimits, + PiSyncError, + fetch_model_ids, + fetch_model_limits, + models_json_path, + provider_block, + sync_models_json, +) + + +class _FakeResponse: + def __init__(self, status_code, payload=None): + self.status_code = status_code + self._payload = payload + + def json(self): + if self._payload is None: + raise ValueError("not json") + return self._payload + + +class TestFetchModelIds: + def test_returns_ids_in_proxy_order_deduped(self): + captured = {} + + def fake_get(url, headers, timeout): + captured["url"] = url + captured["headers"] = headers + return _FakeResponse( + 200, + {"data": [{"id": "m-b"}, {"id": "m-a"}, {"id": "m-b"}]}, + ) + + assert fetch_model_ids("http://localhost:4000/", "sk-key", get=fake_get) == ("m-b", "m-a") + assert captured["url"] == "http://localhost:4000/v1/models" + assert captured["headers"] == {"Authorization": "Bearer sk-key"} + + def test_network_error_is_a_value(self): + def boom(*a, **k): + raise requests.ConnectionError("refused") + + result = fetch_model_ids("http://localhost:4000", "sk-key", get=boom) + assert isinstance(result, PiSyncError) + assert "Could not list models" in result.message + + def test_non_200_is_a_value(self): + result = fetch_model_ids( + "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500) + ) + assert isinstance(result, PiSyncError) + assert "HTTP 500" in result.message + + def test_malformed_body_is_a_value(self): + result = fetch_model_ids( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse(200, {"data": "nope"}), + ) + assert isinstance(result, PiSyncError) + + def test_empty_model_list_is_a_value(self): + result = fetch_model_ids( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse(200, {"data": []}), + ) + assert isinstance(result, PiSyncError) + assert "no models" in result.message + + +class TestFetchModelLimits: + def test_maps_group_limits_and_hits_model_group_info(self): + captured = {} + + def fake_get(url, headers, timeout): + captured["url"] = url + return _FakeResponse( + 200, + { + "data": [ + {"model_group": "m-a", "max_input_tokens": 131072, "max_output_tokens": 8192}, + {"model_group": "m-b", "max_input_tokens": None, "max_output_tokens": None}, + ] + }, + ) + + limits = fetch_model_limits("http://localhost:4000/", "sk-key", get=fake_get) + assert captured["url"] == "http://localhost:4000/model_group/info" + assert limits["m-a"] == ModelLimits(context_window=131072, max_tokens=8192) + assert limits["m-b"] == ModelLimits(context_window=None, max_tokens=None) + + def test_non_200_degrades_to_no_limits(self): + assert fetch_model_limits("http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(403)) == {} + + def test_network_error_degrades_to_no_limits(self): + def boom(*a, **k): + raise requests.ConnectionError("refused") + + assert fetch_model_limits("http://localhost:4000", "sk-key", get=boom) == {} + + def test_malformed_body_degrades_to_no_limits(self): + assert ( + fetch_model_limits( + "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(200, {"data": "nope"}) + ) + == {} + ) + + +class TestModelsJsonPath: + def test_env_override_wins(self): + assert models_json_path({"PI_CODING_AGENT_DIR": "/custom/dir"}) == Path("/custom/dir/models.json") + + def test_defaults_to_home_pi_agent(self): + assert models_json_path({}) == Path.home() / ".pi" / "agent" / "models.json" + + +class TestProviderBlock: + def test_points_pi_at_proxy_with_env_interpolated_key(self): + block = provider_block("http://localhost:4000/", ("m-1", "m-2")) + assert block == { + "baseUrl": "http://localhost:4000/v1", + "api": "openai-completions", + "apiKey": "$LITELLM_PROXY_API_KEY", + "models": [{"id": "m-1"}, {"id": "m-2"}], + } + + def test_known_limits_become_context_window_and_max_tokens(self): + block = provider_block( + "http://localhost:4000", + ("m-1", "m-2"), + { + "m-1": ModelLimits(context_window=131072, max_tokens=8192), + "m-2": ModelLimits(context_window=None, max_tokens=None), + }, + ) + assert block["models"] == [ + {"id": "m-1", "contextWindow": 131072, "maxTokens": 8192}, + {"id": "m-2"}, + ] + + +class TestSyncModelsJson: + def test_creates_file_and_parent_dirs(self, tmp_path): + path = tmp_path / "agent" / "models.json" + assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None + written = json.loads(path.read_text()) + assert written["providers"]["litellm"]["baseUrl"] == "http://localhost:4000/v1" + assert written["providers"]["litellm"]["models"] == [{"id": "m-1"}] + + def test_preserves_other_providers_and_top_level_keys(self, tmp_path): + path = tmp_path / "models.json" + path.write_text( + json.dumps( + { + "somethingElse": True, + "providers": { + "ollama": {"baseUrl": "http://localhost:11434/v1"}, + "litellm": {"baseUrl": "http://stale:1234/v1", "models": []}, + }, + } + ) + ) + assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None + written = json.loads(path.read_text()) + assert written["somethingElse"] is True + assert written["providers"]["ollama"] == {"baseUrl": "http://localhost:11434/v1"} + assert written["providers"]["litellm"]["baseUrl"] == "http://localhost:4000/v1" + assert written["providers"]["litellm"]["models"] == [{"id": "m-1"}] + + def test_write_leaves_no_staging_file_behind(self, tmp_path): + path = tmp_path / "models.json" + assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None + assert [p.name for p in tmp_path.iterdir()] == ["models.json"] + + def test_written_file_is_private(self, tmp_path): + path = tmp_path / "models.json" + assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None + if os.name != "nt": + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + + path.write_text(json.dumps({"providers": {"other": {"apiKey": "literal-secret"}}})) + path.chmod(0o644) + assert sync_models_json(path, "http://localhost:4000", ("m-2",)) is None + if os.name != "nt": + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + + def test_concurrent_syncs_do_not_collide(self, tmp_path): + path = tmp_path / "models.json" + model_lists = (("m-a",), ("m-b",)) + + def sync(model_ids): + return sync_models_json(path, "http://localhost:4000", model_ids) + + with ThreadPoolExecutor(max_workers=2) as executor: + results = [result for _ in range(30) for result in executor.map(sync, model_lists)] + + assert results == [None] * 60 + written = json.loads(path.read_text()) + assert written["providers"]["litellm"]["models"] in ([{"id": "m-a"}], [{"id": "m-b"}]) + assert list(tmp_path.glob("models.json.*.tmp")) == [] + + def test_invalid_json_is_a_value_and_file_untouched(self, tmp_path): + path = tmp_path / "models.json" + path.write_text("{not json") + result = sync_models_json(path, "http://localhost:4000", ("m-1",)) + assert isinstance(result, PiSyncError) + assert path.read_text() == "{not json" + + def test_non_object_providers_is_a_value(self, tmp_path): + path = tmp_path / "models.json" + path.write_text(json.dumps({"providers": ["nope"]})) + result = sync_models_json(path, "http://localhost:4000", ("m-1",)) + assert isinstance(result, PiSyncError) + + def test_top_level_non_object_is_a_value(self, tmp_path): + path = tmp_path / "models.json" + path.write_text(json.dumps(["nope"])) + result = sync_models_json(path, "http://localhost:4000", ("m-1",)) + assert isinstance(result, PiSyncError) + + def test_unwritable_path_is_a_value(self, tmp_path): + blocker = tmp_path / "agent" + blocker.write_text("i am a file, not a directory") + result = sync_models_json(blocker / "models.json", "http://localhost:4000", ("m-1",)) + assert isinstance(result, PiSyncError) + assert "Could not" in result.message