From 4e280ecc344f67f2f04d791b8990886acbfbb83f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 13 Aug 2026 15:33:36 -0700 Subject: [PATCH 1/5] feat(cli): add lite pi to run the pi coding agent through the proxy --- litellm/proxy/client/cli/README.md | 7 +- litellm/proxy/client/cli/commands/agents.py | 71 ++++++- litellm/proxy/client/cli/commands/pi.py | 178 ++++++++++++++++ .../proxy/client/cli/test_agents.py | 147 ++++++++++++- .../test_litellm/proxy/client/cli/test_pi.py | 201 ++++++++++++++++++ 5 files changed, 591 insertions(+), 13 deletions(-) create mode 100644 litellm/proxy/client/cli/commands/pi.py create mode 100644 tests/test_litellm/proxy/client/cli/test_pi.py diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index de9d38963c1..66beecbd2e6 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -467,6 +467,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: @@ -480,17 +481,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. 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). +pi ignores base-URL environment variables entirely, so `lite pi` 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. 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. 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 dfc70a8df7c..f55cf9893e7 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -2,12 +2,22 @@ import os import shutil import sys from collections.abc import Callable, Mapping, Sequence +from types import MappingProxyType from typing import Final import click import requests from .auth import get_stored_api_key, login +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" @@ -17,17 +27,20 @@ OPENAI_API_KEY_ENV: Final = "OPENAI_API_KEY" 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", } CODEX_PROXY_PROVIDER: Final = "litellm" @@ -60,7 +73,9 @@ def build_agent_env( Anthropic clients (Claude Code) append /v1/messages to ANTHROPIC_BASE_URL, so it stays the bare proxy root; OpenAI clients (Codex, OpenCode) expect the /v1 suffix on OPENAI_BASE_URL. ANTHROPIC_API_KEY is dropped so a stray - Anthropic key cannot win over the bearer token we set. + Anthropic key cannot win over the bearer token we 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("/") @@ -71,6 +86,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 @@ -106,6 +123,38 @@ _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, +) -> list[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]}"] + + +_PREPARERS: Final[dict[str, Callable[[str, str, Mapping[str, str]], Sequence[str]]]] = { + "pi": prepare_pi, +} + + def agent_launch_args(command: str, base_url: str) -> list[str]: """Extra CLI args an agent needs to actually honor the proxy. @@ -177,12 +226,14 @@ def run_agent( verify: Callable[[str, str], None] = verify_proxy_key, launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _exec, reattach_terminal: Callable[[], None] | None = None, + preparers: Mapping[str, Callable[[str, str, Mapping[str, str]], Sequence[str]]] = MappingProxyType(_PREPARERS), ) -> None: """Validate, wire the environment, and hand off to the agent. On success this replaces the current process and never returns. Raises - AgentRunError for missing binaries, an unreachable proxy, or a rejected key. - reattach_terminal, when given, runs just before handoff to restore stdin. + 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.") @@ -197,13 +248,12 @@ def run_agent( if not skip_verify: verify(base_url, api_key) - env: Final = build_agent_env( - base_env if base_env is not None else os.environ, - base_url, - api_key, - profiles, - ) - extra_args: Final = agent_launch_args(command[0], base_url) + source_env: Final = base_env if base_env is not None else os.environ + prepare: Final = preparers.get(os.path.basename(command[0])) + prepared_args: Final = list(prepare(base_url, api_key, source_env)) if prepare is not None else [] + + env: Final = build_agent_env(source_env, base_url, api_key, profiles) + 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) @@ -288,6 +338,7 @@ __all__ = [ "agent_launch_args", "agent_profile", "build_agent_env", + "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..bd933874a10 --- /dev/null +++ b/litellm/proxy/client/cli/commands/pi.py @@ -0,0 +1,178 @@ +"""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 +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: list[_Model] + + +class _ModelGroup(BaseModel): + model_group: str + max_input_tokens: float | None = None + max_output_tokens: float | None = None + + +class _ModelGroupList(BaseModel): + data: list[_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}"}, 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}"}, 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]: + limit: Final = limits.get(model_id) + context: Final[dict[str, JsonValue]] = ( + {"contextWindow": limit.context_window} if limit and limit.context_window else {} + ) + output: Final[dict[str, JsonValue]] = {"maxTokens": limit.max_tokens} if limit and limit.max_tokens else {} + return {"id": model_id, **context, **output} + + +def provider_block( + base_url: str, + model_ids: tuple[str, ...], + limits: Mapping[str, ModelLimits] = _NO_LIMITS, +) -> dict[str, JsonValue]: + """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 { + "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], + } + + +_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 = _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", {}) + 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 = { + **current, + "providers": {**existing_providers, PI_PROVIDER_NAME: provider_block(base_url, model_ids, limits)}, + } + try: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(updated, indent=2) + "\n") + except OSError as e: + 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 afd1696a89f..fad0d3842fa 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -34,6 +34,15 @@ class _FakeResponse: self.status_code = status_code +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") @@ -49,6 +58,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" @@ -91,6 +103,15 @@ class TestBuildAgentEnv: assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" assert env["OPENAI_API_KEY"] == "sk-key" + 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( @@ -123,6 +144,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): @@ -223,6 +247,127 @@ 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"]}, + ) + # user args come last so a user-supplied --model wins in pi's parser + 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( @@ -319,7 +464,7 @@ 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_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..99ed2734b76 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -0,0 +1,201 @@ +import json +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_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 From d3dc6f6b12324186f15963f8eb12f02662461784 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 13 Aug 2026 16:15:51 -0700 Subject: [PATCH 2/5] fix(cli): write pi models.json atomically and hide lite pi from --help --- litellm/proxy/client/cli/README.md | 2 +- litellm/proxy/client/cli/commands/agents.py | 11 ++++++++--- litellm/proxy/client/cli/commands/pi.py | 4 +++- tests/test_litellm/proxy/client/cli/test_agents.py | 4 ++++ tests/test_litellm/proxy/client/cli/test_pi.py | 5 +++++ 5 files changed, 21 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 66beecbd2e6..468b8d96123 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -481,7 +481,7 @@ 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. 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). -pi ignores base-URL environment variables entirely, so `lite pi` 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. +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): diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index f55cf9893e7..ba410673ec9 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -3,7 +3,7 @@ import shutil import sys from collections.abc import Callable, Mapping, Sequence from types import MappingProxyType -from typing import Final +from typing import Final, TypeAlias import click import requests @@ -43,6 +43,8 @@ _INSTALL_DOCS: Final[dict[str, str]] = { "pi": "https://pi.dev", } +_HIDDEN_AGENTS: Final = frozenset({"pi"}) + CODEX_PROXY_PROVIDER: Final = "litellm" @@ -150,7 +152,9 @@ def prepare_pi( return ["--model", f"{PI_PROVIDER_NAME}/{ids[0]}"] -_PREPARERS: Final[dict[str, Callable[[str, str, Mapping[str, str]], Sequence[str]]]] = { +_Preparer: TypeAlias = Callable[[str, str, Mapping[str, str]], Sequence[str]] + +_PREPARERS: Final[dict[str, _Preparer]] = { "pi": prepare_pi, } @@ -226,7 +230,7 @@ def run_agent( verify: Callable[[str, str], None] = verify_proxy_key, launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _exec, reattach_terminal: Callable[[], None] | None = None, - preparers: Mapping[str, Callable[[str, str, Mapping[str, str]], Sequence[str]]] = MappingProxyType(_PREPARERS), + preparers: Mapping[str, _Preparer] = MappingProxyType(_PREPARERS), ) -> None: """Validate, wire the environment, and hand off to the agent. @@ -311,6 +315,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) diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index bd933874a10..668b803a33d 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -156,9 +156,11 @@ def sync_models_json( **current, "providers": {**existing_providers, PI_PROVIDER_NAME: provider_block(base_url, model_ids, limits)}, } + staging: Final = path.with_name(path.name + ".tmp") try: path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(updated, indent=2) + "\n") + staging.write_text(json.dumps(updated, indent=2) + "\n") + staging.replace(path) except OSError as e: return PiSyncError(f"Could not write {path}: {e}") return None diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index fad0d3842fa..99b2ff234e2 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -466,6 +466,10 @@ class TestAgentCommands: def test_one_command_per_known_agent(self): 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 index 99ed2734b76..ee3222a77e3 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -174,6 +174,11 @@ class TestSyncModelsJson: 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_invalid_json_is_a_value_and_file_untouched(self, tmp_path): path = tmp_path / "models.json" path.write_text("{not json") From b57cee65cc7b54136a3027a036445b601e5c2756 Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 5 Sep 2026 00:22:42 +0000 Subject: [PATCH 3/5] fix(cli): write pi models atomically and privately Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/client/cli/commands/pi.py | 19 ++++++++++-- .../test_litellm/proxy/client/cli/test_pi.py | 30 +++++++++++++++++++ 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index 668b803a33d..b3b9d520a5a 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -6,6 +6,8 @@ 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 @@ -156,13 +158,24 @@ def sync_models_json( **current, "providers": {**existing_providers, PI_PROVIDER_NAME: provider_block(base_url, model_ids, limits)}, } - staging: Final = path.with_name(path.name + ".tmp") try: path.parent.mkdir(parents=True, exist_ok=True) - staging.write_text(json.dumps(updated, indent=2) + "\n") - staging.replace(path) 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 diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index ee3222a77e3..68c0ac70064 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -1,4 +1,7 @@ import json +import os +import stat +from concurrent.futures import ThreadPoolExecutor from pathlib import Path import requests @@ -179,6 +182,33 @@ class TestSyncModelsJson: 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") From 7a9f466657c457b6c2eca7cceb47b2862db934ad Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 5 Sep 2026 00:37:04 +0000 Subject: [PATCH 4/5] fix(cli): satisfy pi type discipline budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/client/cli/commands/agents.py | 10 ++-- litellm/proxy/client/cli/commands/pi.py | 53 ++++++++++++------- .../proxy/client/cli/test_agents.py | 2 +- 3 files changed, 41 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index d9b8117c994..3dace0b94f4 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -154,7 +154,7 @@ def prepare_pi( base_env: Mapping[str, str], *, get: Callable[..., requests.Response] = requests.get, -) -> list[str]: +) -> 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 @@ -172,14 +172,14 @@ def prepare_pi( 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]}"] + return ("--model", f"{PI_PROVIDER_NAME}/{ids[0]}") _Preparer: TypeAlias = Callable[[str, str, Mapping[str, str]], Sequence[str]] -_PREPARERS: Final[dict[str, _Preparer]] = { - "pi": prepare_pi, -} +_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]: diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index b3b9d520a5a..7b0c1970c4e 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -38,7 +38,7 @@ class _Model(BaseModel): class _ModelList(BaseModel): - data: list[_Model] + data: tuple[_Model, ...] class _ModelGroup(BaseModel): @@ -48,7 +48,7 @@ class _ModelGroup(BaseModel): class _ModelGroupList(BaseModel): - data: list[_ModelGroup] + data: tuple[_ModelGroup, ...] def fetch_model_ids( @@ -59,7 +59,11 @@ def fetch_model_ids( ) -> tuple[str, ...] | PiSyncError: url: Final = base_url.rstrip("/") + "/v1/models" try: - resp: Final = get(url, headers={"Authorization": f"Bearer {api_key}"}, timeout=10) + 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: @@ -87,7 +91,11 @@ def fetch_model_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}"}, timeout=10) + 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()) @@ -110,30 +118,34 @@ def models_json_path(env: Mapping[str, str]) -> Path: return root / "models.json" -def _model_entry(model_id: str, limits: Mapping[str, ModelLimits]) -> dict[str, JsonValue]: +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]] = ( - {"contextWindow": limit.context_window} if limit and limit.context_window else {} + 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]] = {"maxTokens": limit.max_tokens} if limit and limit.max_tokens else {} - return {"id": model_id, **context, **output} + 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]: +) -> 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 { + 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], + "models": [_model_entry(model_id, limits) for model_id in model_ids], # mutable-ok: JSON array } @@ -148,15 +160,20 @@ def sync_models_json( ) -> PiSyncError | None: """Replace only the litellm provider entry, leaving the rest of the file intact.""" try: - current: Final = _MODELS_FILE_ADAPTER.validate_json(path.read_text()) if path.exists() else {} + 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", {}) + 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 = { + updated: Final = { # mutable-ok: JSON serialization requires a mutable object **current, - "providers": {**existing_providers, PI_PROVIDER_NAME: provider_block(base_url, model_ids, limits)}, + "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) @@ -179,7 +196,7 @@ def sync_models_json( return None -__all__ = [ +__all__ = ( "LITELLM_PROXY_API_KEY_ENV", "PI_CONFIG_DIR_ENV", "PI_PROVIDER_NAME", @@ -190,4 +207,4 @@ __all__ = [ "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 bd18a5c6f94..a64adffb3a1 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -625,7 +625,7 @@ class TestRunAgent: get=fake_get, ) - assert pin == ["--model", "litellm/m-first"] + assert pin == ("--model", "litellm/m-first") import json written = json.loads((tmp_path / "models.json").read_text()) From 35ae1ca151643af4c9bc84b649fba7df5d3f7674 Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 5 Sep 2026 01:01:41 +0000 Subject: [PATCH 5/5] test(cli): drop redundant comment in pi arg ordering test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/client/cli/test_agents.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index a64adffb3a1..a8a6659fe9a 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -582,7 +582,6 @@ class TestRunAgent: launcher=lambda p, a, e: calls.update(args=tuple(a), env=dict(e)), preparers={"pi": lambda *a: ["--model", "litellm/m-1"]}, ) - # user args come last so a user-supplied --model wins in pi's parser 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"]