scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #36841 from BerriAI/litellm_lite_pi

Browse source 
feat(cli): add lite pi to run the pi coding agent through the proxy
This commit is contained in:
ryan-crabbe-berri 2026-09-05 17:14:11 -07:00 • committed by GitHub
commit e1fb8affe3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 664 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/proxy/client/cli/README.md
View file

@ -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/<id>`. 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/<other-id>` 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

72
litellm/proxy/client/cli/commands/agents.py
View file

@ -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",

210
litellm/proxy/client/cli/commands/pi.py Normal file
View file

@ -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",
)

150
tests/test_litellm/proxy/client/cli/test_agents.py
View file

@ -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 = {}

236
tests/test_litellm/proxy/client/cli/test_pi.py Normal file
View file

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