From 961f40bb22d436d29cda1eec4715445bb315da9f Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 9 Sep 2026 11:14:15 -0700 Subject: [PATCH] feat(cli): add lite configure codex, --launch, and CLAUDE_CONFIG_DIR/CODEX_HOME Codex gets the same persistent, receipt-based wiring as Claude Code: a litellm model provider in config.toml on the Responses transport with the credential inline or read through lite auth print-token, an optional pinned model with its context window from the proxy, and lite unconfigure codex. The configure/unconfigure core moves into agent_config.py as one codec-agnostic mechanism (JSON settings, tomlkit TOML that keeps comments), configure claude/codex take --launch, plain lite configure lists both agents and offers to start one, and the config directories honor CLAUDE_CONFIG_DIR and CODEX_HOME --- litellm/litellm_core_utils/private_json.py | 33 +- litellm/proxy/client/cli/README.md | 6 +- .../proxy/client/cli/commands/agent_config.py | 650 ++++++++++++++++++ litellm/proxy/client/cli/commands/agents.py | 52 +- .../client/cli/commands/claude_settings.py | 50 +- .../client/cli/commands/codex_settings.py | 168 +++++ .../proxy/client/cli/commands/configure.py | 363 ++++++---- .../client/cli/commands/configured_launch.py | 133 ++++ litellm/proxy/client/cli/commands/up.py | 4 +- pyproject.toml | 2 + .../proxy/client/cli/test_claude_settings.py | 52 +- .../proxy/client/cli/test_codex_settings.py | 332 +++++++++ .../client/cli/test_configure_commands.py | 327 ++++++++- .../client/cli/test_configured_launch.py | 270 ++++++++ uv.lock | 4 + 15 files changed, 2267 insertions(+), 179 deletions(-) create mode 100644 litellm/proxy/client/cli/commands/agent_config.py create mode 100644 litellm/proxy/client/cli/commands/codex_settings.py create mode 100644 litellm/proxy/client/cli/commands/configured_launch.py create mode 100644 tests/test_litellm/proxy/client/cli/test_codex_settings.py create mode 100644 tests/test_litellm/proxy/client/cli/test_configured_launch.py diff --git a/litellm/litellm_core_utils/private_json.py b/litellm/litellm_core_utils/private_json.py index 4cd4a9b4f82..e513d61ff95 100644 --- a/litellm/litellm_core_utils/private_json.py +++ b/litellm/litellm_core_utils/private_json.py @@ -2,9 +2,9 @@ import json import os import stat import tempfile -from collections.abc import Mapping +from collections.abc import Callable, Mapping from pathlib import Path -from typing import Final +from typing import IO, Final PRIVATE_DIR_MODE: Final = 0o700 @@ -16,18 +16,13 @@ def ensure_private_dir(directory: Path) -> None: directory.chmod(PRIVATE_DIR_MODE) -def stage_private_json(path: str, data: Mapping[str, object]) -> str: - """Write JSON to a private temp file beside `path`, ready for `commit_staged_json`. - - Staging is the half that can fail on a read-only or full directory, so callers with something - to lose can find that out before they act on the assumption that the rewrite will land. - """ +def _stage(path: str, write: Callable[[IO[str]], None]) -> str: parent: Final = Path(path).parent parent.mkdir(parents=True, exist_ok=True) - fd, tmp_path = tempfile.mkstemp(dir=str(parent), prefix=".tmp-", suffix=".json") + fd, tmp_path = tempfile.mkstemp(dir=str(parent), prefix=".tmp-", suffix=Path(path).suffix or ".json") try: with os.fdopen(fd, "w") as f: - json.dump(data, f, indent=2) + write(f) f.flush() os.fsync(f.fileno()) except BaseException: @@ -36,6 +31,24 @@ def stage_private_json(path: str, data: Mapping[str, object]) -> str: return tmp_path +def stage_private_text(path: str, text: str) -> str: + """Write text to a private temp file beside `path`, ready for `commit_staged_json`. + + Staging is the half that can fail on a read-only or full directory, so callers with something + to lose can find that out before they act on the assumption that the rewrite will land. + """ + + def write(f: IO[str]) -> None: + f.write(text) + + return _stage(path, write) + + +def stage_private_json(path: str, data: Mapping[str, object]) -> str: + """Write JSON to a private temp file beside `path`, ready for `commit_staged_json`""" + return _stage(path, lambda f: json.dump(data, f, indent=2)) + + def stage_private_bytes(path: str, data: bytes) -> str: parent: Final = Path(path).parent parent.mkdir(parents=True, exist_ok=True) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 2071576a943..2782c94ddcd 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -542,13 +542,15 @@ The key in the file is the login's own, so it expires with it (24h by default): ```bash curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh -lite --base-url https://your-proxy.example.com configure claude --api-key sk-... --model claude-auto +lite --base-url https://your-proxy.example.com configure --api-key sk-... claude --model claude-auto claude ``` The key comes from `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) and is written into `env.ANTHROPIC_AUTH_TOKEN`; without one the command refuses, since a `lite login` credential expires within a day and keeping it fresh would mean Claude Code running `lite` through `apiKeyHelper` on every credential refresh. The command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (under `claude-router-` for a group whose id contains neither `claude` nor `anthropic`, since Claude Code lists only those) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key and as `env.ANTHROPIC_MODEL`, both of which have to be on `/v1/models` for the key. The second one matters for `claude -c` and `claude --resume`: a resumed session otherwise re-sends the model its transcript recorded, which behind an auto-router with `return_raw_model_name: true` is the tier model that answered, and a key scoped to the router alias gets a 403 for it; `ANTHROPIC_MODEL` outranks the transcript on resume. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window (the proxy appends `[1m]` for a group whose configured or known input window reaches 1M) and sends no thinking parameters for it, so name the group like a Claude model id to change that. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control -Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt +Plain `lite configure --api-key `, with no agent named, asks which installed agents to connect and which gateway model each starts on. After both settings files are written, it asks which configured agents to open; one selected agent takes over the current terminal, while Claude Code and Codex selected together open in separate terminal windows. Both clients read the saved configuration directly, so the launch does not need another gateway setup step or expose the key in its command line + +`lite configure --api-key claude --launch` and `lite configure --api-key codex --launch` provide the same single-agent flow without prompts. `lite configure --api-key claude` or `codex` saves the configuration without opening the client What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any request diff --git a/litellm/proxy/client/cli/commands/agent_config.py b/litellm/proxy/client/cli/commands/agent_config.py new file mode 100644 index 00000000000..330d6a30e46 --- /dev/null +++ b/litellm/proxy/client/cli/commands/agent_config.py @@ -0,0 +1,650 @@ +"""Receipt-based configure and unconfigure for Codex's TOML config.""" + +import hashlib +import json +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from datetime import date, datetime, time +from pathlib import Path +from types import MappingProxyType +from typing import Final, Protocol, TypeAlias + +import tomlkit +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError +from tomlkit.container import OutOfOrderTableProxy +from tomlkit.exceptions import ParseError +from tomlkit.items import InlineTable, Item, Table + +from litellm.litellm_core_utils.private_json import ( + commit_staged_json, + discard_staged_json, + ensure_private_dir, + stage_private_text, +) + +ROOT_SECTION: Final = "" + + +class AgentConfigError(Exception): + """Raised for any user-actionable failure while reading or writing an agent's config.""" + + +@dataclass(frozen=True, slots=True) +class SettingsFileOwner: + """A command that takes temporary ownership of a config file and restores it later.""" + + backup_path: Path + start_command: str + stop_command: str + + +class OwnedValue(BaseModel): + """What one key held at a moment in time; `present=False` is an absent key, not a null one.""" + + model_config = ConfigDict(frozen=True) + + present: bool + value: JsonValue = None + + +class SectionReceipt(BaseModel): + """One owned section: whether it existed as an object, what its owned keys held, what was written.""" + + model_config = ConfigDict(frozen=True) + + present: bool + was_object: bool + previous: Mapping[str, OwnedValue] + written: Mapping[str, str] + + +class ConfigureReceipt(BaseModel): + """What configure found and what it wrote, so unconfigure can undo only its own work. + + `previous` values are the ones every owned key had before the first configure; a repeat + configure keeps them, since the values it would otherwise snapshot are its own. `written` + holds fingerprints, so unconfigure can tell a key it still owns from one the user changed + since, without keeping a second copy of a credential on disk. The file and section shapes + are recorded too, so a file that did not exist, or a section that was absent or null, comes + back exactly that way. + """ + + model_config = ConfigDict(frozen=True) + + file_existed: bool + sections: Mapping[str, SectionReceipt] + + +@dataclass(frozen=True, slots=True) +class UnconfigureOutcome: + """Which owned keys unconfigure put back, which it left because the user had changed them, and which + credentials it left removed because the endpoint they belonged to was changed after configure.""" + + restored: tuple[str, ...] + kept: tuple[str, ...] + withheld: tuple[str, ...] = () + file_removed: bool = False + + +@dataclass(frozen=True, slots=True) +class RestoreGroup: + """Keys that only come back together: a credential is never restored next to an endpoint the user changed. + + `anchor` is the (section, key) of the endpoint; `dependents` are the credential slots. When the + anchor no longer holds what configure wrote, the dependents stay removed instead of restored. + """ + + anchor: tuple[str, str] + dependents: tuple[tuple[str, str], ...] + + +class ConfigDocument(Protocol): + """A parsed config file, edited in place so the codec can keep the user's formatting.""" + + def section(self, name: str) -> Mapping[str, JsonValue] | None: ... + + def section_is_object(self, name: str) -> bool: ... + + def section_is_scalar(self, name: str) -> bool: ... + + def set_value(self, section: str, key: str, value: object) -> None: ... + + def delete(self, section: str, key: str) -> None: ... + + def null_section(self, name: str) -> None: ... + + def drop_section(self, name: str) -> None: ... + + def is_empty(self) -> bool: ... + + def dumps(self) -> str: ... + + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +_TOML_MAPPINGS: Final = (Table, InlineTable, OutOfOrderTableProxy) + + +def _plain(value: object) -> JsonValue: + """A tomlkit item as the JSON value the receipt fingerprints; dates become their ISO text.""" + raw: Final = value.unwrap() if isinstance(value, (Item, OutOfOrderTableProxy)) else value + return _json_value(raw) + + +class JsonDocument: + """Claude Code's settings.json: sections are top-level object keys, the root is the file.""" + + def __init__(self, root: dict[str, JsonValue]) -> None: # mutable-ok: the document is edited in place + self._root = root + + @classmethod + def parse(cls, text: bytes, path: Path) -> "JsonDocument": + if not text.strip(): + return cls({}) # mutable-ok: fresh document + try: + return cls(_JSON_OBJECT.validate_json(text)) + except ValidationError: + raise AgentConfigError( + f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely." + ) + + def _container(self, section: str) -> dict[str, JsonValue]: # mutable-ok: the document is edited in place + if section == ROOT_SECTION: + return self._root + current: Final = self._root.get(section) + if isinstance(current, dict): + return current + fresh: Final[dict[str, JsonValue]] = {} # mutable-ok: new section object inserted into the document + self._root[section] = fresh + return fresh + + def section(self, name: str) -> Mapping[str, JsonValue] | None: + value: Final = self._root if name == ROOT_SECTION else self._root.get(name) + return value if isinstance(value, dict) else None + + def section_is_object(self, name: str) -> bool: + return name == ROOT_SECTION or isinstance(self._root.get(name), dict) + + def section_is_scalar(self, name: str) -> bool: + value: Final = self._root.get(name) + return name in self._root and value is not None and not isinstance(value, dict) + + def set_value(self, section: str, key: str, value: object) -> None: + self._container(section)[key] = _json_value(value) + + def delete(self, section: str, key: str) -> None: + self._container(section).pop(key, None) + + def null_section(self, name: str) -> None: + self._root[name] = None + + def drop_section(self, name: str) -> None: + self._root.pop(name, None) + + def is_empty(self) -> bool: + return not self._root + + def dumps(self) -> str: + return json.dumps(self._root, indent=2) + + def root(self) -> Mapping[str, JsonValue]: + return self._root + + +class TomlDocument: + """Codex's config.toml, kept as a tomlkit document so comments and layout survive the rewrite.""" + + def __init__(self, doc: tomlkit.TOMLDocument) -> None: + self._doc = doc + + @classmethod + def parse(cls, text: bytes, path: Path) -> "TomlDocument": + try: + return cls(tomlkit.parse(text.decode("utf-8"))) + except (UnicodeDecodeError, ParseError) as e: + raise AgentConfigError(f"{path} is not valid TOML ({e}); cannot proceed safely.") from e + + def _mapping(self, name: str) -> Table | InlineTable | OutOfOrderTableProxy | None: + """The section as tomlkit holds it: a standard table, an inline table, dotted keys, or fragments + split by other tables all count, since every one of them is an object to TOML.""" + current: Final = self._doc.get(name) + return current if isinstance(current, _TOML_MAPPINGS) else None + + def _container(self, section: str) -> tomlkit.TOMLDocument | Table | InlineTable | OutOfOrderTableProxy: + if section == ROOT_SECTION: + return self._doc + current: Final = self._mapping(section) + if current is not None: + return current + fresh: Final = tomlkit.table() + self._doc[section] = fresh + return fresh + + def section(self, name: str) -> Mapping[str, JsonValue] | None: + source: Final = self._doc if name == ROOT_SECTION else self._mapping(name) + if source is None: + return None + return MappingProxyType({str(key): _plain(item) for key, item in source.items()}) + + def section_is_object(self, name: str) -> bool: + return name == ROOT_SECTION or self._mapping(name) is not None + + def section_is_scalar(self, name: str) -> bool: + return name in self._doc and self._mapping(name) is None + + def set_value(self, section: str, key: str, value: object) -> None: + container: Final = self._container(section) + if section == ROOT_SECTION or not isinstance(value, Mapping): + container[key] = _toml_item(value) + return + self._set_child_table(section, key, value) + + def _set_child_table(self, section: str, key: str, value: Mapping[str, object]) -> None: + container: Final = self._container(section) + """Add a child table so the file still means what it says after tomlkit renders it. + + A `[section.key]` header is the readable form, but under a parent spelled as dotted root keys + tomlkit emits that header before the remaining root keys, which TOML then reads as the child's + members, and inside an inline parent it cannot render at all. The rendered document is checked + against the intended meaning and the child falls back to an inline table when the header form + would change it. + """ + expected: Final = _json_value(value) + container[key] = _toml_item(value) + if self._child_means(section, key, expected): + return + del container[key] + container[key] = _toml_item(value, inline=True) + + def _child_means(self, section: str, key: str, expected: JsonValue) -> bool: + try: + rendered: Final = tomlkit.parse(tomlkit.dumps(self._doc)).unwrap() + except ParseError: + return False + parent: Final = rendered.get(section) + return ( + isinstance(parent, dict) + and _json_value(parent.get(key)) == expected + and _json_value(rendered) == _json_value(self._doc.unwrap()) + ) + + def delete(self, section: str, key: str) -> None: + container: Final = self._container(section) + if key in container: + del container[key] + + def null_section(self, name: str) -> None: + self.drop_section(name) + + def drop_section(self, name: str) -> None: + if name in self._doc: + del self._doc[name] + + def is_empty(self) -> bool: + return not self._doc.unwrap() + + def dumps(self) -> str: + return tomlkit.dumps(self._doc) + + +_MAX_CONFIG_DEPTH: Final = 32 +_JsonContainer: TypeAlias = dict[str, JsonValue] | list[JsonValue] # mutable-ok: a JSON document under construction +_JsonWorklist: TypeAlias = list[tuple[object, _JsonContainer, int]] # mutable-ok: breadth-first worklist +_TomlTable: TypeAlias = Table | InlineTable +_TomlWorklist: TypeAlias = list[ + tuple[Mapping[object, object], _TomlTable, bool, int] +] # mutable-ok: breadth-first worklist + + +def _scalar_json(value: object) -> JsonValue: + if isinstance(value, (date, datetime, time)): + return value.isoformat() + return value # pyright: ignore[reportReturnType] # scalars pass through as the JSON value they already are + + +def _too_deep(depth: int) -> None: + if depth > _MAX_CONFIG_DEPTH: + raise AgentConfigError(f"The config nests more than {_MAX_CONFIG_DEPTH} levels deep; cannot proceed safely.") + + +def _empty_json_like(value: object) -> _JsonContainer: + return {} if isinstance(value, Mapping) else [] # mutable-ok: JSON document under construction + + +def _json_value(value: object) -> JsonValue: + """The JSON shape of a config value, built breadth-first so no config depth can recurse. + + Mappings become dicts, sequences become lists, dates become their ISO text; anything deeper + than a config file could plausibly nest is an error rather than a stack. + """ + if not isinstance(value, (Mapping, list, tuple)): + return _scalar_json(value) + root: Final = _empty_json_like(value) + pending: Final[_JsonWorklist] = [(value, root, 0)] # mutable-ok: worklist + while pending: + source, target, depth = pending.pop() + _too_deep(depth) + entries = source.items() if isinstance(source, Mapping) else enumerate(source) # pyright: ignore[reportUnknownVariableType, reportAttributeAccessIssue] # narrowed by the guards above + for key, item in entries: + nested = isinstance(item, (Mapping, list, tuple)) + converted = _empty_json_like(item) if nested else _scalar_json(item) + if nested: + pending.append((item, converted, depth + 1)) # pyright: ignore[reportArgumentType] # nested is a fresh container + if isinstance(target, dict): + target[str(key)] = converted + else: + target.append(converted) + return root + + +def _toml_item(value: object, inline: bool = False) -> object: + """A tomlkit item for a config value, built breadth-first so no config depth can recurse. + + Mappings become tables (inline ones when `inline`, and everything under an inline table is + inline too), tuples become arrays; scalars pass through for tomlkit to wrap. + """ + if isinstance(value, tuple): + return list(value) # mutable-ok: tomlkit builds its Array from a list + if not isinstance(value, Mapping): + return value + root: Final[_TomlTable] = tomlkit.inline_table() if inline else tomlkit.table() + pending: Final[_TomlWorklist] = [(value, root, inline, 0)] # mutable-ok: worklist + while pending: + source, target, as_inline, depth = pending.pop() + _too_deep(depth) + for key, item in source.items(): + if isinstance(item, Mapping): + child = tomlkit.inline_table() if as_inline else tomlkit.table() + target[str(key)] = child + pending.append((item, child, as_inline, depth + 1)) + elif isinstance(item, tuple): + target[str(key)] = list(item) # mutable-ok: tomlkit builds its Array from a list + else: + target[str(key)] = item + return root + + +def refuse_while_owned(path: Path, owners: Sequence[SettingsFileOwner]) -> None: + for owner in owners: + if owner.backup_path.exists(): + raise AgentConfigError( + f"`{owner.start_command}` is currently managing {path} (backup at " + f"{owner.backup_path}) and will restore it when it stops. " + f"Run `{owner.stop_command}` first, then retry." + ) + + +def read_bytes_or_empty(path: Path) -> bytes: + try: + return path.read_bytes() if path.exists() else b"" + except OSError as e: + raise AgentConfigError(f"Could not read {path}: {e}") from e + + +def write_target(path: Path) -> Path: + """Write through a symlinked config file rather than replacing the link. + + os.replace() would swap the symlink itself for a regular file, silently detaching a config + that is symlinked into a dotfiles repo, and there is no backup to undo that. + """ + return path.resolve() if path.is_symlink() else path + + +def _stage(path: Path, text: str) -> str: + try: + return stage_private_text(str(path), text) + except OSError as e: + raise AgentConfigError(f"Could not write {path}: {e}") from e + + +def _owned(container: Mapping[str, JsonValue] | None, key: str) -> OwnedValue: + if container is None: + return OwnedValue(present=False) + return OwnedValue(present=key in container, value=container.get(key)) + + +def fingerprint(owned: OwnedValue) -> str: + return hashlib.sha256(json.dumps(owned.model_dump(mode="json"), sort_keys=True).encode()).hexdigest() + + +def _section_receipt( + before: ConfigDocument, after: ConfigDocument, name: str, keys: Sequence[str], earlier: SectionReceipt | None +) -> SectionReceipt: + before_section: Final = before.section(name) + if name != ROOT_SECTION and before.section_is_scalar(name): + raise AgentConfigError( + f'The config has a non-object "{name}" value, which this would discard. Fix or remove it, then retry.' + ) + after_section: Final = after.section(name) + if earlier is None: + return SectionReceipt( + present=_raw_present(before, name), + was_object=before.section_is_object(name), + previous=MappingProxyType({key: _owned(before_section, key) for key in keys}), + written=MappingProxyType({key: fingerprint(_owned(after_section, key)) for key in keys}), + ) + return SectionReceipt( + present=earlier.present, + was_object=earlier.was_object, + previous=MappingProxyType({key: earlier.previous.get(key, _owned(before_section, key)) for key in keys}), + written=MappingProxyType( + {key: _written_fingerprint(before_section, after_section, key, earlier) for key in keys} + ), + ) + + +def _written_fingerprint( + before: Mapping[str, JsonValue] | None, after: Mapping[str, JsonValue] | None, key: str, earlier: SectionReceipt +) -> str: + """What the repeat configure counts as its own for one key. + + A key the merge changed is ours at its new value. A key the merge left alone keeps the earlier + fingerprint: if the user edited it since the first configure, that fingerprint no longer matches + and unconfigure will report it kept rather than deleting the user's edit as if it were ours. + """ + after_owned: Final = _owned(after, key) + if fingerprint(_owned(before, key)) != fingerprint(after_owned) or key not in earlier.written: + return fingerprint(after_owned) + return earlier.written[key] + + +def _raw_present(document: ConfigDocument, name: str) -> bool: + root: Final = document.section(ROOT_SECTION) + return name == ROOT_SECTION or (root is not None and name in root) + + +def read_configure_receipt(state_path: Path) -> ConfigureReceipt | None: + if not state_path.exists(): + return None + try: + return ConfigureReceipt.model_validate_json(state_path.read_bytes()) + except (OSError, ValidationError) as e: + raise AgentConfigError( + f"{state_path} is not a readable configure receipt ({e}). " + "Remove it and edit the agent's config by hand if it still points at the proxy." + ) from e + + +def release_keys(document: ConfigDocument, receipt: ConfigureReceipt, keys: Sequence[tuple[str, str]]) -> None: + """Put back the pre-configure value of each key that still holds what configure wrote.""" + for section, key in keys: + recorded = receipt.sections.get(section) + if recorded is None or key not in recorded.written: + continue + if fingerprint(_owned(document.section(section), key)) != recorded.written[key]: + continue + previous = recorded.previous.get(key, OwnedValue(present=False)) + if previous.present: + document.set_value(section, key, previous.value) + else: + document.delete(section, key) + + +def configure_document( + path: Path, + state_path: Path, + owners: Sequence[SettingsFileOwner], + parse: Callable[[bytes, Path], ConfigDocument], + owned: Mapping[str, Sequence[str]], + merge: Callable[[ConfigDocument], None], + release: Sequence[tuple[str, str]] = (), + commit: Callable[[str, str], None] = commit_staged_json, +) -> None: + """Persistently rewrite an agent's config, recording how to undo it. + + Both files are staged before either is committed, so a full disk or a read-only directory + fails before anything changes. The two commits are still two renames, so if the config + rename fails after the receipt landed, the earlier receipt is put back (or the new one + removed on a first configure): the receipt on disk never describes a config that was not + written. A repeat configure keeps the receipt's original `previous` snapshot and only + refreshes what was written, so unconfigure still returns to the pre-configure state. + `release` names owned keys whose earlier pin should be let go of before merging. + """ + refuse_while_owned(path, owners) + text: Final = read_bytes_or_empty(path) + before: Final = parse(text, path) + after: Final = parse(text, path) + earlier: Final = read_configure_receipt(state_path) + if earlier is not None: + release_keys(after, earlier, release) + merge(after) + receipt: Final = ConfigureReceipt( + file_existed=earlier.file_existed if earlier is not None else path.exists(), + sections=MappingProxyType( + { + name: _section_receipt(before, after, name, keys, earlier.sections.get(name) if earlier else None) + for name, keys in owned.items() + } + ), + ) + target: Final = write_target(path) + try: + ensure_private_dir(state_path.parent) + except OSError as e: + raise AgentConfigError(f"Could not write {state_path}: {e}") from e + staged_receipt: Final = _stage(state_path, json.dumps(receipt.model_dump(mode="json"), indent=2)) + try: + staged_config: Final = _stage(target, after.dumps()) + except AgentConfigError: + discard_staged_json(staged_receipt) + raise + try: + commit(staged_receipt, str(state_path)) + except OSError as e: + discard_staged_json(staged_config) + raise AgentConfigError(f"Could not write {state_path}: {e}") from e + try: + commit(staged_config, str(target)) + except OSError as config_error: + try: + _restore_receipt(state_path, earlier, commit) + except (AgentConfigError, OSError) as receipt_error: + raise AgentConfigError( + f"Could not write {target}: {config_error}. The receipt at {state_path} now describes a config " + f"that was not written and could not be put back either ({receipt_error}); remove it before retrying." + ) from config_error + raise AgentConfigError(f"Could not write {target}: {config_error}") from config_error + + +def _restore_receipt(state_path: Path, earlier: ConfigureReceipt | None, commit: Callable[[str, str], None]) -> None: + if earlier is None: + state_path.unlink(missing_ok=True) + return + commit(_stage(state_path, json.dumps(earlier.model_dump(mode="json"), indent=2)), str(state_path)) + + +def _still_ours(document: ConfigDocument, receipt: ConfigureReceipt, section: str, key: str) -> bool: + recorded: Final = receipt.sections.get(section) + return recorded is not None and fingerprint(_owned(document.section(section), key)) == recorded.written.get(key) + + +def unconfigure_document( + path: Path, + state_path: Path, + owners: Sequence[SettingsFileOwner], + parse: Callable[[bytes, Path], ConfigDocument], + label: str, + groups: Sequence[RestoreGroup] = (), +) -> UnconfigureOutcome: + """Undo configure, restoring only the keys the user has not changed since.""" + refuse_while_owned(path, owners) + receipt: Final = read_configure_receipt(state_path) + if receipt is None: + raise AgentConfigError( + f"{label} is not configured by `lite configure` (no receipt at {state_path}); nothing to undo." + ) + document: Final = parse(read_bytes_or_empty(path), path) + for name in receipt.sections: + if name != ROOT_SECTION and document.section_is_scalar(name): + raise AgentConfigError( + f'The config has a non-object "{name}" value, which this would discard. Fix or remove it, then retry.' + ) + withheld_keys: Final = frozenset( + dependent + for group in groups + if not _still_ours(document, receipt, *group.anchor) + for dependent in group.dependents + ) + restored: Final[list[str]] = [] # mutable-ok: report accumulator, frozen into the outcome below + kept: Final[list[str]] = [] # mutable-ok: report accumulator, frozen into the outcome below + withheld: Final[list[str]] = [] # mutable-ok: report accumulator, frozen into the outcome below + for name, section in receipt.sections.items(): + for key in section.written: + label_key = key if name == ROOT_SECTION else f"{name}.{key}" + previous = section.previous.get(key, OwnedValue(present=False)) + if section.written[key] == fingerprint(previous): + continue + if fingerprint(_owned(document.section(name), key)) != section.written[key]: + kept.append(label_key) + continue + if (name, key) in withheld_keys and previous.present: + document.delete(name, key) + withheld.append(label_key) + continue + if previous.present: + document.set_value(name, key, previous.value) + else: + document.delete(name, key) + restored.append(label_key) + if name != ROOT_SECTION and not section.was_object and not document.section(name): + if section.present: + document.null_section(name) + else: + document.drop_section(name) + target: Final = write_target(path) + file_removed: Final = document.is_empty() and not receipt.file_existed + try: + if file_removed: + target.unlink(missing_ok=True) + else: + commit_staged_json(_stage(target, document.dumps()), str(target)) + state_path.unlink(missing_ok=True) + except OSError as e: + raise AgentConfigError(f"Could not write {target}: {e}") from e + return UnconfigureOutcome( + restored=tuple(restored), kept=tuple(kept), withheld=tuple(withheld), file_removed=file_removed + ) + + +__all__ = ( + "ROOT_SECTION", + "AgentConfigError", + "ConfigDocument", + "ConfigureReceipt", + "JsonDocument", + "OwnedValue", + "RestoreGroup", + "SectionReceipt", + "SettingsFileOwner", + "TomlDocument", + "UnconfigureOutcome", + "configure_document", + "fingerprint", + "read_bytes_or_empty", + "read_configure_receipt", + "refuse_while_owned", + "release_keys", + "unconfigure_document", + "write_target", +) diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index ea1eed65505..a04c54abba4 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -404,7 +404,7 @@ _WINDOWS_SHIM_SUFFIXES: Final[frozenset[str]] = frozenset({".cmd", ".bat"}) _CMD_LINE_BREAKS: Final = ("\r", "\n") -def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]: +def windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]: """Build what CreateProcess runs, routing batch shims through cmd.exe. npm installs Claude Code as `claude.cmd`, which PATHEXT lets shutil.which @@ -425,7 +425,7 @@ def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]: rest: Final = tuple(args[1:]) if os.path.splitext(path)[1].lower() not in _WINDOWS_SHIM_SUFFIXES: return (path, *rest) - if any(brk in token for token in rest for brk in _CMD_LINE_BREAKS): + if any(brk in token for token in (path, *rest) for brk in _CMD_LINE_BREAKS): raise AgentRunError( f"Cannot pass an argument containing a line break to `{os.path.basename(path)}` on " "Windows: cmd.exe ends the command line there, so the agent would silently lose it." @@ -465,7 +465,7 @@ def _hand_off( child and exits with its status. """ if platform.startswith("win"): - raise SystemExit(spawn(_windows_command(path, args), env)) + raise SystemExit(spawn(windows_command(path, args), env)) replace(path, list(args), dict(env)) @@ -489,6 +489,10 @@ def _restore_controlling_terminal() -> None: os.close(fd) +hand_off = _hand_off +restore_controlling_terminal = _restore_controlling_terminal + + def run_agent( base_url: str, api_key: str, @@ -502,7 +506,7 @@ def run_agent( agent_model_sync_env ), warn: Callable[[str], None] = _warn, - launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off, + launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = hand_off, reattach_terminal: Callable[[], None] | None = None, preparers: Mapping[str, _Preparer] = MappingProxyType(_PREPARERS), ) -> None: @@ -546,10 +550,14 @@ def run_agent( launcher(binary, [command[0], *extra_args, *command[1:]], env) -def _is_interactive() -> bool: +def is_interactive() -> bool: return sys.stdin.isatty() +def _is_interactive() -> bool: + return is_interactive() + + def resolve_api_key(ctx: click.Context) -> str: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"] @@ -574,15 +582,17 @@ def resolve_api_key(ctx: click.Context) -> str: _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy." -def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: - ctx_obj: Final[CliContextObj] = ctx.obj - base_url: Final = ctx_obj["base_url"] - started_interactive: Final = _is_interactive() - api_key: Final = resolve_api_key(ctx) - +def launch_agent( + base_url: str, + api_key: str, + binary: str, + args: Sequence[str] = (), + *, + skip_verify: bool = False, + started_interactive: bool, +) -> None: display_name, _profiles = agent_profile(binary) click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}") - try: run_agent( base_url, @@ -595,6 +605,19 @@ def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify raise click.ClickException(str(e)) +def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: + ctx_obj: Final[CliContextObj] = ctx.obj + started_interactive: Final = _is_interactive() + launch_agent( + ctx_obj["base_url"], + resolve_api_key(ctx), + binary, + args, + skip_verify=skip_verify, + started_interactive=started_interactive, + ) + + def _make_agent_command(binary: str, display_name: str) -> click.Command: @click.command( name=binary, @@ -631,10 +654,15 @@ __all__ = [ "agent_model_sync_env", "agent_profile", "build_agent_env", + "hand_off", + "is_interactive", + "launch_agent", "opencode_model_sync_env", "opencode_provider_config", "prepare_pi", "resolve_api_key", + "restore_controlling_terminal", "run_agent", "verify_proxy_key", + "windows_command", ] diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index e6231f3cac9..dc348aa15e5 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -12,7 +12,9 @@ owned like any other key and stripped. import hashlib import json +import os import shlex +import shutil import sys from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass @@ -69,8 +71,15 @@ _BASE_URL_PATH: Final = f"{ENV_KEY}.{ANTHROPIC_BASE_URL_KEY}" _MODEL_PATHS: Final = (MODEL_KEY, f"{ENV_KEY}.{ANTHROPIC_MODEL_KEY}") STARTING_MODEL_ROLE: Final = "the /model picker's default row, the model Claude Code starts and resumes on" -CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" CLAUDE_CONFIG_DIR_ENV: Final = "CLAUDE_CONFIG_DIR" + + +def claude_config_dir() -> Path: + override: Final = os.environ.get(CLAUDE_CONFIG_DIR_ENV) + return Path(override).expanduser() if override else Path.home() / ".claude" + + +CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json" CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json" @@ -132,6 +141,16 @@ class StaticToken: token: str +@dataclass(frozen=True, slots=True) +class ApiKeyHelper: + """A `lite auth print-token` command Claude Code runs per request.""" + + command: str + + +ClaudeCredential: TypeAlias = StaticToken | ApiKeyHelper + + @dataclass(frozen=True, slots=True) class KeepModel: """Leave the top-level `model` as it is, the user's or an earlier configure's (a re-login).""" @@ -329,7 +348,7 @@ def with_status_line(settings: Mapping[str, JsonValue], command: str) -> Mapping def merge_claude_settings( settings: Mapping[str, JsonValue], base_url: str, - credential: StaticToken, + credential: ClaudeCredential, default_model: str | None = None, tier_model: str | None = None, *, @@ -353,7 +372,8 @@ def merge_claude_settings( (ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE), ), ((key, value) for key, value in current_env.items() if key not in _CREDENTIAL_ENV_KEYS), - ((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")), (ANTHROPIC_AUTH_TOKEN_KEY, credential.token)), + ((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")),), + ((ANTHROPIC_AUTH_TOKEN_KEY, credential.token),) if isinstance(credential, StaticToken) else (), ((ANTHROPIC_MODEL_KEY, default_model),) if default_model is not None else (), ((key, tier_model) for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS if tier_model is not None), ) @@ -366,11 +386,26 @@ def merge_claude_settings( if key not in (API_KEY_HELPER_KEY, ENV_KEY) ), ((ENV_KEY, env),), + ((API_KEY_HELPER_KEY, credential.command),) if isinstance(credential, ApiKeyHelper) else (), ((MODEL_KEY, default_model),) if default_model is not None else (), ) ) +def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str: + quote: Final = quote_for_cmd if platform.startswith("win") else shlex.quote + return " ".join(quote(token) for token in print_token_command(base_url)) + + +def print_token_command(base_url: str) -> tuple[str, ...]: + lite_path: Final = shutil.which("lite") + if lite_path is None: + raise ClaudeSettingsError( + "Could not find `lite` on your PATH. The agent credential command needs an absolute path to it." + ) + return (str(Path(lite_path).resolve()), "--base-url", base_url, "auth", "print-token") + + def _owned(container: Mapping[str, JsonValue], key: str) -> OwnedValue: return OwnedValue(present=key in container, value=container.get(key)) @@ -535,6 +570,10 @@ def _endpoint_text(endpoint: OwnedValue) -> str: return endpoint.value if isinstance(endpoint.value, str) else json.dumps(endpoint.value) +def _withheld_text(item: WithheldCredential) -> str: + return f"{item.key} (captured with {item.endpoint})" + + def unconfigure_claude_settings( settings_path: Path, state_path: Path, owners: Sequence[SettingsFileOwner] ) -> UnconfigureOutcome: @@ -616,6 +655,8 @@ __all__ = ( "STARTING_MODEL_ROLE", "STATUSLINE_SCRIPT_PATH", "STATUS_LINE_KEY", + "ApiKeyHelper", + "ClaudeCredential", "ClaudeSettingsError", "ConfigureReceipt", "KeepModel", @@ -627,13 +668,16 @@ __all__ = ( "UnconfigureOutcome", "UnpinModel", "WithheldCredential", + "claude_config_dir", "claude_settings_path", "configure_claude_settings", "configure_state_path", "load_json_or_empty", "merge_claude_settings", + "print_token_command", "read_configure_receipt", "refuse_while_owned", + "resolve_api_key_helper", "settings_file_owners", "unconfigure_claude_settings", "write_claude_settings", diff --git a/litellm/proxy/client/cli/commands/codex_settings.py b/litellm/proxy/client/cli/commands/codex_settings.py new file mode 100644 index 00000000000..3aa53eb856e --- /dev/null +++ b/litellm/proxy/client/cli/commands/codex_settings.py @@ -0,0 +1,168 @@ +"""Codex's config.toml: what `lite configure codex` writes there and how it undoes it. + +Codex ignores OPENAI_BASE_URL and routes through a named provider, so the wiring is a +`[model_providers.litellm]` table plus `model_provider` at the root. A Codex launched from the +desktop app never sees the shell environment, which is why the credential is written into the +table (a static key) or read through a command Codex runs itself (the `lite login` credential), +never through an env var. +""" + +import hashlib +import os +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path +from types import MappingProxyType +from typing import Final + +from litellm.litellm_core_utils.private_json import commit_staged_json + +from .agent_config import ( + ROOT_SECTION, + ConfigDocument, + TomlDocument, + UnconfigureOutcome, + configure_document, + unconfigure_document, +) +from .claude_settings import ClaudeCredential, KeepModel, ModelChoice, StartOn, StaticToken + +CODEX_HOME_ENV: Final = "CODEX_HOME" +CODEX_PROVIDER_ID: Final = "litellm" +MODEL_PROVIDERS_KEY: Final = "model_providers" +MODEL_PROVIDER_KEY: Final = "model_provider" +MODEL_KEY: Final = "model" +MODEL_CONTEXT_WINDOW_KEY: Final = "model_context_window" +OWNED_ROOT_KEYS: Final = (MODEL_PROVIDER_KEY, MODEL_KEY, MODEL_CONTEXT_WINDOW_KEY) +OWNED_SECTIONS: Final[Mapping[str, Sequence[str]]] = MappingProxyType( + {ROOT_SECTION: OWNED_ROOT_KEYS, MODEL_PROVIDERS_KEY: (CODEX_PROVIDER_ID,)} +) +PINNED_KEYS: Final = ((ROOT_SECTION, MODEL_KEY), (ROOT_SECTION, MODEL_CONTEXT_WINDOW_KEY)) + + +def codex_home() -> Path: + """Where Codex keeps its config: ~/.codex unless CODEX_HOME relocates it.""" + override: Final = os.environ.get(CODEX_HOME_ENV) + return Path(override).expanduser() if override else Path.home() / ".codex" + + +def codex_config_path(environ: Mapping[str, str]) -> Path: + home: Final = environ.get(CODEX_HOME_ENV) + return Path(home).expanduser() / "config.toml" if home else Path.home() / ".codex" / "config.toml" + + +def codex_configure_state_path(config_path: Path) -> Path: + default: Final = Path.home() / ".codex" / "config.toml" + state_root: Final = Path.home() / ".litellm" / "codex_configure_state" + default_state: Final = state_root.parent / "codex_configure_state.json" + if config_path.resolve() == default.resolve(): + return default_state + digest: Final = hashlib.sha256(str(config_path.resolve()).encode()).hexdigest() + return state_root / f"{digest}.json" + + +CODEX_CONFIG_PATH: Final = codex_home() / "config.toml" +CODEX_CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "codex_configure_state.json" + + +def codex_provider( + base_url: str, credential: ClaudeCredential, print_token: Callable[[], Sequence[str]] +) -> Mapping[str, object]: + """The `[model_providers.litellm]` table: HTTP/SSE Responses transport, credential inline or by command. + + supports_websockets is off because the proxy does not speak the Responses WebSocket protocol. The + print-token argv is resolved only for the login credential, so a static key never needs `lite` + on PATH. + """ + auth: Final = ( + (("experimental_bearer_token", credential.token),) + if isinstance(credential, StaticToken) + else (("auth", _auth_command(print_token())),) + ) + return MappingProxyType( + dict( + ( + ("name", "LiteLLM proxy"), + ("base_url", base_url.rstrip("/") + "/v1"), + ("wire_api", "responses"), + ("supports_websockets", False), + *auth, + ) + ) + ) + + +def _auth_command(argv: Sequence[str]) -> Mapping[str, object]: + return MappingProxyType({"command": argv[0], "args": tuple(argv[1:]), "timeout_ms": 5000}) + + +def apply_codex_merge( + document: ConfigDocument, + base_url: str, + credential: ClaudeCredential, + print_token: Callable[[], Sequence[str]], + model: str | None, + context_window: int | None, +) -> None: + """Point Codex at the proxy in place; `model` and `context_window` are written only when given. + + Codex sizes its context from its own catalog, which knows nothing about proxy models, so the + proxy's limit for the chosen model is written as model_context_window alongside the pin. A pin + with no known window leaves the key to whatever release put there, never a stale one. + """ + document.set_value(ROOT_SECTION, MODEL_PROVIDER_KEY, CODEX_PROVIDER_ID) + document.set_value(MODEL_PROVIDERS_KEY, CODEX_PROVIDER_ID, codex_provider(base_url, credential, print_token)) + if model is not None: + document.set_value(ROOT_SECTION, MODEL_KEY, model) + if context_window is not None: + document.set_value(ROOT_SECTION, MODEL_CONTEXT_WINDOW_KEY, context_window) + + +def configure_codex_config( + base_url: str, + credential: ClaudeCredential, + print_token: Callable[[], Sequence[str]], + model: ModelChoice, + context_window: int | None, + config_path: Path, + state_path: Path, + commit: Callable[[str, str], None] = commit_staged_json, +) -> None: + """Persistently route Codex through base_url, recording how to undo it (see configure_document). + + Both a re-pin and an unpin first let go of the model and window an earlier configure wrote, + so a re-pin to a model whose window the proxy does not report cannot keep the old model's. + """ + pinned: Final = model.model if isinstance(model, StartOn) else None + configure_document( + config_path, + state_path, + (), + TomlDocument.parse, + OWNED_SECTIONS, + lambda document: apply_codex_merge(document, base_url, credential, print_token, pinned, context_window), + release=() if isinstance(model, KeepModel) else PINNED_KEYS, + commit=commit, + ) + + +def unconfigure_codex_config(config_path: Path, state_path: Path) -> UnconfigureOutcome: + """Undo `lite configure codex`, restoring only the keys the user has not changed since.""" + return unconfigure_document(config_path, state_path, (), TomlDocument.parse, "Codex") + + +__all__ = ( + "CODEX_CONFIGURE_STATE_PATH", + "CODEX_CONFIG_PATH", + "CODEX_HOME_ENV", + "CODEX_PROVIDER_ID", + "OWNED_ROOT_KEYS", + "OWNED_SECTIONS", + "PINNED_KEYS", + "apply_codex_merge", + "codex_config_path", + "codex_configure_state_path", + "codex_home", + "codex_provider", + "configure_codex_config", + "unconfigure_codex_config", +) diff --git a/litellm/proxy/client/cli/commands/configure.py b/litellm/proxy/client/cli/commands/configure.py index 4acf94e16f9..5c5573aa39c 100644 --- a/litellm/proxy/client/cli/commands/configure.py +++ b/litellm/proxy/client/cli/commands/configure.py @@ -1,8 +1,8 @@ -"""`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable.""" +"""`lite configure [claude|codex]` and `lite unconfigure [claude|codex]`: persistent agent wiring, undoable.""" import os import sys -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from pathlib import Path from types import MappingProxyType @@ -18,9 +18,12 @@ from litellm.proxy.common_utils.model_listing_utils import ( GATEWAY_CLIENT_HEADER, ) -from .auth import CliContextObj +from .agent_config import AgentConfigError, fingerprint, read_configure_receipt +from .agents import is_interactive +from .auth import CliContextObj, context_secret_vault, get_stored_api_key from .claude_settings import ( - STARTING_MODEL_ROLE, + ApiKeyHelper, + ClaudeCredential, ClaudeSettingsError, ModelChoice, StartOn, @@ -30,35 +33,73 @@ from .claude_settings import ( claude_settings_path, configure_claude_settings, configure_state_path, + print_token_command, refuse_while_owned, + resolve_api_key_helper, settings_file_owners, unconfigure_claude_settings, ) -from .pi import ListedModel, ListingFailure, PiSyncError, fetch_model_listing +from .codex_settings import ( + codex_config_path, + codex_configure_state_path, + configure_codex_config, + unconfigure_codex_config, +) +from .configured_launch import launch_configured_agents +from .pi import ListedModel, ListingFailure, PiSyncError, fetch_model_ids, fetch_model_limits, fetch_model_listing +from .up import UpError, ensure_fresh_login _LISTED_MODELS_SHOWN: Final = 20 _CLAUDE_TARGET: Final = "claude" -_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"),) -_KEEP_DEFAULT_MODEL: Final = "Keep Claude Code's own default" +_CODEX_TARGET: Final = "codex" +_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"), (_CODEX_TARGET, "Codex (CLI and app)")) +_TARGET_LABELS: Final[Mapping[str, str]] = MappingProxyType(dict(_TARGETS)) +_KEEP_DEFAULT_MODEL: Final = "Keep the agent's own default" _CLAUDE_CODE_VIEW: Final = MappingProxyType( {"anthropic-version": "2023-06-01", GATEWAY_CLIENT_HEADER: CLAUDE_CODE_CLIENT} ) _MODEL_OPTION_HELP: Final = ( - f"Proxy model to set as {STARTING_MODEL_ROLE}. Must be listed on /v1/models for the key; without it, " + "Proxy model to set as Claude Code's starting model. Must be listed on /v1/models for the key; without it, " "Claude Code keeps its own default and a pin an earlier configure made is let go of. Nothing pins Claude " "Code's sub-agent or background tiers; `lite autoroute up` is the mode that does." ) +_API_KEY_HELP: Final = ( + "Long-lived LiteLLM virtual key written into the agent's config. Defaults to the `lite --api-key` / " + "LITELLM_PROXY_API_KEY value. Codex can instead use your `lite login` credential through `lite auth print-token`." +) +_LAUNCH_HELP: Final = "Start the agent right after configuring it." -def resolve_credential(ctx: click.Context, api_key: str | None) -> StaticToken: - """The long-lived key written into settings.json: --api-key, `lite --api-key` or LITELLM_PROXY_API_KEY. +@dataclass(frozen=True, slots=True) +class _ClaudeListing: + models: tuple[ListedModel, ...] - A `lite login` credential is never written: it expires within a day, and keeping it fresh would mean - Claude Code running `lite` through `apiKeyHelper` on every credential refresh. - """ + @property + def ids(self) -> tuple[str, ...]: + return tuple(model.id for model in self.models) + + +@dataclass(frozen=True, slots=True) +class _Session: + base_url: str + credential: ClaudeCredential + key: str + started_interactive: bool + + def launch(self, agents: Sequence[str]) -> None: + launch_configured_agents(agents, started_interactive=self.started_interactive) + + +def _explicit_key(ctx: click.Context, api_key: str | None) -> str | None: ctx_obj: Final[CliContextObj] = ctx.obj - explicit: Final = api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key")) - if not explicit: + return api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key")) + + +def _claude_credential(ctx: click.Context, api_key: str | None) -> StaticToken: + settings_path: Final = claude_settings_path(os.environ) + refuse_while_owned(settings_path, settings_file_owners(settings_path)) + explicit: Final = _explicit_key(ctx, api_key) + if explicit is None: raise ClaudeSettingsError( "`lite configure claude` needs a long-lived virtual key: pass --api-key, `lite --api-key`, or set " "LITELLM_PROXY_API_KEY. Your `lite login` credential expires within a day, so it is not written " @@ -67,96 +108,106 @@ def resolve_credential(ctx: click.Context, api_key: str | None) -> StaticToken: return StaticToken(explicit) -@dataclass(frozen=True, slots=True) -class _Listing: - models: tuple[ListedModel, ...] - - @property - def ids(self) -> tuple[str, ...]: - return tuple(model.id for model in self.models) - - -def _start(ctx: click.Context, api_key: str | None) -> tuple[StaticToken, _Listing]: - """Every configure path begins the same way: the local ownership check first, so a `lite up` - session is refused before any request, then the credential, then the listing.""" - settings_path: Final = claude_settings_path(os.environ) +def _codex_credential(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, str]: + explicit: Final = _explicit_key(ctx, api_key) + if explicit is not None: + return StaticToken(explicit), explicit try: - refuse_while_owned(settings_path, settings_file_owners(settings_path)) - credential: Final = resolve_credential(ctx, api_key) - except ClaudeSettingsError as e: - raise click.ClickException(str(e)) - return credential, _listed_models(ctx.obj["base_url"], credential.token) + ensure_fresh_login(ctx, reader="Codex's provider auth command reads this token on every request") + except UpError as e: + raise ClaudeSettingsError(str(e)) from e + base_url: Final = ctx.obj["base_url"] + stored: Final = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx)) + if stored is None: + raise ClaudeSettingsError("Login did not produce a usable token.") + return ApiKeyHelper(resolve_api_key_helper(base_url)), stored -def _listing_error(base_url: str, error: PiSyncError) -> str: - """The hint that fits how the listing failed: only an unreachable proxy gets the "is it running" question.""" +def _listing_error(base_url: str, error: PiSyncError, agent: str) -> str: if error.kind is ListingFailure.REJECTED: return f"LiteLLM rejected your key (HTTP {error.status}). Pass a valid --api-key." if error.kind is ListingFailure.UNREACHABLE: return f"{error.message} Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?" if error.kind is ListingFailure.EMPTY: - return f"{error.message} Claude Code would have nothing to run; give the key access to at least one model." + return f"{error.message} {agent} would have nothing to run; give the key access to at least one model." return f"{error.message} The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy." -def _listed_models(base_url: str, key: str) -> _Listing: +def _claude_listing(base_url: str, key: str) -> _ClaudeListing: listed: Final = fetch_model_listing(base_url, key, headers=_CLAUDE_CODE_VIEW) if isinstance(listed, PiSyncError): - raise click.ClickException(_listing_error(base_url, listed)) - return _Listing(listed) + raise click.ClickException(_listing_error(base_url, listed, "Claude Code")) + return _ClaudeListing(listed) -def _starting_model(model: str, listing: _Listing) -> str | None: - source: Final = next((listed.id for listed in listing.models if listed.source_model == model), None) - return source or next((listed.id for listed in listing.models if listed.id == model), None) +def _codex_listing(base_url: str, key: str) -> tuple[str, ...]: + listed: Final = fetch_model_ids(base_url, key) + if isinstance(listed, PiSyncError): + raise click.ClickException(_listing_error(base_url, listed, "Codex")) + return listed + + +def _session(ctx: click.Context, api_key: str | None, agent: str) -> _Session: + started_interactive: Final = is_interactive() + try: + if agent == _CLAUDE_TARGET: + credential: Final = _claude_credential(ctx, api_key) + return _Session(ctx.obj["base_url"], credential, credential.token, started_interactive) + codex_credential, key = _codex_credential(ctx, api_key) + return _Session(ctx.obj["base_url"], codex_credential, key, started_interactive) + except (AgentConfigError, ClaudeSettingsError) as e: + raise click.ClickException(str(e)) from e def _model_choice(model: str | None) -> ModelChoice: return StartOn(model) if model is not None else UnpinModel() -def _apply_claude(ctx: click.Context, credential: StaticToken, listing: _Listing, model: str | None) -> None: - ctx_obj: Final[CliContextObj] = ctx.obj - base_url: Final = ctx_obj["base_url"] - listed: Final = listing.ids - starting: Final = _starting_model(model, listing) if model is not None else None - if model is not None and starting is None: - shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN]) - more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else "" - raise click.ClickException( - f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}{more}." - ) +def _require_listed(base_url: str, listed: Sequence[str], model: str | None) -> ModelChoice: + if model is None: + return UnpinModel() + if model in listed: + return StartOn(model) + shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN]) + more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else "" + raise click.ClickException(f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}{more}.") + + +def _apply_claude(session: _Session, listing: _ClaudeListing, model: str | None) -> None: + source: Final = ( + next((item.id for item in listing.models if item.source_model == model), None) if model is not None else None + ) + choice: Final = _model_choice(source or model) + if model is not None and source is None and model not in listing.ids: + _require_listed(session.base_url, listing.ids, model) settings_path: Final = claude_settings_path(os.environ) try: configure_claude_settings( - base_url, - credential, - _model_choice(starting), + session.base_url, + session.credential, + choice, settings_path, configure_state_path(settings_path), settings_file_owners(settings_path), ) except ClaudeSettingsError as e: raise click.ClickException(str(e)) - in_picker: Final = sum(1 for listed_model in listed if CLAUDE_CODE_PICKER_PATTERN.search(listed_model)) - click.echo(f"Configured Claude Code: {settings_path} now routes through {base_url}.") - + in_picker: Final = sum(1 for item in listing.ids if CLAUDE_CODE_PICKER_PATTERN.search(item)) + starting: Final = source or model + click.echo(f"Configured Claude Code: {settings_path} now routes through {session.base_url}.") click.echo("Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN.") click.echo( - f"Starting model: {starting} ({STARTING_MODEL_ROLE}); switch any time with /model." + f"Starting model: {starting} (the /model picker's default row, the model Claude Code starts and resumes on); switch any time with /model." if starting is not None - else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or " - "pass --model to start on a proxy model. Without a pin, a resumed session re-sends the model its transcript " - "recorded, which behind a raw-model auto-router is the tier model." + else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or pass --model to start on a proxy model." ) click.echo( - f"/model will list all {len(listed)} of the proxy's models." - if in_picker == len(listed) - else f"/model will list {in_picker} of the proxy's {len(listed)} models: Claude Code shows only ids containing " - "'claude' or 'anthropic', and this proxy does not list the rest under such names." + f"/model will list all {len(listing.ids)} of the proxy's models." + if in_picker == len(listing.ids) + else f"/model will list {in_picker} of the proxy's {len(listing.ids)} models: Claude Code shows only ids containing 'claude' or 'anthropic', and this proxy does not list the rest under such names." ) click.echo("Start `claude` from any terminal. Undo with `lite unconfigure claude`.") - if settings_path.is_symlink(): + if isinstance(session.credential, StaticToken) and settings_path.is_symlink(): click.echo( f"Note: {settings_path} is a symlink to {settings_path.resolve()}, so your key now lives in " "that file; keep it out of version control.", @@ -164,52 +215,120 @@ def _apply_claude(ctx: click.Context, credential: StaticToken, listing: _Listing ) +def _pin_released(state_path: Path, choice: ModelChoice) -> bool: + if not isinstance(choice, UnpinModel): + return False + earlier: Final = read_configure_receipt(state_path) + root: Final = earlier.sections.get("") if earlier is not None else None + return root is not None and "model" in root.written and root.written["model"] != fingerprint(root.previous["model"]) + + +def _apply_codex(session: _Session, listed: tuple[str, ...], model: str | None) -> None: + choice: Final = _require_listed(session.base_url, listed, model) + config_path: Final = codex_config_path(os.environ) + state_path: Final = codex_configure_state_path(config_path) + try: + released: Final = _pin_released(state_path, choice) + limits: Final = fetch_model_limits(session.base_url, session.key) if model is not None else MappingProxyType({}) + context_window: Final = limits[model].context_window if model is not None and model in limits else None + configure_codex_config( + session.base_url, + session.credential, + lambda: print_token_command(session.base_url), + choice, + context_window, + config_path, + state_path, + ) + except (AgentConfigError, ClaudeSettingsError) as e: + raise click.ClickException(str(e)) from e + click.echo(f"Configured Codex: {config_path} now routes through {session.base_url}.") + click.echo( + "Credential: your virtual key, stored in the file as a bearer token." + if isinstance(session.credential, StaticToken) + else "Credential: your `lite login`, which Codex reads through `lite auth print-token` on demand." + ) + window: Final = f", context window {context_window} tokens from the proxy" if context_window is not None else "" + click.echo( + f"Model: {model}{window}." + if model is not None + else "Model: the pin an earlier configure made is released, back to Codex's own default id." + if released + else "Model: not pinned (Codex's own default id)." + ) + click.echo("Start `codex` from any terminal or the Codex app. Undo with `lite unconfigure codex`.") + + +def _target_choices() -> tuple[Choice, ...]: + return tuple(Choice(value, name=label, enabled=value == _CLAUDE_TARGET) for value, label in _TARGETS) + + def _pick_targets() -> tuple[str, ...]: picked: Final = inquirer.checkbox( message="Which agents should route through LiteLLM?", - choices=[Choice(value, name=label, enabled=True) for value, label in _TARGETS], + choices=_target_choices(), validate=lambda chosen: len(chosen) > 0, invalid_message="Pick at least one.", ).execute() return tuple(str(value) for value in picked) -def _pick_model(listed: Sequence[str]) -> str | None: +def _pick_model(agent: str, listed: Sequence[str]) -> str | None: picked: Final = inquirer.fuzzy( - message="Model Claude Code starts on (type to filter; /model switches any time):", - choices=[_KEEP_DEFAULT_MODEL, *listed], + message=f"Model {_TARGET_LABELS[agent]} starts on (type to filter):", + choices=[_KEEP_DEFAULT_MODEL, *listed], # mutable-ok: InquirerPy takes a list ).execute() return None if picked == _KEEP_DEFAULT_MODEL else str(picked) +def _confirm_launch(agent: str) -> bool: + answer: Final[object] = inquirer.confirm(message=f"Start {_TARGET_LABELS[agent]} now?", default=True).execute() + return answer is True + + def interactive_configure( ctx: click.Context, pick_targets: Callable[[], tuple[str, ...]] = _pick_targets, - pick_model: Callable[[Sequence[str]], str | None] = _pick_model, + pick_model: Callable[[str, Sequence[str]], str | None] = _pick_model, + confirm_launch: Callable[[str], bool] = _confirm_launch, ) -> None: - """`lite configure` with no agent named: ask which agents to wire and which model to pin.""" - targets: Final = pick_targets() - if _CLAUDE_TARGET not in targets: + targets: Final = tuple(target for target in pick_targets() if target in (_CLAUDE_TARGET, _CODEX_TARGET)) + if not targets: return - credential, listing = _start(ctx, None) - _apply_claude( - ctx, credential, listing, pick_model(tuple(model.source_model or model.id for model in listing.models)) - ) + session: Final = _session(ctx, None, _CLAUDE_TARGET if _CLAUDE_TARGET in targets else _CODEX_TARGET) + if _CLAUDE_TARGET in targets: + claude_listing: Final = _claude_listing(session.base_url, session.key) + _apply_claude( + session, + claude_listing, + pick_model(_CLAUDE_TARGET, tuple(model.source_model or model.id for model in claude_listing.models)), + ) + if _CODEX_TARGET in targets: + codex_models: Final = _codex_listing(session.base_url, session.key) + _apply_codex(session, codex_models, pick_model(_CODEX_TARGET, codex_models)) + selected: Final = tuple(target for target in targets if confirm_launch(target)) + if selected: + session.launch(selected) + + +def _configure_context(ctx: click.Context, _param: click.Parameter, api_key: str | None) -> None: + ctx.ensure_object(dict) + if api_key is not None: + ctx.obj["api_key"] = api_key # rebind-ok: Click option callback populates its command context + ctx.obj["api_key_from_token_file"] = False # rebind-ok: Click option callback owns this credential marker @click.group(name="configure", invoke_without_command=True) +@click.option("--api-key", "api_key", default=None, help=_API_KEY_HELP, expose_value=False, callback=_configure_context) @click.pass_context def configure_group(ctx: click.Context) -> None: - """Persistently route a coding agent through your LiteLLM proxy. - - With no agent named, asks which agents to wire and which proxy model to pin. - """ + """Persistently route a coding agent through your LiteLLM proxy.""" if ctx.invoked_subcommand is not None: return if not sys.stdin.isatty(): raise click.ClickException( "`lite configure` asks questions, so it needs a terminal. Non-interactively, run " - "`lite configure claude --api-key --model `." + "`lite configure --api-key claude --model ` or `lite configure --api-key codex ...`." ) interactive_configure(ctx) @@ -220,63 +339,69 @@ def unconfigure_group() -> None: @configure_group.command(name="claude") -@click.option( - "--api-key", - "api_key", - default=None, - help="Long-lived LiteLLM virtual key written into Claude Code's settings. Defaults to the `lite --api-key` / " - "LITELLM_PROXY_API_KEY value; required, since a `lite login` credential expires within a day.", -) +@click.option("--api-key", "api_key", default=None, help=_API_KEY_HELP) @click.option("--model", default=None, help=_MODEL_OPTION_HELP) +@click.option("--launch", is_flag=True, default=False, help=_LAUNCH_HELP) @click.pass_context -def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) -> None: - """Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`. +def configure_claude(ctx: click.Context, api_key: str | None, model: str | None, launch: bool) -> None: + """Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`.""" + session: Final = _session(ctx, api_key, _CLAUDE_TARGET) + _apply_claude(session, _claude_listing(session.base_url, session.key), model) + if launch: + session.launch((_CLAUDE_TARGET,)) - Patches ~/.claude/settings.json in place: the proxy URL, your virtual key as a static token, - and gateway model discovery so /model lists the proxy's models; --model picks the one Claude - Code starts on and resumes with. Every other - setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back. - Assumes the proxy is already running. - """ - credential, listing = _start(ctx, api_key) - _apply_claude(ctx, credential, listing, model) + +@configure_group.command(name="codex") +@click.option("--api-key", "api_key", default=None, help=_API_KEY_HELP) +@click.option("--model", default=None, help="Proxy model Codex uses. Must be listed on /v1/models for the key.") +@click.option("--launch", is_flag=True, default=False, help=_LAUNCH_HELP) +@click.pass_context +def configure_codex(ctx: click.Context, api_key: str | None, model: str | None, launch: bool) -> None: + """Route Codex through your LiteLLM proxy until `lite unconfigure codex`.""" + session: Final = _session(ctx, api_key, _CODEX_TARGET) + _apply_codex(session, _codex_listing(session.base_url, session.key), model) + if launch: + session.launch((_CODEX_TARGET,)) @unconfigure_group.command(name="claude") def unconfigure_claude() -> None: - """Return Claude Code's settings to what they were before `lite configure claude`. - - Also undoes `lite login --config-claude`. Only keys still holding what configure wrote are - put back; anything you changed since is left as it is and named in the output. - """ + """Return Claude Code's settings to what they were before `lite configure claude`.""" settings_path: Final = claude_settings_path(os.environ) state_path: Final = configure_state_path(settings_path) try: outcome: Final = unconfigure_claude_settings(settings_path, state_path, settings_file_owners(settings_path)) except ClaudeSettingsError as e: raise click.ClickException(str(e)) - _report_unconfigure(settings_path, state_path, outcome) + _report_unconfigure(settings_path, outcome, state_path) -def _report_unconfigure(settings_path: Path, state_path: Path, outcome: UnconfigureOutcome) -> None: - """Say what unconfigure did, naming only keys whose value it changed.""" +@unconfigure_group.command(name="codex") +def unconfigure_codex() -> None: + """Return Codex's config to what it was before `lite configure codex`.""" + config_path: Final = codex_config_path(os.environ) + try: + outcome: Final = unconfigure_codex_config(config_path, codex_configure_state_path(config_path)) + except (AgentConfigError, ClaudeSettingsError) as e: + raise click.ClickException(str(e)) from e + _report_unconfigure(config_path, outcome, codex_configure_state_path(config_path)) + + +def _report_unconfigure(path: Path, outcome: UnconfigureOutcome, state_path: Path | None = None) -> None: if outcome.file_removed: - click.echo( - f"No settings file remains at {settings_path}; it held nothing but `lite configure claude`'s own keys." - ) + click.echo(f"No settings file remains at {path}; it held nothing but `lite configure`'s own keys.") elif outcome.restored: - click.echo(f"Restored in {settings_path}: {', '.join(outcome.restored)}.") + click.echo(f"Restored in {path}: {', '.join(outcome.restored)}.") else: - click.echo(f"Nothing in {settings_path} was still ours to restore.") + click.echo(f"Nothing in {path} was still ours to restore.") if outcome.kept: click.echo(f"Left as you changed them since: {', '.join(outcome.kept)}.") if outcome.withheld: click.echo( "Left removed, since the file now points at a different server than they were issued for: " - + "; ".join(f"{item.key} (captured with {item.endpoint})" for item in outcome.withheld) - + f". They stay in {state_path}: point env.ANTHROPIC_BASE_URL back and run `lite unconfigure claude` " - "again to put them back, or delete that file to drop them." + + "; ".join(item.key + f" (captured with {item.endpoint})" for item in outcome.withheld) + + f". They stay in {state_path}: point the endpoint back and run `lite unconfigure` again." ) -__all__ = ("configure_group", "interactive_configure", "resolve_credential", "unconfigure_group") +__all__ = ("configure_group", "interactive_configure", "unconfigure_group") diff --git a/litellm/proxy/client/cli/commands/configured_launch.py b/litellm/proxy/client/cli/commands/configured_launch.py new file mode 100644 index 00000000000..91a239b4314 --- /dev/null +++ b/litellm/proxy/client/cli/commands/configured_launch.py @@ -0,0 +1,133 @@ +import os +import shlex +import shutil +import subprocess +import sys +from collections.abc import Callable, Generator, Mapping, Sequence +from contextlib import contextmanager +from pathlib import Path +from typing import Final + +import click + +from .agents import AgentRunError, hand_off, restore_controlling_terminal, windows_command + +_CREATE_NEW_CONSOLE: Final = 0x00000010 + +_ROUTING_ENV: Final = frozenset( + { + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_BASE_URL", + "ANTHROPIC_DEFAULT_FABLE_MODEL", + "ANTHROPIC_DEFAULT_HAIKU_MODEL", + "ANTHROPIC_DEFAULT_OPUS_MODEL", + "ANTHROPIC_DEFAULT_SONNET_MODEL", + "ANTHROPIC_MODEL", + "LITELLM_PROXY_API_KEY", + "LITELLM_PROXY_URL", + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + } +) +_CONFIG_ENV: Final = {"claude": "CLAUDE_CONFIG_DIR", "codex": "CODEX_HOME"} + + +@contextmanager +def _launch_errors(agents: Sequence[str]) -> Generator[None, None, None]: + try: + yield + except (AgentRunError, OSError, subprocess.SubprocessError) as e: + commands: Final = " and ".join(f"`{agent}`" for agent in agents) + raise click.ClickException( + f"Could not open {commands}. Configuration was saved; start {commands} manually. {e}" + ) from e + + +def _agent_env(agent: str, environ: Mapping[str, str]) -> dict[str, str]: + config_key: Final = _CONFIG_ENV[agent] + return { + key: value + for key, value in environ.items() + if key not in _ROUTING_ENV and (key not in _CONFIG_ENV.values() or key == config_key) + } + + +def _shell_command(agent: str, binary: str, cwd: str, environ: Mapping[str, str], quote: Callable[[str], str]) -> str: + config_key: Final = _CONFIG_ENV[agent] + config: Final = environ.get(config_key) + assignment: Final = () if config is None else (f"{config_key}={quote(config)}",) + unsets: Final = tuple(part for key in sorted(_ROUTING_ENV) for part in ("-u", key)) + command: Final = " ".join(("env", *unsets, *assignment, quote(binary))) + return f"cd -- {quote(cwd)} && exec {command}" + + +def _launch_terminal( + agent: str, + binary: str, + cwd: str, + environ: Mapping[str, str], + platform: str, + which: Callable[[str], str | None], + run: Callable[..., subprocess.CompletedProcess[str]], + popen: Callable[..., subprocess.Popen[bytes]], +) -> None: + env: Final = _agent_env(agent, environ) + if platform.startswith("win"): + popen(windows_command(binary, (agent,)), cwd=cwd, env=env, creationflags=_CREATE_NEW_CONSOLE) + return + command: Final = _shell_command(agent, binary, cwd, environ, shlex.quote) + if platform == "darwin": + osascript: Final = which("osascript") + if osascript is None: + raise AgentRunError("Terminal is unavailable.") + script: Final = 'on run argv\ntell application "Terminal" to do script (item 1 of argv)\nend run' + run((str(Path(osascript).resolve()), "-e", script, command), env=env, check=True) + return + terminal: Final = which("x-terminal-emulator") + if terminal is None: + raise AgentRunError("No terminal emulator is available.") + popen((str(Path(terminal).resolve()), "-e", "/bin/sh", "-lc", command), env=env) + + +def launch_configured_agents( + agents: Sequence[str], + *, + started_interactive: bool, + environ: Mapping[str, str] = os.environ, + platform: str = sys.platform, + cwd: str | None = None, + which: Callable[[str], str | None] = shutil.which, + hand_off: Callable[[str, Sequence[str], Mapping[str, str]], None] = hand_off, + restore_terminal: Callable[[], None] = restore_controlling_terminal, + run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run, + popen: Callable[..., subprocess.Popen[bytes]] = subprocess.Popen, +) -> None: + selected: Final = tuple(dict.fromkeys(agents)) + if not selected or any(agent not in _CONFIG_ENV for agent in selected): + raise click.ClickException("Choose Claude Code, Codex, or both.") + with _launch_errors(selected): + binaries: Final = tuple((agent, which(agent)) for agent in selected) + missing: Final = tuple(agent for agent, binary in binaries if binary is None) + if missing: + raise AgentRunError(f"Could not find {', '.join(missing)} on PATH.") + resolved: Final = tuple( + (agent, str(Path(binary).resolve()) if not platform.startswith("win") else binary) + for agent, binary in binaries + if binary is not None + ) + launch_cwd: Final = cwd or os.getcwd() + errors: Final[list[str]] = [] # mutable-ok: accumulate independent launch failures before reporting them together + for agent, binary in resolved: + try: + with _launch_errors((agent,)): + if len(resolved) == 1: + if started_interactive: + restore_terminal() + hand_off(binary, (agent,), _agent_env(agent, environ)) + else: + _launch_terminal(agent, binary, launch_cwd, environ, platform, which, run, popen) + except click.ClickException as e: + errors.append(e.message) + if errors: + raise click.ClickException("\n".join(errors)) diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index f2624797a5f..c9871053f94 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -124,7 +124,7 @@ def _stored_login_is_pkce(vault: SecretVault) -> bool: return token_data is not None and token_data.get("refresh_token") is not None -def ensure_fresh_login(ctx: click.Context) -> None: +def ensure_fresh_login(ctx: click.Context, reader: str = "the agent") -> None: ctx_obj: Final[CliContextObj] = ctx.obj base_url: Final = ctx_obj["base_url"].rstrip("/") vault: Final = context_secret_vault(ctx) @@ -134,7 +134,7 @@ def ensure_fresh_login(ctx: click.Context) -> None: pkce: Final = _stored_login_is_pkce(vault) login_command: Final = "lite login --pkce" if pkce else "lite login" if not sys.stdin.isatty(): - raise UpError(f"No fresh LiteLLM login found for this proxy. Run `{login_command}` first.") + raise UpError(f"No fresh LiteLLM login found for this proxy. Run `{login_command}` first ({reader}).") click.echo("No fresh LiteLLM login found for this proxy; starting login...") ctx.invoke(login, config_claude=False, pkce=pkce) diff --git a/pyproject.toml b/pyproject.toml index 448451f7f93..b33c56094a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,6 +72,7 @@ proxy = [ "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", + "tomlkit>=0.13.0,<1.0", "polars>=1.38.1,<2.0", "soundfile>=0.12.1,<1.0", "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", @@ -88,6 +89,7 @@ cli = [ "requests>=2.32.0,<3.0", "InquirerPy>=0.3.4,<1.0", "keyring>=25.6.0,<26.0", + "tomlkit>=0.13.0,<1.0", ] extra_proxy = [ "prisma>=0.11.0,<1.0", diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index cf52d41e963..c2d3837e236 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -3,13 +3,12 @@ import os import pathlib import shlex import stat +import subprocess import sys -import time from pathlib import Path from unittest.mock import patch import pytest -from click.testing import CliRunner from litellm.litellm_core_utils.private_json import commit_staged_json from litellm.proxy.client.cli.commands.claude_settings import ( @@ -17,7 +16,6 @@ from litellm.proxy.client.cli.commands.claude_settings import ( AUTOROUTE_BACKUP_PATH, BACKUP_PATH, CLAUDE_SETTINGS_PATH, - CONFIGURE_STATE_PATH, OWNED_ENV_KEYS, OWNED_TOP_LEVEL_KEYS, SETTINGS_FILE_OWNERS, @@ -29,9 +27,10 @@ from litellm.proxy.client.cli.commands.claude_settings import ( UnpinModel, claude_settings_path, configure_claude_settings, - install_statusline_script, configure_state_path, + install_statusline_script, merge_claude_settings, + print_token_command, statusline_command, unconfigure_claude_settings, with_status_line, @@ -45,6 +44,8 @@ def _owners(*backup_paths): CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" AUTH_MODULE = "litellm.proxy.client.cli.commands.auth" + + @pytest.fixture def paths(tmp_path): return tmp_path / "claude" / "settings.json", tmp_path / "backup.json" @@ -53,7 +54,21 @@ def paths(tmp_path): def _static_configure(base_url, settings_path, owners, state_path=None): """`lite configure claude --api-key`'s shape: a virtual key as a static token, no pinned model.""" state = state_path if state_path is not None else settings_path.parent.parent / "state.json" - configure_claude_settings(base_url.rstrip("/"), StaticToken("sk-virtual-key"), KeepModel(), settings_path, state, owners) + configure_claude_settings( + base_url.rstrip("/"), StaticToken("sk-virtual-key"), KeepModel(), settings_path, state, owners + ) + + +def test_print_token_command_persists_an_absolute_executable(monkeypatch, tmp_path): + worktree = tmp_path / "untrusted-project" + relative_bin = worktree / "tools" + relative_bin.mkdir(parents=True) + lite = relative_bin / "lite" + lite.touch() + lite.chmod(0o700) + monkeypatch.chdir(worktree) + monkeypatch.setenv("PATH", f"tools{os.pathsep}{os.environ.get('PATH', '')}") + assert print_token_command("https://proxy.example.com")[0] == str(lite.resolve()) class TestConfigureClaudeSettings: @@ -98,7 +113,9 @@ class TestConfigureClaudeSettings: settings_path, backup_path = paths settings_path.parent.mkdir(parents=True) settings_path.write_text( - json.dumps({"apiKeyHelper": "/usr/local/bin/lite auth print-token", "env": {"ANTHROPIC_API_KEY": "sk-leaked"}}) + json.dumps( + {"apiKeyHelper": "/usr/local/bin/lite auth print-token", "env": {"ANTHROPIC_API_KEY": "sk-leaked"}} + ) ) _static_configure("https://proxy.example.com", settings_path, _owners(backup_path)) @@ -245,6 +262,29 @@ class TestClaudeSettingsPath: Path.home() / ".claude-work" / "settings.json" ) + def test_an_import_time_override_does_not_redefine_the_default_file(self, tmp_path): + home = tmp_path / "home" + profile = tmp_path / "profile" + script = """ +import os +from pathlib import Path +from litellm.proxy.client.cli.commands.claude_settings import ( + CLAUDE_SETTINGS_PATH, + configure_state_path, + settings_file_owners, +) +profile = Path(os.environ["CLAUDE_CONFIG_DIR"]) / "settings.json" +print(CLAUDE_SETTINGS_PATH) +print(configure_state_path(profile)) +print(len(settings_file_owners(profile))) +""" + env = {**os.environ, "HOME": str(home), "USERPROFILE": str(home), "CLAUDE_CONFIG_DIR": str(profile)} + result = subprocess.run([sys.executable, "-c", script], env=env, check=True, capture_output=True, text=True) + default, receipt, owners = result.stdout.splitlines() + assert Path(default) == home / ".claude" / "settings.json" + assert Path(receipt).parent == home / ".litellm" / "claude_configure_state" + assert owners == "0" + class TestConfigureStatePath: """Each settings file gets its own undo receipt: the default file keeps the long-standing path, and diff --git a/tests/test_litellm/proxy/client/cli/test_codex_settings.py b/tests/test_litellm/proxy/client/cli/test_codex_settings.py new file mode 100644 index 00000000000..7080667040f --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_codex_settings.py @@ -0,0 +1,332 @@ +import stat +from pathlib import Path + +import pytest +import tomlkit + +from litellm.proxy.client.cli.commands.agent_config import AgentConfigError, ConfigureReceipt +from litellm.proxy.client.cli.commands.claude_settings import ApiKeyHelper, KeepModel, StartOn, StaticToken, UnpinModel +from litellm.proxy.client.cli.commands.codex_settings import ( + OWNED_ROOT_KEYS, + codex_configure_state_path, + codex_home, + configure_codex_config, + unconfigure_codex_config, +) + +BASE_URL = "http://127.0.0.1:4000" +PRINT_TOKEN = ("/usr/local/bin/lite", "--base-url", BASE_URL, "auth", "print-token") +ORIGINAL = """# my codex config +model = "gpt-6-astra" # pinned by hand +model_reasoning_effort = "xhigh" + +[model_providers.other] +name = "Other" +base_url = "https://other.example.com/v1" + +[projects."/Users/me"] +trust_level = "trusted" +""" + + +@pytest.fixture +def paths(tmp_path): + return tmp_path / "codex" / "config.toml", tmp_path / "state" / "codex_configure_state.json" + + +def _configure(paths, model=StartOn("claude-auto"), credential=StaticToken("sk-virtual-key"), context_window=200000): + config_path, state_path = paths + configure_codex_config(BASE_URL, credential, lambda: PRINT_TOKEN, model, context_window, config_path, state_path) + + +def _no_lite_on_path(): + raise AgentConfigError("Could not find `lite` on your PATH.") + + +def _failing_commit(config_path, state_path, break_restore): + landed_receipt = False + + def commit(staged, path): + nonlocal landed_receipt + from litellm.litellm_core_utils.private_json import commit_staged_json + + if path == str(config_path): + Path(staged).unlink() + if break_restore: + state_path.unlink() + state_path.mkdir() + raise OSError("config rename failed") + if break_restore and path == str(state_path) and landed_receipt: + Path(staged).unlink() + raise OSError("receipt rollback failed") + commit_staged_json(staged, path) + if path == str(state_path): + landed_receipt = True + + return commit + + +class TestConfigureCodex: + def test_wires_the_provider_and_pins_the_model_keeping_comments_and_other_tables(self, paths): + config_path, state_path = paths + config_path.parent.mkdir(parents=True) + config_path.write_text(ORIGINAL) + + _configure(paths) + + text = config_path.read_text() + doc = tomlkit.parse(text).unwrap() + assert doc["model_provider"] == "litellm" + assert doc["model"] == "claude-auto" + assert doc["model_context_window"] == 200000 + assert doc["model_providers"]["litellm"] == { + "name": "LiteLLM proxy", + "base_url": f"{BASE_URL}/v1", + "wire_api": "responses", + "supports_websockets": False, + "experimental_bearer_token": "sk-virtual-key", + } + assert doc["model_providers"]["other"] == {"name": "Other", "base_url": "https://other.example.com/v1"} + assert doc["projects"]["/Users/me"] == {"trust_level": "trusted"} + assert "# my codex config" in text and "# pinned by hand" in text + assert stat.S_IMODE(config_path.stat().st_mode) == 0o600 + assert stat.S_IMODE(state_path.stat().st_mode) == 0o600 + + def test_the_login_credential_becomes_a_command_codex_runs_itself(self, paths): + config_path, _ = paths + _configure(paths, credential=ApiKeyHelper("ignored for codex"), model=UnpinModel(), context_window=None) + provider = tomlkit.parse(config_path.read_text()).unwrap()["model_providers"]["litellm"] + assert provider["auth"] == {"command": PRINT_TOKEN[0], "args": list(PRINT_TOKEN[1:]), "timeout_ms": 5000} + assert "experimental_bearer_token" not in provider + assert "model" not in tomlkit.parse(config_path.read_text()).unwrap() + + def test_the_receipt_never_holds_the_key(self, paths): + _, state_path = paths + _configure(paths, credential=StaticToken("sk-never-on-disk-twice")) + assert "sk-never-on-disk-twice" not in state_path.read_text() + assert set(ConfigureReceipt.model_validate_json(state_path.read_bytes()).sections[""].written) == set( + OWNED_ROOT_KEYS + ) + + def test_reports_invalid_toml_without_touching_the_file(self, paths): + config_path, _ = paths + config_path.parent.mkdir(parents=True) + config_path.write_text("model = [unterminated\n") + with pytest.raises(AgentConfigError, match="not valid TOML"): + _configure(paths) + assert config_path.read_text() == "model = [unterminated\n" + + def test_a_static_key_never_needs_lite_on_path(self, paths): + # A desktop-launched Codex reads the key from the file; only the login credential runs lite. + config_path, state_path = paths + configure_codex_config( + BASE_URL, + StaticToken("sk-virtual-key"), + _no_lite_on_path, + StartOn("claude-auto"), + None, + config_path, + state_path, + ) + assert tomlkit.parse(config_path.read_text()).unwrap()["model_providers"]["litellm"][ + "experimental_bearer_token" + ] + with pytest.raises(AgentConfigError, match="lite"): + configure_codex_config( + BASE_URL, + ApiKeyHelper("ignored for codex"), + _no_lite_on_path, + StartOn("claude-auto"), + None, + config_path, + state_path, + ) + + def test_a_repin_to_a_model_with_no_known_window_drops_the_old_window(self, paths): + # Codex sizes its context from model_context_window; a stale 200k from the previous pin would be wrong. + config_path, _ = paths + _configure(paths, model=StartOn("claude-auto"), context_window=200000) + _configure(paths, model=StartOn("gpt-5.6-luna"), context_window=None) + after = tomlkit.parse(config_path.read_text()).unwrap() + assert after["model"] == "gpt-5.6-luna" and "model_context_window" not in after + + @pytest.mark.parametrize( + ("shape", "text"), + [ + ("dotted", 'model_providers.other.name = "Other"\nmodel = "gpt-6-astra"\n'), + ("inline", 'model_providers = { other = { name = "Other" } }\nmodel = "gpt-6-astra"\n'), + ( + "split", + '[model_providers.a]\nname = "a"\n\n[projects."/Users/me"]\ntrust_level = "trusted"\n\n[model_providers.b]\nname = "b"\n', + ), + ], + ) + def test_every_toml_object_shape_of_model_providers_round_trips(self, paths, shape, text): + # TOML spells an object as a standard table, dotted keys, an inline table, or fragments split by + # other tables; configure must extend all of them and unconfigure must restore all of them. + config_path, state_path = paths + config_path.parent.mkdir(parents=True) + config_path.write_text(text) + _configure(paths) + configured = tomlkit.parse(config_path.read_text()).unwrap() + assert configured["model_providers"]["litellm"]["wire_api"] == "responses" + assert configured["model_provider"] == "litellm" + original_providers = tomlkit.parse(text).unwrap()["model_providers"] + assert {k: v for k, v in configured["model_providers"].items() if k != "litellm"} == original_providers + unconfigure_codex_config(config_path, state_path) + assert tomlkit.parse(config_path.read_text()).unwrap() == tomlkit.parse(text).unwrap() + + def test_toml_dates_in_owned_keys_are_fingerprinted_not_rejected(self, paths): + config_path, state_path = paths + config_path.parent.mkdir(parents=True) + config_path.write_text("model = 2024-01-01\n[model_providers.litellm]\nsince = 1979-05-27T07:32:00Z\n") + _configure(paths) + unconfigure_codex_config(config_path, state_path) + after = tomlkit.parse(config_path.read_text()).unwrap() + assert str(after["model"]) == "2024-01-01" and "since" in after["model_providers"]["litellm"] + + +class TestConfigureRollback: + @pytest.mark.parametrize("repeat", [False, True], ids=["first", "repeat"]) + def test_config_failure_restores_the_previous_receipt(self, paths, repeat): + config_path, state_path = paths + if repeat: + _configure(paths) + receipt_before = state_path.read_bytes() + config_before = config_path.read_bytes() + else: + receipt_before = None + config_before = None + with pytest.raises(AgentConfigError, match="config rename failed"): + configure_codex_config( + BASE_URL, + StaticToken("sk-rotated"), + lambda: PRINT_TOKEN, + StartOn("claude-auto"), + 200000, + config_path, + state_path, + commit=_failing_commit(config_path, state_path, False), + ) + assert (state_path.read_bytes() if state_path.exists() else None) == receipt_before + assert (config_path.read_bytes() if config_path.exists() else None) == config_before + assert not list(state_path.parent.glob(".tmp-*")) and not list(config_path.parent.glob(".tmp-*")) + + @pytest.mark.parametrize("repeat", [False, True], ids=["first", "repeat"]) + def test_double_failure_names_the_stale_receipt_and_keeps_original_cause(self, paths, repeat): + config_path, state_path = paths + config_before = None + if repeat: + _configure(paths) + config_before = config_path.read_bytes() + with pytest.raises(AgentConfigError, match="could not be put back") as exc_info: + configure_codex_config( + BASE_URL, + StaticToken("sk-rotated"), + lambda: PRINT_TOKEN, + StartOn("claude-auto"), + 200000, + config_path, + state_path, + commit=_failing_commit(config_path, state_path, True), + ) + assert str(state_path) in str(exc_info.value) + assert "remove it before retrying" in str(exc_info.value) + assert isinstance(exc_info.value.__cause__, OSError) + assert str(exc_info.value.__cause__) == "config rename failed" + assert state_path.exists() + assert (config_path.read_bytes() if config_path.exists() else None) == config_before + assert not list(config_path.parent.glob(".tmp-*")) + + +class TestUnconfigureCodex: + def test_returns_the_file_to_its_original_bytes(self, paths): + config_path, state_path = paths + config_path.parent.mkdir(parents=True) + config_path.write_text(ORIGINAL) + _configure(paths) + + outcome = unconfigure_codex_config(config_path, state_path) + + assert tomlkit.parse(config_path.read_text()).unwrap() == tomlkit.parse(ORIGINAL).unwrap() + assert "# pinned by hand" in config_path.read_text() + assert not state_path.exists() + assert outcome.kept == () + assert set(outcome.restored) == {"model_provider", "model", "model_context_window", "model_providers.litellm"} + + def test_removes_a_config_that_only_configure_created(self, paths): + config_path, state_path = paths + _configure(paths) + unconfigure_codex_config(config_path, state_path) + assert not config_path.exists() + + def test_leaves_keys_the_user_changed_since(self, paths): + config_path, state_path = paths + config_path.parent.mkdir(parents=True) + config_path.write_text(ORIGINAL) + _configure(paths) + doc = tomlkit.parse(config_path.read_text()) + doc["model"] = "claude-sonnet-4-6" + config_path.write_text(tomlkit.dumps(doc)) + + outcome = unconfigure_codex_config(config_path, state_path) + after = tomlkit.parse(config_path.read_text()).unwrap() + assert after["model"] == "claude-sonnet-4-6" + assert "model_provider" not in after and "litellm" not in after["model_providers"] + assert outcome.kept == ("model",) + + def test_a_key_the_user_edited_between_two_configures_stays_theirs(self, paths): + # A repeat configure (a re-login is one) must not adopt the user's edit as its own write and + # then delete it on unconfigure. + config_path, state_path = paths + config_path.parent.mkdir(parents=True) + config_path.write_text(ORIGINAL) + _configure(paths, model=StartOn("claude-auto")) + doc = tomlkit.parse(config_path.read_text()) + doc["model"] = "my-favourite" + config_path.write_text(tomlkit.dumps(doc)) + _configure(paths, model=KeepModel(), credential=StaticToken("sk-rotated"), context_window=None) + assert tomlkit.parse(config_path.read_text()).unwrap()["model"] == "my-favourite" + outcome = unconfigure_codex_config(config_path, state_path) + after = tomlkit.parse(config_path.read_text()).unwrap() + assert after["model"] == "my-favourite" + assert "litellm" not in after["model_providers"] and "model_provider" not in after + assert "model" in outcome.kept + + def test_refuses_when_the_provider_section_became_a_scalar(self, paths): + config_path, state_path = paths + _configure(paths) + config_path.write_text('model_providers = "oops"\nmodel_provider = "litellm"\n') + with pytest.raises(AgentConfigError, match="non-object"): + unconfigure_codex_config(config_path, state_path) + assert state_path.exists() + + def test_a_relogin_keeps_the_pin_and_a_repeat_without_a_model_releases_it(self, paths): + config_path, state_path = paths + config_path.parent.mkdir(parents=True) + config_path.write_text(ORIGINAL) + _configure(paths, model=StartOn("claude-auto")) + _configure(paths, model=KeepModel(), credential=ApiKeyHelper("ignored for codex"), context_window=None) + assert tomlkit.parse(config_path.read_text()).unwrap()["model"] == "claude-auto" + _configure(paths, model=UnpinModel(), context_window=None) + after = tomlkit.parse(config_path.read_text()).unwrap() + assert after["model"] == "gpt-6-astra" and "model_context_window" not in after + unconfigure_codex_config(config_path, state_path) + assert tomlkit.parse(config_path.read_text()).unwrap() == tomlkit.parse(ORIGINAL).unwrap() + + +def test_relocated_codex_receipt_stays_under_litellm_home(tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path / "home")) + config_path = tmp_path / "codex-home" / "config.toml" + receipt = codex_configure_state_path(config_path) + assert receipt.parent.parent == Path.home() / ".litellm" + assert receipt.parent.name == "codex_configure_state" + assert receipt.name.endswith(".json") + assert receipt != config_path.parent / ".litellm-configure-state.json" + + +def test_codex_home_honors_the_env_override(monkeypatch, tmp_path): + monkeypatch.setenv("CODEX_HOME", str(tmp_path / "elsewhere")) + assert codex_home() == tmp_path / "elsewhere" + monkeypatch.delenv("CODEX_HOME") + assert codex_home() == Path.home() / ".codex" diff --git a/tests/test_litellm/proxy/client/cli/test_configure_commands.py b/tests/test_litellm/proxy/client/cli/test_configure_commands.py index 8ed188af737..97de17e2093 100644 --- a/tests/test_litellm/proxy/client/cli/test_configure_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_configure_commands.py @@ -1,16 +1,26 @@ +import asyncio +import io import json import os +import shlex import stat +import subprocess +import sys import click import pytest import requests import responses +import tomlkit from click.testing import CliRunner +from prompt_toolkit.application import create_app_session +from prompt_toolkit.input import create_pipe_input +from prompt_toolkit.output import DummyOutput from litellm.proxy.client.cli import cli from litellm.proxy.client.cli.commands import claude_settings as claude_settings_module from litellm.proxy.client.cli.commands import configure as configure_module +from litellm.proxy.client.cli.commands.agent_config import AgentConfigError from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner from litellm.proxy.client.cli.commands.configure import configure_claude, configure_group, interactive_configure @@ -68,7 +78,122 @@ def lite_up_backup(monkeypatch, tmp_path): def _configure(runner, *args): - return runner.invoke(cli, ["--base-url", PROXY, "configure", "claude", *args]) + return runner.invoke(cli, ["--base-url", PROXY, "configure", "--api-key", VALID_KEY, "claude", *args]) + + +@responses.activate +@pytest.mark.parametrize("reconfigure", [False, True], ids=["direct-undo", "reconfigure-then-undo"]) +def test_claude_commands_accept_pre_codex_flat_receipts(runner, paths, reconfigure): + settings_path, state_path = paths + settings_path.parent.mkdir(parents=True) + state_path.parent.mkdir(parents=True) + settings_path.write_text( + json.dumps( + { + "theme": "dark", + "model": "claude-auto", + "env": {"ANTHROPIC_BASE_URL": PROXY, "ANTHROPIC_AUTH_TOKEN": "sk-legacy-login"}, + } + ) + ) + state_path.write_text( + json.dumps( + { + "file_existed": True, + "env_present": True, + "env_was_object": True, + "previous": { + "model": {"present": True, "value": "original-model"}, + "env.ANTHROPIC_BASE_URL": {"present": True, "value": "https://old-proxy.example"}, + "env.ANTHROPIC_AUTH_TOKEN": {"present": True, "value": "original-token"}, + }, + "written": { + "model": "bbeb00a33788020610852e74a0af54a7ff3262d60f35be0cc38eb763ad9d1b23", + "env.ANTHROPIC_BASE_URL": "162f7a3bcb1750dd476b303bd2207abe025abd249bf00612580b58d1005b3bda", + "env.ANTHROPIC_AUTH_TOKEN": "b8edb3ee4e755b11d291a7e83af0536797d2ab5f315d13e52f31d6dee858dec1", + }, + "endpoints": { + "env.ANTHROPIC_AUTH_TOKEN": {"present": True, "value": "https://old-proxy.example"}, + }, + } + ) + ) + if reconfigure: + _mock_models() + result = _configure(runner, "--model", "claude-auto") + assert result.exit_code == 0, result.output + assert "sections" not in json.loads(state_path.read_text()) + result = runner.invoke(cli, ["unconfigure", "claude"]) + assert result.exit_code == 0, result.output + assert json.loads(settings_path.read_text()) == { + "theme": "dark", + "model": "original-model", + "env": {"ANTHROPIC_BASE_URL": "https://old-proxy.example", "ANTHROPIC_AUTH_TOKEN": "original-token"}, + } + assert not state_path.exists() + + +class TestConfigureCodexErrors: + @responses.activate + @pytest.mark.parametrize("command", ["configure", "unconfigure"]) + def test_reports_agent_config_errors_without_a_traceback(self, runner, monkeypatch, command): + if command == "configure": + _mock_models() + monkeypatch.setattr( + configure_module, + "configure_codex_config", + lambda *args, **kwargs: (_ for _ in ()).throw(AgentConfigError("bad codex config")), + ) + result = runner.invoke(cli, ["--base-url", PROXY, "configure", "--api-key", VALID_KEY, "codex"]) + else: + monkeypatch.setattr( + configure_module, + "unconfigure_codex_config", + lambda *args, **kwargs: (_ for _ in ()).throw(AgentConfigError("bad codex config")), + ) + result = runner.invoke(cli, ["unconfigure", "codex"]) + assert result.exit_code != 0 + assert "Error: bad codex config" in result.output + assert "Traceback" not in result.output + + @responses.activate + def test_reports_an_invalid_codex_receipt_without_a_traceback(self, runner, monkeypatch, tmp_path): + _mock_models() + config_path = tmp_path / "codex" / "config.toml" + state_path = tmp_path / "codex" / ".litellm-configure-state.json" + state_path.parent.mkdir(parents=True) + state_path.write_text("not json") + monkeypatch.setattr(configure_module, "codex_config_path", lambda environ: config_path) + monkeypatch.setattr(configure_module, "codex_configure_state_path", lambda path: state_path) + + result = runner.invoke(cli, ["--base-url", PROXY, "configure", "--api-key", VALID_KEY, "codex"]) + + assert result.exit_code != 0 + assert "Error:" in result.output + assert "Traceback" not in result.output + + +class TestApiKeyOptionOrdering: + @responses.activate + @pytest.mark.parametrize( + "args", + [ + ("configure", "--api-key", VALID_KEY, "claude"), + ("configure", "claude", "--api-key", VALID_KEY), + ("configure", "--api-key", VALID_KEY, "codex"), + ("configure", "codex", "--api-key", VALID_KEY), + ], + ) + def test_accepts_parent_and_legacy_subcommand_order(self, runner, paths, tmp_path, monkeypatch, args): + _mock_models() + monkeypatch.setenv("CODEX_HOME", str(tmp_path / "codex")) + result = runner.invoke(cli, ["--base-url", PROXY, *args]) + assert result.exit_code == 0, result.output + if "claude" in args: + assert json.loads(paths[0].read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + else: + provider = tomlkit.parse((tmp_path / "codex" / "config.toml").read_text())["model_providers"]["litellm"] + assert provider["experimental_bearer_token"] == VALID_KEY class TestConfigureClaudeWithAVirtualKey: @@ -76,7 +201,7 @@ class TestConfigureClaudeWithAVirtualKey: def test_writes_settings_and_reports_without_echoing_the_key(self, runner, paths): _mock_models() settings_path, state_path = paths - result = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto") + result = _configure(runner, "--model", "claude-auto") assert result.exit_code == 0, result.output written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == PROXY @@ -107,7 +232,7 @@ class TestConfigureClaudeWithAVirtualKey: def test_refuses_a_model_the_proxy_does_not_list(self, runner, paths): _mock_models() settings_path, _ = paths - result = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-nope") + result = _configure(runner, "--model", "claude-nope") assert result.exit_code != 0 assert "'claude-nope' is not served" in result.output assert "claude-auto, gpt-5.6-luna" in result.output @@ -117,7 +242,7 @@ class TestConfigureClaudeWithAVirtualKey: def test_refuses_a_key_the_proxy_rejects(self, runner, paths): _mock_models() settings_path, _ = paths - result = _configure(runner, "--api-key", "sk-wrong") + result = runner.invoke(cli, ["--base-url", PROXY, "configure", "--api-key", "sk-wrong", "claude"]) assert result.exit_code != 0 assert "rejected your key (HTTP 401)" in result.output assert not settings_path.exists() @@ -154,7 +279,7 @@ class TestConfigureClaudeWithAVirtualKey: # empty list prove it is up, and the hint says so instead. mock() settings_path, _ = paths - result = _configure(runner, "--api-key", VALID_KEY) + result = _configure(runner) assert result.exit_code != 0 assert expected in result.output and unexpected not in result.output assert not settings_path.exists() @@ -166,10 +291,15 @@ class TestConfigureClaudeWithAVirtualKey: if entry == "interactive": ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY}) with pytest.raises(click.ClickException, match="lite down"): - interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=lambda listed: None) + interactive_configure( + ctx, + pick_targets=lambda: ("claude",), + pick_model=lambda _agent, listed: None, + confirm_launch=lambda _agent: False, + ) else: - args = ["--api-key", VALID_KEY] if entry == "virtual-key" else [] - result = runner.invoke(configure_claude, args, obj={"base_url": PROXY, "api_key": None}) + args = ["--api-key", VALID_KEY, "claude"] if entry == "virtual-key" else ["claude"] + result = runner.invoke(configure_group, args, obj={"base_url": PROXY, "api_key": None}) assert result.exit_code != 0 and "lite down" in result.output assert len(responses.calls) == 0 assert not paths[0].exists() @@ -183,7 +313,7 @@ class TestConfigureClaudeWithAVirtualKey: target.write_text("{}") settings_path.parent.mkdir(parents=True) settings_path.symlink_to(target) - result = _configure(runner, "--api-key", VALID_KEY) + result = _configure(runner) assert result.exit_code == 0, result.output assert "keep it out of version control" in result.output assert json.loads(target.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY @@ -212,8 +342,8 @@ class TestConfigureClaudeWithoutAKey: _mock_models() settings_path, _ = paths result = runner.invoke( - configure_claude, - ["--api-key", VALID_KEY], + configure_group, + ["--api-key", VALID_KEY, "claude"], obj={"base_url": PROXY, "api_key": "sk-login-jwt", "api_key_from_token_file": True}, ) assert result.exit_code == 0, result.output @@ -221,6 +351,147 @@ class TestConfigureClaudeWithoutAKey: assert written["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY and "apiKeyHelper" not in written +class _TerminalInput(io.StringIO): + def isatty(self): + return True + + +@pytest.mark.timeout(20) +@responses.activate +@pytest.mark.parametrize( + ("keys", "claude_model", "codex_model", "launched"), + [ + (("\r", "\r", "n"), None, None, None), + (("\r", "\x1b[B\r", "y"), "claude-router", None, "claude"), + ((" \x1b[B \r", "\x1b[B\r", "n"), None, "route", None), + (("\x1b[B \r", "\x1b[B\r", "\x1b[B\r", "n", "y"), "claude-router", "route", "codex"), + ], + ids=["default-decline", "claude-launch", "codex-decline", "both-launch-second"], +) +def test_bare_configure_drives_real_prompts(keys, claude_model, codex_model, launched, paths, tmp_path, monkeypatch): + codex_home = tmp_path / "codex" + monkeypatch.setenv("CODEX_HOME", str(codex_home)) + monkeypatch.setattr(sys, "stdin", _TerminalInput()) + responses.get( + f"{PROXY}/v1/models", + json={"data": [{"id": "claude-router", "source_model": "route"}]}, + match=[responses.matchers.header_matcher({"x-gateway-client": "claude-code"})], + ) + responses.get(f"{PROXY}/v1/models", json={"data": [{"id": "route"}]}) + responses.get( + f"{PROXY}/model_group/info", + json={"data": [{"model_group": "route", "max_input_tokens": 200000, "max_output_tokens": 64000}]}, + ) + launches = [] + + def launch(agents, *, started_interactive): + assert started_interactive + assert paths[0].exists() or (codex_home / "config.toml").exists() + launches.extend(agents) + + monkeypatch.setattr(configure_module, "launch_configured_agents", launch) + + async def drive(): + with create_pipe_input() as pipe, create_app_session(input=pipe, output=DummyOutput()) as session: + task = asyncio.create_task( + asyncio.to_thread( + cli.main, args=["--base-url", PROXY, "--api-key", VALID_KEY, "configure"], standalone_mode=False + ) + ) + previous = None + try: + for text in keys: + + async def next_prompt(previous=previous): + while (app := session.app) is None or app is previous or not app.is_running: + if task.done(): + await task + raise AssertionError("CLI ended before the next prompt") + await asyncio.sleep(0.01) + return app + + previous = await asyncio.wait_for(next_prompt(), timeout=5) + pipe.send_text(text) + await asyncio.wait_for(task, timeout=5) + finally: + pipe.close() + + asyncio.run(drive()) + assert launches == ([] if launched is None else [launched]) + if claude_model is not None or len(keys) == 3 and keys[1] == "\r": + settings = json.loads(paths[0].read_text()) + assert settings.get("model") == claude_model + assert settings["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY + else: + assert not paths[0].exists() + if codex_model is not None: + config = tomlkit.parse((codex_home / "config.toml").read_text()) + assert config["model"] == codex_model and config["model_context_window"] == 200000 + assert config["model_providers"]["litellm"]["experimental_bearer_token"] == VALID_KEY + else: + assert not (codex_home / "config.toml").exists() + + +@responses.activate +@pytest.mark.parametrize("entry", ["subcommand", "interactive"]) +def test_codex_login_writes_executable_argv(entry, paths, tmp_path, monkeypatch, runner): + codex_home = tmp_path / "codex" + bin_dir = tmp_path / "bin with spaces" + bin_dir.mkdir() + lite = bin_dir / "lite" + lite.write_text( + f"#!/bin/sh\nexec {shlex.quote(sys.executable)} -c 'import json, sys; print(json.dumps(sys.argv[1:]))' \"$@\"\n" + ) + lite.chmod(0o700) + monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ.get('PATH', '')}") + monkeypatch.setenv("CODEX_HOME", str(codex_home)) + monkeypatch.setattr(sys, "stdin", _TerminalInput()) + + def login(ctx, reader): + monkeypatch.setattr(sys, "stdin", io.StringIO()) + + monkeypatch.setattr(configure_module, "ensure_fresh_login", login) + monkeypatch.setattr(configure_module, "get_stored_api_key", lambda **kwargs: "short-lived-login") + responses.get(f"{PROXY}/v1/models", json={"data": [{"id": "route"}]}) + responses.get(f"{PROXY}/model_group/info", json={"data": []}) + launches = [] + monkeypatch.setattr( + configure_module, + "launch_configured_agents", + lambda agents, *, started_interactive: launches.extend((tuple(agents), started_interactive)), + ) + if entry == "interactive": + ctx = click.Context(configure_group, obj={"base_url": PROXY, "api_key": None}) + interactive_configure( + ctx, + pick_targets=lambda: ("codex",), + pick_model=lambda agent, listed: "route", + confirm_launch=lambda agent: True, + ) + assert launches == [("codex",), True] + else: + result = runner.invoke(configure_module.configure_codex, [], obj={"base_url": PROXY, "api_key": None}) + assert result.exit_code == 0, result.output + provider = tomlkit.parse((codex_home / "config.toml").read_text())["model_providers"]["litellm"] + auth = provider["auth"] + assert auth["command"] == str(lite) + executed = subprocess.run([auth["command"], *auth["args"]], capture_output=True, text=True, check=True, timeout=10) + assert json.loads(executed.stdout) == ["--base-url", PROXY, "auth", "print-token"] + assert "short-lived-login" not in (codex_home / "config.toml").read_text() + assert "experimental_bearer_token" not in provider + + +@responses.activate +@pytest.mark.parametrize("agent,label", [("claude", "Claude Code"), ("codex", "Codex")]) +def test_empty_listing_names_selected_agent(agent, label, runner, paths, tmp_path, monkeypatch): + monkeypatch.setenv("CODEX_HOME", str(tmp_path / "codex")) + responses.get(f"{PROXY}/v1/models", json={"data": []}) + result = runner.invoke(cli, ["--base-url", PROXY, "--api-key", VALID_KEY, "configure", agent]) + assert result.exit_code == 1 + assert f"{label} would have nothing to run" in result.output + assert not paths[0].exists() and not (tmp_path / "codex" / "config.toml").exists() + + class TestInteractiveConfigure: @responses.activate def test_asks_for_targets_and_a_starting_model_then_configures(self, paths): @@ -228,14 +499,17 @@ class TestInteractiveConfigure: settings_path, _ = paths asked = {} - def pick_model(listed): + def pick_model(agent, listed): + assert agent == "claude" asked["listed"] = tuple(listed) return "claude-auto" ctx = click.Context( configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} ) - interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=pick_model) + interactive_configure( + ctx, pick_targets=lambda: ("claude",), pick_model=pick_model, confirm_launch=lambda _agent: False + ) assert asked["listed"] == LISTED_MODELS assert json.loads(settings_path.read_text())["model"] == "claude-auto" @@ -244,13 +518,13 @@ class TestInteractiveConfigure: ctx = click.Context( configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} ) - interactive_configure(ctx, pick_targets=lambda: (), pick_model=lambda listed: None) + interactive_configure(ctx, pick_targets=lambda: (), pick_model=lambda _agent, listed: None) assert not settings_path.exists() def test_bare_configure_without_a_terminal_names_the_non_interactive_command(self, runner, paths): result = runner.invoke(cli, ["--base-url", PROXY, "configure"]) assert result.exit_code != 0 - assert "lite configure claude --api-key" in result.output + assert "lite configure --api-key" in result.output class TestUnconfigureClaude: @@ -261,7 +535,7 @@ class TestUnconfigureClaude: settings_path.parent.mkdir(parents=True) original = {"theme": "dark", "model": "claude-opus-5"} settings_path.write_text(json.dumps(original)) - assert _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto").exit_code == 0 + assert _configure(runner, "--model", "claude-auto").exit_code == 0 result = runner.invoke(cli, ["unconfigure", "claude"]) assert result.exit_code == 0, result.output @@ -274,7 +548,7 @@ class TestUnconfigureClaude: def test_a_file_only_configure_created_is_reported_removed_not_restored(self, runner, paths): _mock_models() settings_path, _ = paths - assert _configure(runner, "--api-key", VALID_KEY).exit_code == 0 + assert _configure(runner).exit_code == 0 result = runner.invoke(cli, ["unconfigure", "claude"]) assert result.exit_code == 0, result.output assert not settings_path.exists() @@ -286,7 +560,7 @@ class TestUnconfigureClaude: settings_path, _ = paths settings_path.parent.mkdir(parents=True) settings_path.write_text(json.dumps({"theme": "dark"})) - assert _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto").exit_code == 0 + assert _configure(runner, "--model", "claude-auto").exit_code == 0 edited = json.loads(settings_path.read_text()) edited["env"] = {key: f"{value}-edited" for key, value in edited["env"].items()} edited["model"] = "mine" @@ -306,7 +580,7 @@ class TestUnconfigureClaude: settings_path.write_text( json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://api.anthropic.com", "ANTHROPIC_API_KEY": "sk-ant"}}) ) - assert _configure(runner, "--api-key", VALID_KEY).exit_code == 0 + assert _configure(runner).exit_code == 0 edited = json.loads(settings_path.read_text()) edited["env"]["ANTHROPIC_BASE_URL"] = "http://other-proxy:4000" settings_path.write_text(json.dumps(edited)) @@ -332,7 +606,7 @@ class TestUnconfigureClaude: (work_dir / "settings.json").write_text(json.dumps(original)) monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(work_dir)) - configured = _configure(runner, "--api-key", VALID_KEY, "--model", "claude-auto") + configured = _configure(runner, "--model", "claude-auto") assert configured.exit_code == 0, configured.output assert f"Configured Claude Code: {work_dir / 'settings.json'}" in configured.output assert json.loads((work_dir / "settings.json").read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == VALID_KEY @@ -381,7 +655,7 @@ class TestClaudeCodeView: ] ) settings_path, _ = paths - result = _configure(runner, "--api-key", VALID_KEY, "--model", model) + result = _configure(runner, "--model", model) assert result.exit_code == 0, result.output assert json.loads(settings_path.read_text())["model"] == pinned assert f"Starting model: {pinned}" in result.output @@ -391,7 +665,7 @@ class TestClaudeCodeView: def test_refuses_unknown_short_suffix(self, runner, paths): self._mock([{"id": "emitted-router-source", "source_model": "literal-router-source"}]) settings_path, _ = paths - result = _configure(runner, "--api-key", VALID_KEY, "--model", "source") + result = _configure(runner, "--model", "source") assert result.exit_code != 0 assert "'source' is not served" in result.output assert not settings_path.exists() @@ -405,17 +679,20 @@ class TestClaudeCodeView: configure_group, obj={"base_url": PROXY, "api_key": VALID_KEY, "api_key_from_token_file": False} ) - def pick_model(listed): + def pick_model(agent, listed): + assert agent == "claude" asked["listed"] = tuple(listed) return "source" - interactive_configure(ctx, pick_targets=lambda: ("claude",), pick_model=pick_model) + interactive_configure( + ctx, pick_targets=lambda: ("claude",), pick_model=pick_model, confirm_launch=lambda _agent: False + ) assert asked["listed"] == ("source",) assert json.loads(settings_path.read_text())["model"] == "emitted" @responses.activate def test_counts_what_an_older_proxy_lets_the_picker_show(self, runner, paths): _mock_models() - result = _configure(runner, "--api-key", VALID_KEY) + result = _configure(runner) assert result.exit_code == 0, result.output assert "/model will list 1 of the proxy's 2 models: Claude Code shows only ids containing" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_configured_launch.py b/tests/test_litellm/proxy/client/cli/test_configured_launch.py new file mode 100644 index 00000000000..a46020cb88b --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_configured_launch.py @@ -0,0 +1,270 @@ +import os +import subprocess +import sys +from pathlib import Path + +import click +import pytest +from click.testing import CliRunner + +from litellm.proxy.client.cli.commands.agents import AgentRunError +from litellm.proxy.client.cli.commands.configured_launch import launch_configured_agents + + +def test_single_agent_hands_off_with_only_its_persisted_config(tmp_path): + calls = [] + restored = [] + binary = tmp_path / "bin with spaces" / "claude" + launch_configured_agents( + ("claude",), + started_interactive=True, + environ={ + "PATH": "/usr/bin", + "CLAUDE_CONFIG_DIR": "/config with spaces/claude", + "CODEX_HOME": "/config/codex", + "ANTHROPIC_AUTH_TOKEN": "secret", + "ANTHROPIC_MODEL": "wrong-model", + "ANTHROPIC_DEFAULT_SONNET_MODEL": "wrong-tier", + "OPENAI_API_KEY": "secret-two", + }, + which=lambda command: str(binary) if command == "claude" else None, + hand_off=lambda path, argv, env: calls.append((path, argv, env)), + restore_terminal=lambda: restored.append(True), + ) + assert restored == [True] + assert calls == [ + ( + str(binary.resolve()), + ("claude",), + {"PATH": "/usr/bin", "CLAUDE_CONFIG_DIR": "/config with spaces/claude"}, + ) + ] + + +def test_two_agents_open_separate_macos_terminals_without_secrets(tmp_path): + calls = [] + binaries = { + "claude": str(tmp_path / "claude binary"), + "codex": str(tmp_path / "codex;binary"), + "osascript": str(tmp_path / "osascript"), + } + environ = { + "PATH": "/usr/bin", + "CLAUDE_CONFIG_DIR": "/claude profile", + "CODEX_HOME": "/codex;profile", + "ANTHROPIC_AUTH_TOKEN": "secret-a", + "OPENAI_API_KEY": "secret-b", + } + + def run(argv, **kwargs): + calls.append((argv, kwargs)) + + launch_configured_agents( + ("claude", "codex"), + started_interactive=True, + environ=environ, + platform="darwin", + cwd="/project; touch /tmp/nope", + which=binaries.get, + run=run, + ) + + assert len(calls) == 2 + assert all(call[0][0] == str(Path(binaries["osascript"]).resolve()) for call in calls) + assert all( + call[0][1:3] == ("-e", 'on run argv\ntell application "Terminal" to do script (item 1 of argv)\nend run') + for call in calls + ) + assert "secret" not in repr(calls) + assert "'" in calls[0][0][3] and "/project; touch /tmp/nope" in calls[0][0][3] + assert calls[0][1]["env"] == {"PATH": "/usr/bin", "CLAUDE_CONFIG_DIR": "/claude profile"} + assert calls[1][1]["env"] == {"PATH": "/usr/bin", "CODEX_HOME": "/codex;profile"} + + +def test_two_agents_use_linux_terminal_argv(tmp_path): + calls = [] + binaries = { + "claude": str(tmp_path / "claude"), + "codex": str(tmp_path / "codex"), + "x-terminal-emulator": str(tmp_path / "terminal"), + } + launch_configured_agents( + ("claude", "codex"), + started_interactive=False, + environ={"PATH": "/usr/bin"}, + platform="linux", + cwd="/work", + which=binaries.get, + popen=lambda argv, **kwargs: calls.append((argv, kwargs)), + ) + assert [call[0][:4] for call in calls] == [ + (str(Path(binaries["x-terminal-emulator"]).resolve()), "-e", "/bin/sh", "-lc"), + (str(Path(binaries["x-terminal-emulator"]).resolve()), "-e", "/bin/sh", "-lc"), + ] + + +@pytest.mark.parametrize("suffix", ["exe", "cmd", "bat"]) +def test_two_windows_agents_use_new_consoles_and_the_shared_command_builder(suffix): + calls = [] + binaries = { + "claude": rf"C:\Program Files\100% & !name!\claude.{suffix}", + "codex": rf"C:\Users\me\tools\codex.{suffix}", + } + workdir = r"C:\work & %PATH% !name!" + launch_configured_agents( + ("claude", "codex"), + started_interactive=True, + environ={"PATH": r"C:\Windows\System32", "ANTHROPIC_MODEL": "stale", "OPENAI_API_KEY": "secret"}, + platform="win32", + cwd=workdir, + which=binaries.get, + popen=lambda command, **kwargs: calls.append((command, kwargs)), + ) + assert len(calls) == 2 + for agent, (command, kwargs) in zip(("claude", "codex"), calls): + assert kwargs == { + "cwd": workdir, + "env": {"PATH": r"C:\Windows\System32"}, + "creationflags": 0x00000010, + } + if suffix == "exe": + assert command == (binaries[agent],) + else: + assert isinstance(command, str) + assert command.startswith('cmd.exe /d /e:on /v:off /s /c "') + assert command.endswith('""') + if agent == "claude": + assert "100%%cd:~,% & !name!" in command + assert "start" not in command + + +@pytest.mark.parametrize("newline", ["\n", "\r"]) +def test_windows_shim_path_cannot_inject_a_second_command(newline): + calls = [] + with pytest.raises(click.ClickException, match="line break"): + launch_configured_agents( + ("claude", "codex"), + started_interactive=True, + platform="win32", + which=lambda agent: f"C:\\tools\\{agent}{newline}injected.cmd", + popen=lambda *args, **kwargs: calls.append(args), + ) + assert calls == [] + + +@pytest.mark.skipif(sys.platform != "win32", reason="Requires Windows CreateProcess and cmd.exe") +@pytest.mark.parametrize("suffix", ["cmd", "bat"]) +def test_windows_new_console_executes_real_clients(tmp_path, suffix): + workdir = tmp_path / "work & %PATH% !name!" + workdir.mkdir() + binaries = {} + for agent in ("claude", "codex"): + directory = tmp_path / agent / "bin with spaces & %PATH% !name!" + directory.mkdir(parents=True) + binary = directory / f"{agent}.{suffix}" + binary.write_text('@echo off\nif defined ANTHROPIC_MODEL exit /b 1\ncd > "%~dp0launched.txt"\n') + binaries[agent] = str(binary) + processes = [] + + def spawn(command, **kwargs): + process = subprocess.Popen(command, **kwargs) + processes.append(process) + return process + + launch_configured_agents( + ("claude", "codex"), + started_interactive=True, + environ={**os.environ, "ANTHROPIC_MODEL": "must-not-leak"}, + cwd=str(workdir), + which=binaries.get, + popen=spawn, + ) + for process in processes: + assert process.wait(timeout=20) == 0 + for binary in binaries.values(): + assert Path(binary).with_name("launched.txt").read_text().strip() == str(workdir) + + +@pytest.mark.parametrize("platform", ["darwin", "linux", "win32"]) +@pytest.mark.parametrize("stage", ["which", "restore", "handoff", "first-terminal", "second-terminal"]) +@pytest.mark.parametrize( + "error", + [PermissionError("launch denied"), AgentRunError("invalid launch"), subprocess.CalledProcessError(1, "launcher")], + ids=["os-error", "agent-error", "process-error"], +) +def test_launch_errors_preserve_configs_and_report_recovery(tmp_path, platform, stage, error): + settings = tmp_path / "settings.json" + config = tmp_path / "config.toml" + settings.write_text('{"env":{"ANTHROPIC_BASE_URL":"http://gateway.test"}}') + config.write_text('model_provider = "litellm"\n') + original = settings.read_bytes(), config.read_bytes() + selected = ("claude", "codex") if stage.endswith("terminal") else ("claude",) + calls = [] + + def fail(): + raise error + + def which(name): + if stage == "which": + fail() + return str(tmp_path / name) + + def restore(): + if stage == "restore": + fail() + + def handoff(*args): + calls.append("handoff") + if stage == "handoff": + fail() + + def spawn(*args, **kwargs): + calls.append("terminal") + if stage == "first-terminal" or (stage == "second-terminal" and len(calls) == 2): + fail() + + @click.command() + def launch(): + launch_configured_agents( + selected, + started_interactive=True, + platform=platform, + cwd=str(tmp_path), + environ={}, + which=which, + restore_terminal=restore, + hand_off=handoff, + run=spawn, + popen=spawn, + ) + + result = CliRunner().invoke(launch) + failed_agent = "codex" if stage == "second-terminal" else "claude" + assert result.exit_code == 1 + assert f"Error: Could not open `{failed_agent}`." in result.output + assert f"Configuration was saved; start `{failed_agent}` manually." in result.output + assert "Traceback" not in result.output + assert (settings.read_bytes(), config.read_bytes()) == original + assert len(calls) == (2 if stage.endswith("terminal") else 0 if stage in ("which", "restore") else 1) + + +@pytest.mark.parametrize("error", [SystemExit(7), KeyboardInterrupt(), TypeError("programming error")]) +def test_launch_does_not_relabel_exit_interrupt_or_programming_errors(tmp_path, error): + def handoff(*args): + raise error + + with pytest.raises(type(error), match=r"7|programming error|^$"): + launch_configured_agents( + ("claude",), started_interactive=False, which=lambda _: str(tmp_path / "claude"), hand_off=handoff + ) + + +def test_missing_terminal_keeps_configuration_and_names_manual_fallback(tmp_path): + binaries = {"claude": str(tmp_path / "claude"), "codex": str(tmp_path / "codex")} + with pytest.raises(click.ClickException, match="Configuration was saved; start `claude` manually"): + launch_configured_agents( + ("claude", "codex"), + started_interactive=True, + platform="linux", + which=binaries.get, + ) diff --git a/uv.lock b/uv.lock index 88f558221e7..e3c15725b36 100644 --- a/uv.lock +++ b/uv.lock @@ -4390,6 +4390,7 @@ cli = [ { name = "pyyaml" }, { name = "requests" }, { name = "rich" }, + { name = "tomlkit" }, ] extra-proxy = [ { name = "a2a-sdk" }, @@ -4444,6 +4445,7 @@ proxy = [ { name = "rq" }, { name = "soundfile" }, { name = "starlette" }, + { name = "tomlkit" }, { name = "uvicorn" }, { name = "uvloop", marker = "sys_platform != 'win32'" }, { name = "websockets" }, @@ -4669,6 +4671,8 @@ requires-dist = [ { name = "starlette", marker = "extra == 'proxy'", specifier = ">=1.0.1,<2.0" }, { name = "tiktoken", specifier = ">=0.8.0,<1.0" }, { name = "tokenizers", specifier = ">=0.21.0,<1.0" }, + { name = "tomlkit", marker = "extra == 'cli'", specifier = ">=0.13.0,<1.0" }, + { name = "tomlkit", marker = "extra == 'proxy'", specifier = ">=0.13.0,<1.0" }, { name = "uvicorn", marker = "extra == 'proxy'", specifier = ">=0.33.0,<1.0" }, { name = "uvloop", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.22.1,<1.0" }, { name = "websockets", marker = "extra == 'proxy'", specifier = ">=15.0.1,<16.0" },