mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
Merge remote-tracking branch 'origin/main' into litellm_registry_audit_20261024
This commit is contained in:
Devin AI
commit
50ae72aef8
10 changed files with 691 additions and 88 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 <agent>` 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),
|
||||
}
|
||||
)
|
||||
|
|
|
|||
165
tests/integration/configuration/test_openapi_docs_route.py
Normal file
165
tests/integration/configuration/test_openapi_docs_route.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
import time
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, list_value, object_value
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy_process
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm.proxy._lazy_features import LAZY_FEATURES
|
||||
|
||||
DOCS_SECONDS: Final = 120
|
||||
COLD_DOCS: Final = 4
|
||||
CHATS: Final = 8
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
OWNED_PROXY_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
TWO_OWNED_PROXIES_TIMEOUT: Final = pytest.mark.timeout(2 * (2 * graceful_stop_seconds() + 120))
|
||||
LAZY_ROUTES_FLAG: Final = "LITELLM_DISABLE_LAZY_ROUTES"
|
||||
WARMED_FEATURE: Final = next(feature for feature in LAZY_FEATURES if feature.name == "policies")
|
||||
|
||||
|
||||
def _timed_get(candidate: Gateway, path: str) -> tuple[httpx.Response, float]:
|
||||
started: Final = time.perf_counter()
|
||||
response: Final = candidate.request("GET", path)
|
||||
return response, time.perf_counter() - started
|
||||
|
||||
|
||||
def _liveliness_until(done: Future[tuple[httpx.Response, float]], candidate: Gateway) -> Iterator[tuple[int, float]]:
|
||||
while not done.done():
|
||||
response, seconds = _timed_get(candidate, "/health/liveliness")
|
||||
yield response.status_code, seconds
|
||||
|
||||
|
||||
def _schema(response: httpx.Response) -> dict[str, JsonValue]:
|
||||
assert response.status_code == 200, response.text[:500]
|
||||
return JSON_OBJECT.validate_json(response.content)
|
||||
|
||||
|
||||
def _servers(response: httpx.Response) -> list[JsonValue]:
|
||||
return list(list_value(_schema(response).get("servers", [])))
|
||||
|
||||
|
||||
def _chat(candidate: Gateway, model: str, key: str, content: str) -> httpx.Response:
|
||||
return candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": content}], "cache": {"no-cache": True}},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _warmed_feature_operations(response: httpx.Response) -> dict[str, JsonValue]:
|
||||
paths: Final = object_value(_schema(response)["paths"])
|
||||
return {path: operations for path, operations in paths.items() if WARMED_FEATURE.matches(path)}
|
||||
|
||||
|
||||
@OWNED_PROXY_TIMEOUT
|
||||
def test_liveliness_answers_while_a_cold_openapi_schema_builds(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with (
|
||||
owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned,
|
||||
ThreadPoolExecutor(1) as pool,
|
||||
):
|
||||
docs: Final = pool.submit(_timed_get, owned.gateway, "/openapi.json")
|
||||
probes: Final = tuple(_liveliness_until(docs, owned.gateway))
|
||||
docs_response, docs_seconds = docs.result()
|
||||
assert "/v1/chat/completions" in object_value(_schema(docs_response)["paths"])
|
||||
assert probes, "no liveliness probe ran while the docs request was pending"
|
||||
assert {status for status, _ in probes} == {200}, probes
|
||||
slowest: Final = max(seconds for _, seconds in probes)
|
||||
assert slowest < docs_seconds / 3, (
|
||||
f"a liveliness probe waited {slowest:.2f}s of the {docs_seconds:.2f}s cold docs request"
|
||||
)
|
||||
|
||||
|
||||
@OWNED_PROXY_TIMEOUT
|
||||
def test_concurrent_cold_openapi_requests_share_the_schema_later_requests_get(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with (
|
||||
owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned,
|
||||
ThreadPoolExecutor(COLD_DOCS) as pool,
|
||||
):
|
||||
cold: Final = tuple(pool.map(lambda _: owned.gateway.request("GET", "/openapi.json"), range(COLD_DOCS)))
|
||||
warm: Final = owned.gateway.request("GET", "/openapi.json")
|
||||
assert [response.status_code for response in (*cold, warm)] == [200] * (COLD_DOCS + 1)
|
||||
assert {response.content for response in cold} == {warm.content}
|
||||
assert "/v1/chat/completions" in object_value(_schema(warm)["paths"])
|
||||
|
||||
|
||||
@OWNED_PROXY_TIMEOUT
|
||||
def test_chat_completions_overlapping_a_cold_openapi_build_all_succeed(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with (
|
||||
owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
ThreadPoolExecutor(CHATS + 1) as pool,
|
||||
):
|
||||
model: Final = scenario.model()
|
||||
key: Final = scenario.key(models=[model])
|
||||
first_chat: Final = _chat(owned.gateway, model, key, "before the docs build")
|
||||
docs: Final = pool.submit(owned.gateway.request, "GET", "/openapi.json")
|
||||
chats: Final = tuple(
|
||||
pool.map(lambda index: _chat(owned.gateway, model, key, f"overlapping docs build {index}"), range(CHATS))
|
||||
)
|
||||
chats_returned_during_build: Final = not docs.done()
|
||||
docs_response: Final = docs.result()
|
||||
assert first_chat.status_code == 200, first_chat.text[:300]
|
||||
assert chats_returned_during_build, "the chat calls returned only after the cold docs build finished"
|
||||
assert [(response.status_code, response.text[:300]) for response in chats if response.status_code != 200] == []
|
||||
replies: Final = tuple(JSON_OBJECT.validate_json(response.content) for response in chats)
|
||||
assert {reply["object"] for reply in replies} == {"chat.completion"}
|
||||
assert all(list_value(reply["choices"]) for reply in replies), replies
|
||||
assert "/v1/chat/completions" in object_value(_schema(docs_response)["paths"])
|
||||
|
||||
|
||||
@OWNED_PROXY_TIMEOUT
|
||||
def test_docs_filtered_mode_serves_only_llm_routes(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {"DOCS_FILTERED": "True"}, client_timeout=DOCS_SECONDS) as owned:
|
||||
paths: Final = object_value(_schema(owned.gateway.request("GET", "/openapi.json"))["paths"])
|
||||
assert "/v1/chat/completions" in paths
|
||||
assert {"/key/generate", "/user/new", "/team/new"}.isdisjoint(paths), sorted(paths)
|
||||
|
||||
|
||||
@OWNED_PROXY_TIMEOUT
|
||||
def test_each_root_path_prefix_lists_itself_first_in_servers(gateway: Gateway, tmp_path: Path) -> None:
|
||||
prefixes: Final = ("/tenant-a", "/tenant-b", "/tenant-a")
|
||||
with owned_proxy_process(
|
||||
gateway, tmp_path, {"SERVER_ROOT_PATHS": "/tenant-a,/tenant-b"}, client_timeout=DOCS_SECONDS
|
||||
) as owned:
|
||||
served: Final = tuple(
|
||||
(prefix, _servers(owned.gateway.request("GET", f"{prefix}/openapi.json"))) for prefix in prefixes
|
||||
)
|
||||
assert served == tuple((prefix, [{"url": prefix}]) for prefix in prefixes)
|
||||
|
||||
|
||||
@OWNED_PROXY_TIMEOUT
|
||||
def test_a_scalar_server_root_path_is_listed_once_in_servers(gateway: Gateway, tmp_path: Path) -> None:
|
||||
overrides: Final = {"SERVER_ROOT_PATH": "/gateway", "LITELLM_UI_PATH": str(tmp_path / "ui")}
|
||||
with owned_proxy_process(gateway, tmp_path, overrides, client_timeout=DOCS_SECONDS) as owned:
|
||||
servers: Final = _servers(owned.gateway.request("GET", "/gateway/openapi.json"))
|
||||
assert servers == [{"url": "/gateway"}]
|
||||
|
||||
|
||||
@TWO_OWNED_PROXIES_TIMEOUT
|
||||
def test_a_lazy_feature_loaded_during_a_cold_openapi_build_shows_its_live_routes_in_the_next_schema(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {LAZY_ROUTES_FLAG: "true"}, client_timeout=DOCS_SECONDS) as eager:
|
||||
live: Final = _warmed_feature_operations(eager.gateway.request("GET", "/openapi.json"))
|
||||
with (
|
||||
owned_proxy_process(
|
||||
gateway, tmp_path, {}, remove_environment=(LAZY_ROUTES_FLAG,), client_timeout=DOCS_SECONDS
|
||||
) as lazy,
|
||||
ThreadPoolExecutor(1) as pool,
|
||||
):
|
||||
docs: Final = pool.submit(lazy.gateway.request, "GET", "/openapi.json")
|
||||
warmed: Final = lazy.gateway.request("POST", f"/lazy/warm/{WARMED_FEATURE.name}")
|
||||
warmed_during_build: Final = not docs.done()
|
||||
during: Final = _warmed_feature_operations(docs.result())
|
||||
after: Final = _warmed_feature_operations(lazy.gateway.request("GET", "/openapi.json"))
|
||||
assert warmed.status_code == 200, warmed.text
|
||||
assert warmed_during_build, "the feature finished loading only after the cold docs build finished"
|
||||
assert live, "the eager proxy published no route of the warmed feature"
|
||||
assert during != live, "the cold docs build started after the feature had already loaded"
|
||||
assert after == live
|
||||
|
|
@ -691,9 +691,6 @@ def test_scim_group_put_returns_the_updated_members(gateway: Gateway, tmp_path:
|
|||
|
||||
@pytest.mark.timeout(300)
|
||||
def test_scim_admin_group_delete_revokes_the_members_admin_key(gateway: Gateway, tmp_path: Path) -> None:
|
||||
pytest.skip(
|
||||
"BUG: DELETE /scim/v2/Groups/{id} demotes the admin group's members in the DB but their existing keys keep proxy_admin access on /key/generate"
|
||||
)
|
||||
group_name: Final = f"scim-admin-delete-{uuid.uuid4().hex}"
|
||||
config: Final = _owned_scim_config(tmp_path / "scim-admin-group-delete.yaml", {"scim_admin_group": group_name})
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned, owned.gateway.scenario() as scenario:
|
||||
|
|
|
|||
|
|
@ -50,7 +50,7 @@ from integration._support.scim import (
|
|||
)
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
|
||||
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS, TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS
|
||||
|
||||
pytestmark: Final = pytest.mark.shares_worker_state
|
||||
|
||||
|
|
@ -70,6 +70,17 @@ FROM "LiteLLM_AuditLog"
|
|||
WHERE object_id IN (%s, %s) AND action = 'deleted'
|
||||
ORDER BY table_name
|
||||
"""
|
||||
AUDIT_ROWS_NAMING_SQL: Final = """
|
||||
SELECT table_name, object_id, changed_by_api_key
|
||||
FROM "LiteLLM_AuditLog"
|
||||
WHERE object_id = ANY(%s::text[]) AND action = 'deleted'
|
||||
ORDER BY table_name, object_id
|
||||
"""
|
||||
OLDEST_TRANSACTION_BLOCKED_BY_SQL: Final = """
|
||||
SELECT coalesce(max(extract(epoch FROM clock_timestamp() - xact_start)), 0)::float AS age
|
||||
FROM pg_stat_activity
|
||||
WHERE %s::int = ANY(pg_blocking_pids(pid))
|
||||
"""
|
||||
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
|
||||
_PIDS: Final = TypeAdapter(tuple[int, ...])
|
||||
|
||||
|
|
@ -115,15 +126,28 @@ def _deleted_audit_rows(team: str, hashed_key: str) -> list[dict[str, JsonValue]
|
|||
return read_rows(DELETED_AUDIT_ROWS_SQL, (team, hashed_key))
|
||||
|
||||
|
||||
def _audit_rows_naming(object_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(AUDIT_ROWS_NAMING_SQL, ("{" + ",".join(object_ids) + "}",))
|
||||
|
||||
|
||||
def _oldest_transaction_blocked_by(holder: int) -> float:
|
||||
age: Final = read_rows(OLDEST_TRANSACTION_BLOCKED_BY_SQL, (str(holder),))[0]["age"]
|
||||
assert isinstance(age, float), age
|
||||
return age
|
||||
|
||||
|
||||
def _worker_startups(log: Path) -> tuple[tuple[int, ...], int]:
|
||||
text: Final = log.read_text()
|
||||
return _PIDS.validate_python(_STARTED_WORKER.findall(text)), text.count("Application startup complete.")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _fresh_connection(candidate: Gateway) -> Iterator[Gateway]:
|
||||
def _fresh_connection(candidate: Gateway, timeout: httpx.Timeout | float | None = None) -> Iterator[Gateway]:
|
||||
with httpx.Client(
|
||||
base_url=candidate.client.base_url, timeout=candidate.client.timeout, trust_env=False, limits=GATEWAY_LIMITS
|
||||
base_url=candidate.client.base_url,
|
||||
timeout=candidate.client.timeout if timeout is None else timeout,
|
||||
trust_env=False,
|
||||
limits=GATEWAY_LIMITS,
|
||||
) as client:
|
||||
yield Gateway(client, candidate.key, candidate.upstream_url)
|
||||
|
||||
|
|
@ -294,6 +318,82 @@ def test_a_group_delete_revokes_member_keys_and_a_re_created_group_gets_its_serv
|
|||
assert _models_status(gateway, member_key) == 401
|
||||
|
||||
|
||||
@pytest.mark.timeout(OWNED_PROXY_TIMEOUT)
|
||||
def test_a_group_delete_sent_with_the_idps_scim_token_revokes_member_keys_and_audits_them_under_that_token(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with (
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"},
|
||||
remove_environment=("DATABASE_URL_READ_REPLICA",),
|
||||
) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
candidate: Final = owned.gateway
|
||||
scim_token: Final = scenario.key(allowed_routes=["/scim/*"], models=[])
|
||||
users: Final = seed_users(scenario, 3)
|
||||
team: Final = created_team(scenario, create_group(candidate, users))
|
||||
member_keys: Final = (team_key(candidate, team, users[0]), team_key(candidate, team, users[1]))
|
||||
service_key: Final = service_account_key(candidate, scenario, team)
|
||||
|
||||
deleted: Final = delete_group(candidate, team, key=scim_token)
|
||||
assert deleted.status_code == 204, deleted.text
|
||||
|
||||
assert keys_of(team) == frozenset({_hashed(service_key)})
|
||||
for key in member_keys:
|
||||
_assert_revoked_everywhere((candidate,), key)
|
||||
audited: Final = (team, *(_hashed(key) for key in member_keys), _hashed(service_key))
|
||||
rows: Final = eventually(lambda: _audit_rows_naming(audited), lambda found: len(found) >= 3, seconds=30)
|
||||
token: Final = _hashed(scim_token)
|
||||
assert rows == [
|
||||
{"table_name": "LiteLLM_TeamTable", "object_id": team, "changed_by_api_key": token},
|
||||
*(
|
||||
{"table_name": "LiteLLM_VerificationToken", "object_id": hashed, "changed_by_api_key": token}
|
||||
for hashed in sorted(_hashed(key) for key in member_keys)
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.timeout(TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 120)
|
||||
def test_a_group_delete_that_outlives_its_transaction_on_a_locked_member_row_changes_nothing_and_a_retry_lands(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
patient_timeout: Final = TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 90
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
_fresh_connection(gateway, timeout=patient_timeout) as patient,
|
||||
ThreadPoolExecutor(max_workers=1) as pool,
|
||||
):
|
||||
users: Final = seed_users(scenario, 3)
|
||||
team: Final = created_team(scenario, create_group(gateway, users))
|
||||
key: Final = team_key(gateway, team, users[0])
|
||||
|
||||
with held_user_row(users[0]) as holder:
|
||||
delete: Final = pool.submit(delete_group, patient, team)
|
||||
eventually(
|
||||
lambda: (
|
||||
delete.done()
|
||||
or _oldest_transaction_blocked_by(holder) > TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 1
|
||||
),
|
||||
lambda outlived: outlived,
|
||||
seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 30,
|
||||
)
|
||||
failed: Final = delete.result(timeout=patient_timeout)
|
||||
|
||||
assert failed.status_code >= 500, failed.text
|
||||
assert_landed(gateway, team, users)
|
||||
assert keys_of(team) == frozenset({_hashed(key)})
|
||||
assert read_rows(DELETED_TEAM_RECORDS_SQL, (team,)) == []
|
||||
assert read_rows(DELETED_KEY_RECORDS_SQL, (_hashed(key),)) == []
|
||||
assert _models_status(gateway, key) == 200
|
||||
|
||||
retried: Final = delete_group(gateway, team)
|
||||
assert retried.status_code == 204, retried.text
|
||||
assert_gone(gateway, team, users, key)
|
||||
|
||||
|
||||
def test_deleting_an_unknown_group_is_404(gateway: Gateway) -> None:
|
||||
missing: Final = f"missing-{uuid.uuid4().hex}"
|
||||
response: Final = delete_group(gateway, missing)
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ from tests.integration._support.client import (
|
|||
string_value,
|
||||
)
|
||||
from tests.integration._support.database import read_rows, write_rows
|
||||
from tests.integration._support.process import owned_proxy_process
|
||||
from tests.integration._support.process import graceful_stop_seconds, owned_proxy_process
|
||||
|
||||
LARGE_ROSTER: Final = 250
|
||||
POOL_LIMIT: Final = 5
|
||||
|
|
@ -131,6 +131,13 @@ def _roster_user_ids(roster: JsonValue) -> list[str]:
|
|||
return sorted(string_value(object_value(member)["user_id"]) for member in roster)
|
||||
|
||||
|
||||
def _deleted_audit_rows(object_id: str, table_name: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT before_value FROM "LiteLLM_AuditLog" WHERE object_id = %s AND table_name = %s AND action = %s',
|
||||
(object_id, table_name, "deleted"),
|
||||
)
|
||||
|
||||
|
||||
def _waiters_on_lock_held_by(backend_pid: int) -> int:
|
||||
rows: Final = read_rows(WAITERS_ON_HELD_LOCK_SQL, ("%pg_advisory_xact_lock%", str(backend_pid)))
|
||||
waiting: Final = rows[0]["waiting"]
|
||||
|
|
@ -551,6 +558,41 @@ def test_delete_writes_one_audit_row_for_the_team_and_one_per_key(gateway: Gatew
|
|||
]
|
||||
|
||||
|
||||
@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120))
|
||||
def test_two_deletes_racing_on_one_team_write_one_tombstone_and_one_audit_row(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with (
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"},
|
||||
remove_environment=("DATABASE_URL_READ_REPLICA",),
|
||||
) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
ThreadPoolExecutor(max_workers=2) as pool,
|
||||
):
|
||||
user: Final = scenario.user()
|
||||
team: Final = _create_team(scenario)
|
||||
_bulk_member_add(owned.gateway, team, [user])
|
||||
|
||||
with psycopg.connect(os.environ["DATABASE_URL"]) as holder:
|
||||
holder.execute(TAKE_TEAM_LOCK_SQL, (team,))
|
||||
deletes: Final = tuple(
|
||||
_post_in_background(pool, owned.gateway, "/team/delete", {"team_ids": [team]}) for _ in range(2)
|
||||
)
|
||||
eventually(
|
||||
lambda: _waiters_on_lock_held_by(holder.info.backend_pid), lambda waiting: waiting == 2, seconds=30
|
||||
)
|
||||
|
||||
assert [delete.result(timeout=60).status_code for delete in deletes] == [200, 200]
|
||||
tombstones: Final = _tombstones(team)
|
||||
assert [row["team_id"] for row in tombstones] == [team]
|
||||
assert _roster_user_ids(tombstones[0]["members_with_roles"]) == sorted([PROXY_ADMIN, user])
|
||||
|
||||
team_audit_rows: Final = _deleted_audit_rows(team, "LiteLLM_TeamTable")
|
||||
assert len(team_audit_rows) == 1, team_audit_rows
|
||||
assert object_value(team_audit_rows[0]["before_value"])["team_id"] == team
|
||||
|
||||
|
||||
def test_second_delete_of_the_same_team_is_404_with_one_tombstone(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
user: Final = scenario.user()
|
||||
|
|
|
|||
|
|
@ -950,6 +950,45 @@ def test_nonstream_failed_follow_up_bills_round_one_on_the_failure_row(
|
|||
)
|
||||
|
||||
|
||||
NO_CACHE: Final = {"no-cache": True}
|
||||
|
||||
|
||||
def _upstream_rounds(rig: Rig) -> tuple[int, int]:
|
||||
answered: Final = tuple(
|
||||
_has_tool_result(JSON_OBJECT.validate_json(request.body))
|
||||
for request in rig.wire.drain()
|
||||
if request.method == "POST"
|
||||
)
|
||||
return answered.count(False), answered.count(True)
|
||||
|
||||
|
||||
def _billed_rows(call_id: str, call_type: str, billed: float) -> list[dict[str, JsonValue]]:
|
||||
return eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s',
|
||||
(call_id, call_type),
|
||||
),
|
||||
lambda found: sum(float(str(row["spend"])) for row in found) >= billed * (1 - 1e-6),
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("surface", "call_type"), (("responses", "aresponses"), ("chat", "acompletion")))
|
||||
def test_nonstream_follow_up_failing_on_every_retry_bills_each_attempts_first_round_on_one_failure_row(
|
||||
gateway: Gateway, surface: Surface, call_type: str
|
||||
) -> None:
|
||||
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True)
|
||||
with _rig(gateway, surface, billing=billing, prices={**PRICES, "num_retries": 2}) as rig:
|
||||
response: Final = rig.send(_granted_key(rig), [AUTO], cache=NO_CACHE)
|
||||
assert response.status_code == 500 and FOLLOW_UP_OUTAGE in response.text, response.text
|
||||
first_rounds, follow_ups = _upstream_rounds(rig)
|
||||
assert first_rounds == follow_ups > 1, (first_rounds, follow_ups)
|
||||
paid: Final = _price(*(FIRST_ROUND,) * first_rounds)
|
||||
rows: Final = _billed_rows(response.headers["x-litellm-call-id"], call_type, paid)
|
||||
assert [row["status"] for row in rows] == ["failure"], rows
|
||||
assert float(str(rows[0]["spend"])) == pytest.approx(paid), rows
|
||||
|
||||
|
||||
def _tool_rows(call_id: str, expected: int) -> list[dict[str, JsonValue]]:
|
||||
return eventually(
|
||||
lambda: read_rows(
|
||||
|
|
|
|||
|
|
@ -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 = {}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue