diff --git a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md index c79632e9195..e1f606a2eb7 100644 --- a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md +++ b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md @@ -40,6 +40,8 @@ env_http_headers = { "x-litellm-api-key" = "LITELLM_API_KEY" } Replace `https://litellm.example.com` with your proxy URL. A local proxy normally uses `http://localhost:4000/v1`. Launch Codex, sign into ChatGPT if needed, and send a prompt with a model available to your subscription. The selected model must match the name exposed by the gateway. +The `lite` CLI can set up this provider each time it starts Codex, with no `config.toml` edit. Export `LITELLM_PROXY_URL` and `LITELLM_PROXY_API_KEY`, then run `lite codex --subscription`. Run `lite config set codex_subscription true` to make it the default. + For a managed workspace rollout, use OpenAI's [gateway configuration guide](https://learn.chatgpt.com/docs/enterprise/sign-in-with-chatgpt-through-a-gateway). ## Native compatibility and Pi diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 52e36964bc1..4d09e4d1e0f 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -34,7 +34,7 @@ from .claude_settings import ( load_json_or_empty, ) from .cmd_quoting import quote_for_cmd -from .config import claude_subscription_default +from .config import subscription_default from .pi import ( LITELLM_PROXY_API_KEY_ENV, PI_PROVIDER_NAME, @@ -83,9 +83,11 @@ _INSTALL_DOCS: Final[dict[str, str]] = { } _HIDDEN_AGENTS: Final = frozenset({"pi"}) -_SUBSCRIPTION_AGENTS: Final = frozenset({"claude"}) +_SUBSCRIPTION_LOGINS: Final = MappingProxyType({"claude": "claude.ai", "codex": "ChatGPT"}) +_SUBSCRIPTION_AGENTS: Final = frozenset(_SUBSCRIPTION_LOGINS) CODEX_PROXY_PROVIDER: Final = "litellm" +CODEX_SUBSCRIPTION_PROVIDER: Final = "lite_subscription" CODEX_HOME_ENV: Final = "CODEX_HOME" CODEX_MODEL_CATALOG_FILENAME: Final = "litellm-models.json" _CODEX_PREFLIGHT_TIMEOUT_SECONDS: Final = 10.0 @@ -143,6 +145,10 @@ def build_agent_env( travels in the x-litellm-api-key header through ANTHROPIC_CUSTOM_HEADERS, and no ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY is left in the environment, since either one makes Claude Code stop using the subscription (and its connectors). + With subscription=True the OpenAI profile keeps Codex on its ChatGPT login: the key goes to + LITELLM_PROXY_API_KEY, which Codex sends in the x-litellm-api-key header, and OPENAI_API_KEY and + OPENAI_BASE_URL are left exactly as the caller had them, so the key is never offered to Codex as an + OpenAI credential and a child process never pairs the caller's own OpenAI key with the proxy URL. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") @@ -161,21 +167,24 @@ def build_agent_env( if ENABLE_GATEWAY_MODEL_DISCOVERY_ENV not in env: env[ENABLE_GATEWAY_MODEL_DISCOVERY_ENV] = ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE if PROFILE_OPENAI in profiles: - env[OPENAI_BASE_URL_ENV] = root + "/v1" - env[OPENAI_API_KEY_ENV] = api_key + if subscription: + env[LITELLM_PROXY_API_KEY_ENV] = api_key + else: + env[OPENAI_BASE_URL_ENV] = root + "/v1" + env[OPENAI_API_KEY_ENV] = api_key if PROFILE_LITELLM in profiles: env[LITELLM_PROXY_API_KEY_ENV] = api_key return env -def codex_proxy_provider(base_url: str) -> Mapping[str, str | bool]: +def codex_proxy_provider(base_url: str, *, subscription: bool = False) -> Mapping[str, str | bool]: return MappingProxyType( { "name": "LiteLLM proxy", "base_url": base_url.rstrip("/") + "/v1", "wire_api": "responses", "supports_websockets": False, - "requires_openai_auth": False, + "requires_openai_auth": subscription, } ) @@ -304,7 +313,7 @@ def subscription_launch_args( return (_SETTINGS_FLAG, str(overlay)), rest -def _codex_proxy_args(base_url: str) -> list[str]: +def _codex_proxy_args(base_url: str, subscription: bool) -> list[str]: """Codex `-c` overrides that point it at the proxy. Codex ignores OPENAI_BASE_URL (it always dials api.openai.com), so the env @@ -312,24 +321,35 @@ def _codex_proxy_args(base_url: str) -> list[str]: inline; supports_websockets=false forces the HTTP/SSE Responses transport because the proxy does not speak the Responses WebSocket protocol. The key is read from OPENAI_API_KEY, which build_agent_env already exports. + With subscription=True the provider is `lite_subscription` with requires_openai_auth, so Codex + signs in with its own ChatGPT login, and the key travels in the x-litellm-api-key header from + LITELLM_PROXY_API_KEY. It sets no env_key, which would win over the login. A `-c` override cannot + remove keys, so the id is one no config.toml defines: `lite configure codex` writes `litellm`, and + the cookbook tells users to write `litellm_subscription`, whose env_key or headers would otherwise apply. """ - provider: Final = f"model_providers.{CODEX_PROXY_PROVIDER}" + provider_id: Final = CODEX_SUBSCRIPTION_PROVIDER if subscription else CODEX_PROXY_PROVIDER + provider: Final = f"model_providers.{provider_id}" + credential: Final = ( + f"{provider}.env_http_headers={{{json.dumps(LITELLM_KEY_HEADER)}={json.dumps(LITELLM_PROXY_API_KEY_ENV)}}}" + if subscription + else f'{provider}.env_key="{OPENAI_API_KEY_ENV}"' + ) return [ "-c", - f'model_provider="{CODEX_PROXY_PROVIDER}"', + f'model_provider="{provider_id}"', *( argument - for key, value in codex_proxy_provider(base_url).items() + for key, value in codex_proxy_provider(base_url, subscription=subscription).items() for argument in ("-c", f"{provider}.{key}={json.dumps(value)}") ), "-c", - f'{provider}.env_key="{OPENAI_API_KEY_ENV}"', + credential, "-c", f"{provider}.http_headers={{}}", ] -_PROXY_ARGS: Final[dict[str, Callable[[str], list[str]]]] = { +_PROXY_ARGS: Final[dict[str, Callable[[str, bool], list[str]]]] = { "codex": _codex_proxy_args, } @@ -410,14 +430,14 @@ _Preparer: TypeAlias = Callable[[str, str, Mapping[str, str]], Sequence[str]] _PREPARERS: Final[Mapping[str, _Preparer]] = MappingProxyType({"pi": prepare_pi, "codex": prepare_codex}) -def agent_launch_args(command: str, base_url: str) -> list[str]: +def agent_launch_args(command: str, base_url: str, *, subscription: bool = False) -> list[str]: """Extra CLI args an agent needs to actually honor the proxy. Claude Code and OpenCode respect the exported env vars, so they get nothing here; Codex needs its provider pointed via config overrides. """ builder: Final = _PROXY_ARGS.get(os.path.basename(command)) - return builder(base_url) if builder else [] + return builder(base_url, subscription) if builder else [] class ListedModel(LiteLLMBaseModel): @@ -889,11 +909,13 @@ def run_agent( (verify_subscription if subscription else verify)(base_url, api_key) env_before_sync: Final = base_env if base_env is not None else os.environ - synced: Final = sync_models(binary, env_before_sync, base_url, api_key, skip_verify) + synced: Final = ( + _NO_EXTRA_ENV if subscription else sync_models(binary, env_before_sync, base_url, api_key, skip_verify) + ) if isinstance(synced, ModelSyncSkipped): warn(f"litellm: not syncing {display_name} models from the proxy: {synced.reason}") - prepare: Final = preparers.get(os.path.basename(command[0])) + prepare: Final = None if subscription else preparers.get(os.path.basename(command[0])) prepared_args: Final = tuple(prepare(base_url, api_key, env_before_sync)) if prepare is not None else () env: Final = MappingProxyType( @@ -903,7 +925,11 @@ def run_agent( } ) synced_args: Final = synced.args if isinstance(synced, ModelSyncArgs) else () - extra_args: Final = (*agent_launch_args(command[0], base_url), *synced_args, *prepared_args) + extra_args: Final = ( + *agent_launch_args(command[0], base_url, subscription=subscription), + *synced_args, + *prepared_args, + ) settings_args, forwarded = ( subscription_launch_args( command, @@ -914,7 +940,7 @@ def run_agent( ), write_overlay, ) - if subscription and os.path.basename(command[0]) in _SUBSCRIPTION_AGENTS + if subscription and os.path.basename(command[0]) == "claude" else ((), tuple(command[1:])) ) if reattach_terminal is not None: @@ -948,10 +974,19 @@ def resolve_api_key(ctx: click.Context) -> str: _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy." -_SUBSCRIPTION_HELP: Final = ( - "Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the " - "x-litellm-api-key header instead of replacing Claude Code's own credential. " - "Make it the default with `lite config set claude_subscription true`." +_SUBSCRIPTION_HELP: Final = MappingProxyType( + { + "claude": ( + "Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the " + "x-litellm-api-key header instead of replacing Claude Code's own credential. " + "Make it the default with `lite config set claude_subscription true`." + ), + "codex": ( + "Keep using your ChatGPT login. The LiteLLM key is sent in the " + "x-litellm-api-key header instead of replacing Codex's own credential. " + "Make it the default with `lite config set codex_subscription true`." + ), + } ) @@ -964,10 +999,8 @@ def _launch( api_key: Final = resolve_api_key(ctx) display_name, _profiles = agent_profile(binary) - use_subscription: Final = ( - subscription if subscription is not None else binary in _SUBSCRIPTION_AGENTS and claude_subscription_default() - ) - login_note: Final = " (keeping your claude.ai login)" if use_subscription else "" + use_subscription: Final = subscription if subscription is not None else subscription_default(binary) + login_note: Final = f" (keeping your {_SUBSCRIPTION_LOGINS[binary]} login)" if use_subscription else "" click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}{login_note}") try: @@ -989,7 +1022,7 @@ _Decorated: Final = TypeVar("_Decorated", bound=Callable[..., object] | click.Co def _subscription_option(binary: str) -> Callable[[_Decorated], _Decorated]: if binary not in _SUBSCRIPTION_AGENTS: return lambda func: func - return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP) + return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP[binary]) def _make_agent_command(binary: str, display_name: str) -> click.Command: diff --git a/litellm/proxy/client/cli/commands/config.py b/litellm/proxy/client/cli/commands/config.py index 539d88dce4a..1689a08c3de 100644 --- a/litellm/proxy/client/cli/commands/config.py +++ b/litellm/proxy/client/cli/commands/config.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.private_json import ensure_private_dir, write_pr HIDDEN_COMMANDS_KEY: Final = "hidden_commands" CLAUDE_SUBSCRIPTION_KEY: Final = "claude_subscription" +CODEX_SUBSCRIPTION_KEY: Final = "codex_subscription" PKCE_KEY: Final = "pkce" _config_adapter: Final[TypeAdapter[Mapping[str, str]]] = TypeAdapter(Mapping[str, str]) @@ -114,9 +115,15 @@ def _normalize_boolean(key: str) -> Callable[[str], str]: return normalize -def claude_subscription_default() -> bool: - """Whether `lite claude` keeps the claude.ai login unless told otherwise.""" - return get_config_value(CLAUDE_SUBSCRIPTION_KEY) == "true" +_SUBSCRIPTION_KEYS: Final[Mapping[str, str]] = MappingProxyType( + {"claude": CLAUDE_SUBSCRIPTION_KEY, "codex": CODEX_SUBSCRIPTION_KEY} +) + + +def subscription_default(agent: str) -> bool: + """Whether `lite ` keeps the agent's own subscription login unless told otherwise.""" + key: Final = _SUBSCRIPTION_KEYS.get(agent) + return key is not None and get_config_value(key) == "true" def pkce_default() -> bool: @@ -129,6 +136,7 @@ _NORMALIZERS: Final[Mapping[str, Callable[[str], str]]] = MappingProxyType( "base_url": normalize_base_url, HIDDEN_COMMANDS_KEY: _normalize_hidden_commands, CLAUDE_SUBSCRIPTION_KEY: _normalize_boolean(CLAUDE_SUBSCRIPTION_KEY), + CODEX_SUBSCRIPTION_KEY: _normalize_boolean(CODEX_SUBSCRIPTION_KEY), PKCE_KEY: _normalize_boolean(PKCE_KEY), } ) diff --git a/tests/unit/proxy/client/cli/test_agents.py b/tests/unit/proxy/client/cli/test_agents.py index 3481346076f..3c79b916202 100644 --- a/tests/unit/proxy/client/cli/test_agents.py +++ b/tests/unit/proxy/client/cli/test_agents.py @@ -1,10 +1,13 @@ +import functools import inspect import json import os import stat import subprocess import sys +from collections.abc import Mapping, Sequence from pathlib import Path +from typing import Final from unittest.mock import patch import click @@ -31,8 +34,14 @@ from litellm.proxy.client.cli.commands.agents import ( verify_subscription_key, ) from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError +from litellm.proxy.client.cli.commands.config import save_config from litellm.proxy.client.cli.main import cli as lite +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents" CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" @@ -189,6 +198,30 @@ def _overlay(args): return json.loads(Path(args[2]).read_text()) +def _codex_config(args): + assert set(args[0::2]) == {"-c"} + return tomllib.loads("\n".join(args[1::2])) + + +def _set_codex_key(config: Mapping[str, object], keys: Sequence[str], value: object) -> Mapping[str, object]: + head, *rest = keys + if not rest: + return {**config, head: value} + nested: Final = config.get(head, {}) + assert isinstance(nested, Mapping) + return {**config, head: _set_codex_key(nested, rest, value)} + + +def _with_codex_override(config: Mapping[str, object], override: str) -> Mapping[str, object]: + path, value = override.split("=", 1) + return _set_codex_key(config, path.split("."), tomllib.loads(f"value = {value}")["value"]) + + +def _codex_config_over(user_config_toml: str, args: Sequence[str]) -> Mapping[str, object]: + assert set(args[0::2]) == {"-c"} + return functools.reduce(_with_codex_override, args[1::2], tomllib.loads(user_config_toml)) + + class TestAgentProfile: def test_claude_is_anthropic(self): name, profiles = agent_profile("claude") @@ -254,6 +287,7 @@ class TestBuildAgentEnv: env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"})) assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" assert env["OPENAI_API_KEY"] == "sk-key" + assert "LITELLM_PROXY_API_KEY" not in env assert "ANTHROPIC_BASE_URL" not in env assert "ENABLE_TOOL_SEARCH" not in env assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in env @@ -317,11 +351,21 @@ class TestBuildAgentEnv: ) assert env["ANTHROPIC_CUSTOM_HEADERS"].splitlines() == ["x-team: blue", "x-litellm-api-key: Bearer sk-new"] - def test_subscription_leaves_openai_agents_on_the_bearer_key(self): - env = build_agent_env({}, "http://localhost:4000", "sk-key", frozenset({"openai"}), subscription=True) - assert env["OPENAI_API_KEY"] == "sk-key" + def test_subscription_routes_openai_agents_without_offering_the_key_as_an_openai_credential(self): + env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"}), subscription=True) + assert "OPENAI_BASE_URL" not in env + assert env["LITELLM_PROXY_API_KEY"] == "sk-key" + assert "OPENAI_API_KEY" not in env assert "ANTHROPIC_CUSTOM_HEADERS" not in env + def test_subscription_leaves_the_users_own_openai_key_and_base_url_as_they_were(self): + base = {"OPENAI_API_KEY": "sk-users-openai", "OPENAI_BASE_URL": "https://api.openai.com/v1"} + env = build_agent_env(base, "http://localhost:4000", "sk-key", frozenset({"openai"}), subscription=True) + assert env["OPENAI_API_KEY"] == "sk-users-openai" + assert env["OPENAI_BASE_URL"] == "https://api.openai.com/v1" + assert env["LITELLM_PROXY_API_KEY"] == "sk-key" + assert base == {"OPENAI_API_KEY": "sk-users-openai", "OPENAI_BASE_URL": "https://api.openai.com/v1"} + def test_default_mode_does_not_touch_custom_headers(self): env = build_agent_env( {"ANTHROPIC_CUSTOM_HEADERS": "x-team: blue"}, "http://localhost:4000", "sk-key", frozenset({"anthropic"}) @@ -354,6 +398,67 @@ class TestAgentLaunchArgs: agent_launch_args("codex", "http://localhost:4000") ) + def test_codex_default_provider_reads_the_key_from_openai_api_key(self): + assert _codex_config(agent_launch_args("codex", "http://localhost:4000/")) == { + "model_provider": "litellm", + "model_providers": { + "litellm": { + "name": "LiteLLM proxy", + "base_url": "http://localhost:4000/v1", + "wire_api": "responses", + "supports_websockets": False, + "requires_openai_auth": False, + "env_key": "OPENAI_API_KEY", + "http_headers": {}, + } + }, + } + + def test_codex_subscription_provider_keeps_the_chatgpt_login_and_sends_the_key_in_the_litellm_header(self): + assert _codex_config( + agent_launch_args("/usr/local/bin/codex", "http://localhost:4000/", subscription=True) + ) == { + "model_provider": "lite_subscription", + "model_providers": { + "lite_subscription": { + "name": "LiteLLM proxy", + "base_url": "http://localhost:4000/v1", + "wire_api": "responses", + "supports_websockets": False, + "requires_openai_auth": True, + "env_http_headers": {"x-litellm-api-key": "LITELLM_PROXY_API_KEY"}, + "http_headers": {}, + } + }, + } + + def test_codex_subscription_leaves_the_cookbook_provider_table_out_of_the_launch(self): + user_config = """ +model_provider = "litellm_subscription" + +[model_providers.litellm_subscription] +name = "LiteLLM subscription" +base_url = "https://litellm.example.com/v1" +wire_api = "responses" +requires_openai_auth = true +env_key = "OPENAI_API_KEY" +http_headers = { "x-team" = "qa" } +""" + config = _codex_config_over( + user_config, agent_launch_args("codex", "http://localhost:4000/", subscription=True) + ) + launched = config["model_providers"][config["model_provider"]] + assert "env_key" not in launched + assert launched["requires_openai_auth"] is True + assert launched["base_url"] == "http://localhost:4000/v1" + assert config["model_providers"]["litellm_subscription"] == tomllib.loads(user_config)["model_providers"][ + "litellm_subscription" + ] + + @pytest.mark.parametrize("agent", ["claude", "opencode", "pi", "mytool"]) + def test_subscription_adds_no_args_for_agents_without_a_proxy_provider(self, agent): + assert agent_launch_args(agent, "http://localhost:4000", subscription=True) == [] + def test_pi_gets_no_static_args(self): assert agent_launch_args("pi", "http://localhost:4000") == [] @@ -1286,6 +1391,25 @@ class TestRunAgent: with pytest.raises(AgentRunError): run_agent("http://localhost:4000", "sk-key", []) + @pytest.mark.parametrize("subscription", [True, False]) + def test_codex_stop_hook_is_added_only_when_codex_sends_the_key_as_its_openai_key(self, subscription): + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex"], + subscription=subscription, + base_env={"OPENAI_API_KEY": "sk-users-own-openai-key"}, + which=lambda name: "/usr/local/bin/codex", + verify=lambda *a: None, + verify_subscription=lambda *a: None, + sync_models=lambda *a: ModelSyncSkipped("not needed"), + warn=lambda message: None, + launcher=lambda p, a, e: launched.update(args=tuple(a)), + preparers={"codex": lambda *a: ("-c", "hooks.Stop=lite")}, + ) + assert ("hooks.Stop=lite" in launched["args"]) is not subscription + def test_subscription_launch_probes_with_the_header_and_exports_it(self): probed = [] launched = {} @@ -1574,6 +1698,56 @@ class TestRunAgent: read_settings=lambda path: {}, ) + def test_codex_subscription_launch_keeps_the_chatgpt_login_and_needs_no_settings_overlay(self): + probed = [] + overlays = [] + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex", "exec", "--settings", "kept"], + subscription=True, + base_env={"PATH": "/usr/bin"}, + which=lambda name: "/usr/local/bin/codex", + verify=lambda *a: probed.append(("bearer", a)), + verify_subscription=lambda *a: probed.append(("header", a)), + sync_models=lambda *a: {}, + preparers={}, + write_overlay=lambda text, environ: overlays.append(text), + launcher=lambda p, a, e: launched.update(path=p, args=tuple(a), env=dict(e)), + ) + assert probed == [("header", ("http://localhost:4000", "sk-key"))] + assert overlays == [] + assert launched["path"] == "/usr/local/bin/codex" + assert launched["args"][0] == "codex" + assert launched["args"][-3:] == ("exec", "--settings", "kept") + assert _codex_config(launched["args"][1:-3])["model_provider"] == "lite_subscription" + assert "OPENAI_BASE_URL" not in launched["env"] + assert launched["env"]["LITELLM_PROXY_API_KEY"] == "sk-key" + assert "OPENAI_API_KEY" not in launched["env"] + + @pytest.mark.parametrize( + "sync_result", + [ModelSyncArgs(("-c", 'model_catalog_json="/tmp/c.json"')), ModelSyncSkipped("lists no chat models")], + ) + def test_codex_subscription_launch_leaves_the_model_list_to_codex(self, sync_result): + warnings = [] + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex"], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/codex", + verify_subscription=lambda *a: None, + sync_models=lambda *a: sync_result, + warn=warnings.append, + launcher=lambda p, a, e: launched.update(args=tuple(a)), + ) + assert warnings == [] + assert "model_catalog_json" not in json.dumps(launched["args"]) + def test_default_launch_still_probes_with_the_bearer(self): probed = [] run_agent( @@ -1964,7 +2138,11 @@ class TestAgentCommands: assert captured["command"] == ["opencode"] assert "routing OpenCode through proxy at http://localhost:4000" in result.output - def test_claude_subscription_flag_reaches_run_agent_and_is_not_forwarded(self): + @pytest.mark.parametrize( + ("agent", "user_args", "login"), + [("claude", ["--resume"], "claude.ai"), ("codex", ["exec", "hi"], "ChatGPT")], + ) + def test_subscription_flag_reaches_run_agent_and_is_not_forwarded(self, agent, user_args, login): captured = {} def fake_run_agent(base_url, api_key, command, **kwargs): @@ -1973,61 +2151,85 @@ class TestAgentCommands: with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent): result = self.runner.invoke( - _agent_command("claude"), - ["--subscription", "--resume"], + _agent_command(agent), + ["--subscription", *user_args], obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, ) assert result.exit_code == 0, result.output assert captured["subscription"] is True - assert captured["command"] == ["claude", "--resume"] - assert "keeping your claude.ai login" in result.output + assert captured["command"] == [agent, *user_args] + assert f"keeping your {login} login" in result.output + @pytest.mark.parametrize(("agent", "login"), [("claude", "claude.ai"), ("codex", "ChatGPT")]) @pytest.mark.parametrize( - ("configured", "args", "expected"), + ("saved", "args", "expected"), [ - (False, [], False), - (True, [], True), - (True, ["--no-subscription"], False), - (False, ["--subscription"], True), + (None, [], False), + ("false", [], False), + ("true", [], True), + ("true", ["--no-subscription"], False), + ("false", ["--subscription"], True), ], ) - def test_subscription_comes_from_the_flag_else_the_saved_setting(self, configured, args, expected): + def test_subscription_comes_from_the_flag_else_the_saved_setting(self, agent, login, saved, args, expected): + if saved is not None: + save_config({f"{agent}_subscription": saved}) captured = {} - with ( - patch(f"{AGENTS_MODULE}.claude_subscription_default", return_value=configured), - patch( - f"{AGENTS_MODULE}.run_agent", - side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription"), command=list(c)), - ), + with patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription"), command=list(c)), ): result = self.runner.invoke( - _agent_command("claude"), args, obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} + _agent_command(agent), args, obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} ) assert result.exit_code == 0, result.output assert captured["sub"] is expected - assert captured["command"] == ["claude"] - assert ("keeping your claude.ai login" in result.output) is expected + assert captured["command"] == [agent] + assert (f"keeping your {login} login" in result.output) is expected - @pytest.mark.parametrize("agent", ["codex", "opencode"]) - def test_the_saved_claude_setting_does_not_change_other_agents(self, agent): + @pytest.mark.parametrize( + ("saved_key", "agent"), + [ + ("claude_subscription", "codex"), + ("claude_subscription", "opencode"), + ("codex_subscription", "claude"), + ("codex_subscription", "opencode"), + ], + ) + def test_a_saved_subscription_setting_only_changes_its_own_agent(self, saved_key, agent): + save_config({saved_key: "true"}) captured = {} - with ( - patch(f"{AGENTS_MODULE}.claude_subscription_default", return_value=True), - patch( - f"{AGENTS_MODULE}.run_agent", - side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription")), - ), + with patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription")), ): result = self.runner.invoke( _agent_command(agent), [], obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} ) assert result.exit_code == 0, result.output assert captured["sub"] is False + assert "keeping your" not in result.output - @pytest.mark.parametrize("agent", ["codex", "opencode", "pi"]) - def test_only_claude_offers_the_subscription_flag(self, agent): - assert "--subscription" not in {opt for param in _agent_command(agent).params for opt in param.opts} + def test_only_claude_and_codex_offer_the_subscription_flag(self): + offering = { + command.name + for command in agent_commands() + if "--subscription" in {opt for param in command.params for opt in param.opts} + } + assert offering == {"claude", "codex"} + + @pytest.mark.parametrize( + ("agent", "login", "own_credential"), + [("claude", "claude.ai", "Claude Code's"), ("codex", "ChatGPT", "Codex's")], + ) + def test_subscription_help_names_the_agents_own_login_and_its_setting(self, agent, login, own_credential): + result = self.runner.invoke(lite, [agent, "--help"], prog_name="lite") + assert result.exit_code == 0, result.output + text = " ".join(result.output.split()) + assert f"Keep using your {login} login" in text + assert f"instead of replacing {own_credential} own credential" in text + assert f"lite config set {agent}_subscription true" in text def test_skip_verify_is_consumed_not_forwarded(self): captured = {} diff --git a/tests/unit/proxy/client/cli/test_config_commands.py b/tests/unit/proxy/client/cli/test_config_commands.py index 2cd7a76ab95..e99a4311caf 100644 --- a/tests/unit/proxy/client/cli/test_config_commands.py +++ b/tests/unit/proxy/client/cli/test_config_commands.py @@ -9,12 +9,12 @@ from click.testing import CliRunner from litellm.litellm_core_utils.private_json import write_private_json from litellm.proxy.client.cli import cli from litellm.proxy.client.cli.commands.config import ( - claude_subscription_default, get_config_file_path, get_config_value, load_config, pkce_default, save_config, + subscription_default, ) from litellm.proxy.client.cli.interface import show_commands @@ -362,33 +362,48 @@ class TestWritePrivateJson: assert list(tmp_path.glob(".tmp-*")) == [] -class TestClaudeSubscriptionSetting: +@pytest.mark.parametrize(("agent", "key"), [("claude", "claude_subscription"), ("codex", "codex_subscription")]) +class TestSubscriptionSetting: @pytest.mark.parametrize("word", ["true", "TRUE", "on", "1", "yes"]) - def test_truthy_words_are_stored_as_true(self, cli_runner, isolated_home, word): - result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", word]) + def test_truthy_words_are_stored_as_true(self, cli_runner, isolated_home, agent, key, word): + result = cli_runner.invoke(cli, ["config", "set", key, word]) assert result.exit_code == 0, result.output - assert load_config()["claude_subscription"] == "true" - assert claude_subscription_default() is True + assert load_config()[key] == "true" + assert subscription_default(agent) is True @pytest.mark.parametrize("word", ["false", "Off", "0", "no"]) - def test_falsy_words_are_stored_as_false(self, cli_runner, isolated_home, word): - cli_runner.invoke(cli, ["config", "set", "claude_subscription", "true"]) - result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", word]) + def test_falsy_words_are_stored_as_false(self, cli_runner, isolated_home, agent, key, word): + cli_runner.invoke(cli, ["config", "set", key, "true"]) + result = cli_runner.invoke(cli, ["config", "set", key, word]) assert result.exit_code == 0, result.output - assert load_config()["claude_subscription"] == "false" - assert claude_subscription_default() is False + assert load_config()[key] == "false" + assert subscription_default(agent) is False - def test_anything_else_is_rejected_and_nothing_is_saved(self, cli_runner, isolated_home): - result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", "maybe"]) + def test_anything_else_is_rejected_and_nothing_is_saved(self, cli_runner, isolated_home, agent, key): + result = cli_runner.invoke(cli, ["config", "set", key, "maybe"]) assert result.exit_code != 0 - assert "true or false" in result.output - assert "claude_subscription" not in load_config() + assert f"{key} must be true or false" in result.output + assert key not in load_config() - def test_unset_means_off(self, isolated_home): - assert claude_subscription_default() is False + def test_unset_means_off(self, cli_runner, isolated_home, agent, key): + cli_runner.invoke(cli, ["config", "set", key, "true"]) + result = cli_runner.invoke(cli, ["config", "unset", key]) + + assert result.exit_code == 0, result.output + assert subscription_default(agent) is False + + def test_turning_it_on_leaves_every_other_agent_off(self, cli_runner, isolated_home, agent, key): + cli_runner.invoke(cli, ["config", "set", key, "true"]) + + assert {name: subscription_default(name) for name in ("claude", "codex", "opencode", "pi")} == { + "claude": agent == "claude", + "codex": agent == "codex", + "opencode": False, + "pi": False, + } class TestPkceSetting: