From 4846d9c2ed38795b931cd397e70705e45090b280 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 27 Oct 2026 20:24:56 +0000 Subject: [PATCH 01/83] fix(cost-map): add 2026-11-06 deprecation_date to two together_ai rows (#48651) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 ++ model_prices_and_context_window.json | 2 ++ 2 files changed, 4 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7cd4eadca6e..39500330096 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -47102,6 +47102,7 @@ }, "together_ai/MiniMaxAI/MiniMax-M3": { "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-11-06", "input_cost_per_token": 3e-07, "litellm_provider": "together_ai", "max_input_tokens": 524288, @@ -47295,6 +47296,7 @@ }, "together_ai/thinkingmachines/Inkling": { "cache_read_input_token_cost": 1.7e-07, + "deprecation_date": "2026-11-06", "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", "max_input_tokens": 524288, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7cd4eadca6e..39500330096 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -47102,6 +47102,7 @@ }, "together_ai/MiniMaxAI/MiniMax-M3": { "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-11-06", "input_cost_per_token": 3e-07, "litellm_provider": "together_ai", "max_input_tokens": 524288, @@ -47295,6 +47296,7 @@ }, "together_ai/thinkingmachines/Inkling": { "cache_read_input_token_cost": 1.7e-07, + "deprecation_date": "2026-11-06", "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", "max_input_tokens": 524288, From d6edd37810944752561ed49d9a8e85085f8dce81 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 27 Oct 2026 13:33:03 -0700 Subject: [PATCH 02/83] fix(aiohttp_transport): close an owned aiohttp session when its transport is collected (#48650) An OpenAI or Azure SDK client litellm built and evicted from the client cache is usually garbage collected inside the closer's 900 second grace window, so the closer's weak reference dies and the client's aiohttp session is never closed. LiteLLMAiohttpTransport.__del__ now hands an owned, still-open ClientSession to _close_recycled_session, which already covers the current-loop, foreign-loop and closed-loop cases. Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: BangerSoul --- .../llms/custom_httpx/aiohttp_transport.py | 7 ++ .../sdk/test_aiohttp_session_rebuild_wire.py | 86 ++++++++++++++++++- .../unit/caching/test_llm_caching_handler.py | 33 +++++++ .../custom_httpx/test_aiohttp_transport.py | 73 ++++++++++++++++ 4 files changed, 197 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index b8cb183c02f..0e4d97fd770 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -289,6 +289,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport): cls._background_close_tasks.add(task) task.add_done_callback(cls._on_close_task_done) + def __del__(self) -> None: + session: Final = getattr(self, "client", None) + if not getattr(self, "_owns_session", False) or not isinstance(session, ClientSession): + return + with contextlib.suppress(Exception): + self._close_recycled_session(session) + def get_valid_client_session(self) -> ClientSession: """ Helper to get a valid ClientSession for the current event loop. diff --git a/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py b/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py index 1ec1f486bed..f1a18ba389a 100644 --- a/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py +++ b/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py @@ -11,11 +11,30 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Final import pytest -from pydantic import JsonValue, TypeAdapter +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter CONFIGURED_KEEPALIVE_SECONDS: Final = 1 IDLE_SECONDS: Final = 2 +EVICTED_CLIENTS: Final = 30 RESPONSES: Final = TypeAdapter(list[dict[str, JsonValue]]) +CHAT_COMPLETION: Final = { + "id": "chatcmpl-scripted", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5.4-nano", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class EvictionReport(BaseModel): + model_config = ConfigDict(frozen=True) + + cache_size: int + cached_clients: int + established_connections: int + unclosed_messages: tuple[str, ...] + REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent( """ @@ -38,6 +57,47 @@ REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent( ) +EVICTED_CLIENTS_EXCHANGE: Final = textwrap.dedent( + """ + import asyncio, gc, json, sys + import psutil + import litellm + from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport + + unclosed = [] + + def established(port: int) -> int: + connections = psutil.Process().net_connections("tcp") + return sum(1 for c in connections if c.status == "ESTABLISHED" and c.raddr and c.raddr.port == port) + + async def one_call(base_url: str, index: int) -> None: + await litellm.acompletion( + model="openai/gpt-5.4-nano", + api_key=f"sk-distinct-{index}", + api_base=base_url, + messages=[{"role": "user", "content": "hi"}], + ) + + async def main(base_url: str, port: int, evicted: int) -> None: + asyncio.get_running_loop().set_exception_handler(lambda loop, ctx: unclosed.append(ctx["message"])) + cache = litellm.in_memory_llm_clients_cache + for index in range(cache.max_size_in_memory + evicted): + await one_call(base_url, index) + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + await asyncio.sleep(0) + print(json.dumps({ + "cache_size": cache.max_size_in_memory, + "cached_clients": len(cache.cache_dict), + "established_connections": established(port), + "unclosed_messages": unclosed, + })) + + asyncio.run(main(sys.argv[1], int(sys.argv[2]), int(sys.argv[3]))) + """ +) + + class _ConnectionCountingPeer(ThreadingHTTPServer): daemon_threads = True @@ -62,7 +122,8 @@ class _ConnectionHandler(BaseHTTPRequestHandler): def do_POST(self) -> None: self.rfile.read(int(self.headers["Content-Length"])) - body: Final = json.dumps({"connection": self.connection_number}).encode() + payload: Final = CHAT_COMPLETION if self.path.endswith("/chat/completions") else {"connection": self.connection_number} + body: Final = json.dumps(payload).encode() self.send_response(200) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(body))) @@ -107,3 +168,24 @@ def test_rebuilt_shared_session_drops_idle_connection_after_configured_keepalive ) -> None: observed: Final = _rebuilt_session_exchange(connection_counting_peer) assert observed == [{"connection": 1}, {"connection": 2}], observed + + +def _evicted_clients_exchange(base_url: str) -> EvictionReport: + completed: Final = subprocess.run( + [sys.executable, "-P", "-c", EVICTED_CLIENTS_EXCHANGE, base_url, base_url.rsplit(":", 1)[1], str(EVICTED_CLIENTS)], + capture_output=True, + text=True, + timeout=300, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return EvictionReport.model_validate_json(completed.stdout) + + +def test_sdk_clients_evicted_from_a_full_client_cache_leave_no_open_session_or_connection( + connection_counting_peer: str, +) -> None: + report: Final = _evicted_clients_exchange(connection_counting_peer) + assert report.cached_clients == report.cache_size, report + assert report.unclosed_messages == (), report + assert report.established_connections <= report.cache_size, report diff --git a/tests/unit/caching/test_llm_caching_handler.py b/tests/unit/caching/test_llm_caching_handler.py index 0c0fe8ca3c5..a3efd5d97aa 100644 --- a/tests/unit/caching/test_llm_caching_handler.py +++ b/tests/unit/caching/test_llm_caching_handler.py @@ -11,11 +11,14 @@ See: https://github.com/BerriAI/litellm/pull/22247 import asyncio, gc, httpx, importlib, litellm, os, threading, time, weakref import warnings +import aiohttp import pytest +from openai import AsyncOpenAI from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import( @@ -204,6 +207,36 @@ async def test_evicted_litellm_owned_client_is_closed_once_the_grace_window_elap assert client.closed is True +@pytest.mark.asyncio +async def test_an_evicted_owned_client_collected_inside_its_grace_window_still_closes_its_aiohttp_session(): + """ + The closer holds an evicted client weakly for the whole grace window. An SDK + process that exits, or a proxy that evicts and never calls the cache again, + collects the client first, so only its own finalizer can still close the + aiohttp session (LIT-9767). + """ + clock = _FakeClock() + cache = LLMClientCache( + max_size_in_memory=2, + evicted_client_closer=EvictedClientCloser(grace_seconds=900.0, clock=clock), + ) + session = aiohttp.ClientSession() + client = AsyncOpenAI( + api_key="sk-evicted", + http_client=httpx.AsyncClient(transport=LiteLLMAiohttpTransport(client=session)), + ) + cache.set_cache("client-key", client, litellm_owned_client=True, ttl=600) + + cache.ttl_dict = {key: 0 for key in cache.ttl_dict} + cache.expiration_heap = [(0, key) for _, key in cache.expiration_heap] + cache.evict_cache() + del client + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert session.closed, "the evicted client was collected inside the grace window and left its session open" + + @pytest.mark.asyncio async def test_evicted_caller_supplied_client_is_never_closed(): """litellm does not own a client the caller passed in, so it must stay open.""" diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py index 0765371dc4a..c7d5518ef5b 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py @@ -1,8 +1,10 @@ import aiohttp as aiohttp_aiohttp_handler, asyncio, importlib, litellm import concurrent.futures +import gc import socket import sys import struct +import threading from typing import Final import aiohttp @@ -1310,3 +1312,74 @@ async def test_reset_after_request_is_not_replayed(monkeypatch: pytest.MonkeyPat assert len(received_requests) == 1 finally: await handler.close() + + +async def _new_session() -> aiohttp.ClientSession: + return aiohttp.ClientSession() + + +async def _until_closed(session: aiohttp.ClientSession) -> None: + while not session.closed: + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_a_collected_transport_closes_the_session_it_owns(): + """An SDK client evicted from the cache is collected long before the closer's grace window ends.""" + session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=session) + + del transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert session.closed, "the owned session outlived the transport that was collected before closing it" + + +@pytest.mark.asyncio +async def test_a_collected_transport_leaves_a_shared_session_open(): + session = aiohttp.ClientSession() + try: + transport = LiteLLMAiohttpTransport(client=session, owns_session=False) + + del transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert not session.closed, "a transport closed a session it does not own" + finally: + await session.close() + + +def test_a_transport_collected_after_its_loop_closed_still_closes_its_session(caplog: pytest.LogCaptureFixture): + """asyncio.run() closes the loop the session lives on before the evicted client is collected.""" + loop = asyncio.new_event_loop() + session = loop.run_until_complete(_new_session()) + loop.close() + transport = LiteLLMAiohttpTransport(client=session) + + del transport + gc.collect() + + assert session.closed, "the owned session stayed open after its transport was collected on a closed loop" + del session + gc.collect() + assert "Unclosed client session" not in caplog.text, caplog.text + + +def test_a_transport_collected_while_its_loop_runs_elsewhere_closes_its_session_on_that_loop(): + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + session = asyncio.run_coroutine_threadsafe(_new_session(), loop).result(timeout=10) + transport = LiteLLMAiohttpTransport(client=session) + try: + del transport + gc.collect() + asyncio.run_coroutine_threadsafe(_until_closed(session), loop).result(timeout=10) + + assert session.closed, "the owned session stayed open after its transport was collected from another thread" + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=10) + loop.close() From 9e01b3c37b0736806aa70dcde72f5d85079be3e8 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 13:33:32 -0700 Subject: [PATCH 03/83] fix(complexity_router): keep the held user_turn route for Codex compaction requests (#48627) * test(complexity_router): Codex compaction requests hold the user_turn decision Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(complexity_router): treat Codex compaction traffic as user_turn continuation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(complexity_router): validate Codex turn metadata with frozen models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(complexity_router): flatten body system text through _message_text Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(complexity_router): drop unrelated reformatting from the mapped test file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(complexity_router): skip pin escalation on Codex compaction turns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../complexity_router/complexity_router.py | 97 ++++++++++++++-- ...test_complexity_router_codex_compaction.py | 109 ++++++++++++++++++ .../router_strategy/test_complexity_router.py | 96 +++++++++++++++ 3 files changed, 290 insertions(+), 12 deletions(-) create mode 100644 tests/integration/routing/test_complexity_router_codex_compaction.py diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index cc592bbcbc5..12968bf7d4a 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -28,7 +28,7 @@ from threading import Lock from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast -from pydantic import BaseModel, TypeAdapter, ValidationError, create_model +from pydantic import BaseModel, Field, TypeAdapter, ValidationError, create_model from pydantic_core import ErrorDetails from typing_extensions import ReadOnly, TypedDict @@ -773,6 +773,56 @@ def _newest_turn_is_human_ask( return _last_human_ask_index(messages, marker_pairs) == newest_activity +_CODEX_TURN_METADATA_KEY: Final = "x-codex-turn-metadata" +_CODEX_COMPACTION_SUMMARY_PREFIX: Final = ( + "Another language model started to solve this problem and produced a summary of its thinking process." +) + + +class _CodexTurnMetadata(BaseModel, frozen=True): + request_kind: str | None = None + + +class _CodexTurnMetadataCarrier(BaseModel, frozen=True): + turn_metadata: str | None = Field(default=None, alias=_CODEX_TURN_METADATA_KEY) + + +class _CodexProxyRequest(BaseModel, frozen=True): + headers: _CodexTurnMetadataCarrier | None = None + + +class _CodexRequestMetadata(BaseModel, frozen=True): + """Codex sends its turn metadata in the Responses body's `client_metadata` and projects it onto a + header of the same name for transports without that field""" + + client_metadata: _CodexTurnMetadataCarrier | None = None + proxy_server_request: _CodexProxyRequest | None = None + + def turn_metadata_values(self) -> tuple[str | None, ...]: + headers: Final = self.proxy_server_request.headers if self.proxy_server_request else None + return ( + self.client_metadata.turn_metadata if self.client_metadata else None, + headers.turn_metadata if headers else None, + ) + + +def _codex_request_kind(value: str | None) -> str | None: + if value is None: + return None + try: + return _CodexTurnMetadata.model_validate_json(value).request_kind + except ValidationError: + return None + + +def _is_codex_compaction_request(request_kwargs: Mapping[str, object]) -> bool: + try: + request_metadata: Final = _CodexRequestMetadata.model_validate(request_kwargs) + except ValidationError: + return False + return any(_codex_request_kind(value) == "compaction" for value in request_metadata.turn_metadata_values()) + + def _iter_system_scope_texts( body_system: object, messages: Sequence[Mapping[str, object]], @@ -787,10 +837,9 @@ def _iter_system_scope_texts( must go through the staleness-aware tail scan instead -- scanning it here would floor every turn of a session that once planned, for any pattern whose client injects mid-conversation. """ - if isinstance(body_system, str): - yield body_system - elif isinstance(body_system, list): - yield _message_text(body_system) + body_system_text: Final = _message_text(body_system) + if body_system_text: + yield body_system_text for msg in messages: if msg.get("role") != "system": return @@ -4263,14 +4312,31 @@ class ComplexityRouter(CustomLogger): def _reminder_markers_for_request(self, request_kwargs: Mapping[str, object]) -> tuple[tuple[str, str], ...]: if self.config.reminder_markers is not None: return self._reminder_markers - if any( + return _CODEX_REMINDER_MARKERS if self._is_codex_request(request_kwargs) else _DEFAULT_REMINDER_MARKERS + + @staticmethod + def _is_codex_request(request_kwargs: Mapping[str, object]) -> bool: + return any( is_codex_user_agent(user_agent) for metadata_key in ("litellm_metadata", "metadata") if isinstance(metadata := request_kwargs.get(metadata_key), Mapping) if isinstance(user_agent := metadata.get("user_agent"), str) - ): - return _CODEX_REMINDER_MARKERS - return _DEFAULT_REMINDER_MARKERS + ) + + def _is_codex_compaction_lifecycle( + self, + request_kwargs: Mapping[str, object], + resolved_messages: Sequence[Mapping[str, object]] | None, + marker_pairs: tuple[tuple[str, str], ...], + ) -> bool: + """Codex's compaction call and the first request after it carry a synthetic user turn (the + summarization prompt, then the summary), so `user_turn` must read both as continuations.""" + if _is_codex_compaction_request(request_kwargs): + return True + if not resolved_messages or not self._is_codex_request(request_kwargs): + return False + newest_ask: Final = _newest_turn_ask(resolved_messages, marker_pairs) + return newest_ask is not None and newest_ask.startswith(_CODEX_COMPACTION_SUMMARY_PREFIX) @staticmethod def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: @@ -4380,15 +4446,22 @@ class ComplexityRouter(CustomLogger): # In 'user_turn' mode a held pin is replayed only on continuation turns; a new human # ask falls through and re-classifies. session_affinity restores pin-first for asks too. - pin_replay_allowed: Final = bool(self.config.session_affinity) or not _newest_turn_is_human_ask( - resolved_messages, marker_pairs + codex_compaction: Final = self._is_codex_compaction_lifecycle(request_kwargs, resolved_messages, marker_pairs) + pin_replay_allowed: Final = ( + bool(self.config.session_affinity) + or not _newest_turn_is_human_ask(resolved_messages, marker_pairs) + or codex_compaction ) if cache_key is not None and pin_replay_allowed: pinned_value: Final = await self.litellm_router_instance.cache.async_get_cache(key=cache_key) pinned_pin: Final = _parse_session_affinity_pin(pinned_value, self.config.tier_names()) if pinned_pin is not None: - user_message: Final = _newest_turn_ask(resolved_messages, marker_pairs) if resolved_messages else None + user_message: Final = ( + _newest_turn_ask(resolved_messages, marker_pairs) + if resolved_messages and not codex_compaction + else None + ) pin_escalation_keyword: Final = ( self._matched_escalation_keyword(user_message) if user_message is not None else None ) diff --git a/tests/integration/routing/test_complexity_router_codex_compaction.py b/tests/integration/routing/test_complexity_router_codex_compaction.py new file mode 100644 index 00000000000..2f436901cd0 --- /dev/null +++ b/tests/integration/routing/test_complexity_router_codex_compaction.py @@ -0,0 +1,109 @@ +import json +import uuid +from collections.abc import Callable +from pathlib import Path +from typing import Final + +import yaml +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +_ASK: Final = {"role": "user", "content": "Design a Byzantine fault tolerant consensus protocol and prove it safe"} +_COMPACTION_PROMPT: Final = {"role": "user", "content": "Hello! Summarize the conversation so far."} + + +def _completion(content: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex[:8], + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12}, + } + ).encode() + ) + + +def _answering(content: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + return _completion(content) + + return respond + + +def _classifier(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + body: Final = request.body.decode() + tier: Final = "SIMPLE" if "Summarize" in body else "REASONING" + return _completion(json.dumps({"tier": tier})) + + +def test_codex_compaction_request_stays_on_the_held_user_turn_model(gateway: Gateway, tmp_path: Path) -> None: + with ( + wire_server(_classifier) as judge, + wire_server(_answering("simple answer")) as simple, + wire_server(_answering("reasoning answer")) as reasoning, + ): + router: Final = "router-" + uuid.uuid4().hex[:8] + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = [ + { + "model_name": name, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_base": url, "api_key": "synthetic"}, + } + for name, url in (("judge", judge.url), ("simple", simple.url), ("reasoning", reasoning.url)) + ] + [ + { + "model_name": router, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "judge", "timeout_ms": 20000}, + "classification_mode": "user_turn", + "tiers": { + "SIMPLE": "simple", + "MEDIUM": "simple", + "COMPLEX": "reasoning", + "REASONING": "reasoning", + }, + }, + }, + } + ] + path: Final = tmp_path / "complexity_router.yaml" + path.write_text(yaml.safe_dump(config)) + session: Final = uuid.uuid4().hex + + def codex_headers(request_kind: str) -> dict[str, str]: + return { + "User-Agent": "codex-tui/0.164.0", + "session_id": session, + "x-codex-turn-metadata": json.dumps({"session_id": session, "request_kind": request_kind}), + } + + with owned_proxy(gateway, tmp_path, {}, config=path) as candidate: + ask: Final = candidate.request( + "POST", "/v1/chat/completions", {"model": router, "messages": [_ASK]}, headers=codex_headers("turn") + ) + assert ask.status_code == 200, ask.text + assert ask.json()["choices"][0]["message"]["content"] == "reasoning answer", ask.text + compaction: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": router, "messages": [_ASK, {"role": "assistant", "content": "done"}, _COMPACTION_PROMPT]}, + headers=codex_headers("compaction"), + ) + assert compaction.status_code == 200, compaction.text + assert compaction.json()["choices"][0]["message"]["content"] == "reasoning answer", compaction.text + assert len([call for call in judge.drain() if call.method == "POST"]) == 1 + assert [call for call in simple.drain() if call.method == "POST"] == [] diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index a46f415e10a..2d27a025ffa 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -7820,6 +7820,102 @@ class TestClassificationMode: assert pure.routing_decision["cause"] == "user_turn_continuation" assert hybrid.model == "gpt-4o-mini" + CODEX_USER_AGENT = "codex-tui/0.164.0 (Ubuntu 22.4.0; x86_64) screen (codex-tui; 0.164.0)" + COMPACTION_PROMPT = {"role": "user", "content": "Hello! Summarize the conversation so far."} + COMPACTION_SUMMARY = { + "role": "user", + "content": "Another language model started to solve this problem and produced a summary of its thinking " + "process. You also have access to the state of the tools that were used by that language model.\nHi.", + } + + def _codex_kwargs( + self, session_id: str, request_kind: str, transport: str = "client_metadata" + ) -> dict[str, object]: + turn_metadata: Final = json.dumps({"session_id": session_id, "request_kind": request_kind}) + carrier: Final = ( + {"client_metadata": {"x-codex-turn-metadata": turn_metadata}} + if transport == "client_metadata" + else {"proxy_server_request": {"headers": {"x-codex-turn-metadata": turn_metadata}}} + ) + return {"metadata": {"session_id": session_id, "user_agent": self.CODEX_USER_AGENT}, **carrier} + + @pytest.mark.asyncio + @pytest.mark.parametrize("transport", ["client_metadata", "header"]) + async def test_codex_compaction_request_replays_held_decision( + self, mock_router_instance, user_turn_config, transport + ): + """Codex's compaction call carries a synthetic user prompt; it must hold the session's + decision instead of classifying the prompt as a new human ask.""" + router = self._router(mock_router_instance, user_turn_config) + session = f"codex-compact-{transport}" + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._codex_kwargs(session, "turn"), messages=[self.REASONING_ASK] + ) + compaction = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs(session, "compaction", transport), + messages=[self.REASONING_ASK, self.ASSISTANT_ANSWER, self.COMPACTION_PROMPT], + ) + assert compaction.model == "o1-preview" + assert compaction.routing_decision["cause"] == "user_turn_continuation" + + @pytest.mark.asyncio + async def test_codex_post_compaction_summary_is_a_continuation(self, mock_router_instance, user_turn_config): + """The first request after a mid-turn compaction ends on Codex's summary fragment, which + is not a human ask, so the held decision carries on.""" + router = self._router(mock_router_instance, user_turn_config) + await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs("codex-summary", "turn"), + messages=[self.REASONING_ASK], + ) + resumed = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs("codex-summary", "turn"), + messages=[self.REASONING_ASK, self.COMPACTION_SUMMARY], + ) + assert resumed.model == "o1-preview" + assert resumed.routing_decision["cause"] == "user_turn_continuation" + + @pytest.mark.asyncio + async def test_codex_summary_quoting_escalation_keyword_keeps_held_tier( + self, mock_router_instance, user_turn_config + ): + router = self._router(mock_router_instance, user_turn_config) + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._codex_kwargs("codex-quote", "turn"), messages=[self.SIMPLE_ASK] + ) + quoted_summary: Final = { + "role": "user", + "content": f"{self.COMPACTION_SUMMARY['content']} The user said LITELLM ESCALATE", + } + resumed = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs("codex-quote", "turn"), + messages=[self.SIMPLE_ASK, quoted_summary], + ) + assert resumed.model == "gpt-4o-mini" + assert resumed.routing_decision["cause"] == "user_turn_continuation" + + @pytest.mark.asyncio + async def test_codex_human_ask_after_compaction_reclassifies(self, mock_router_instance, user_turn_config): + router = self._router(mock_router_instance, user_turn_config) + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._codex_kwargs("codex-next", "turn"), messages=[self.REASONING_ASK] + ) + await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs("codex-next", "compaction"), + messages=[self.REASONING_ASK, self.ASSISTANT_ANSWER, self.COMPACTION_PROMPT], + ) + next_ask = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._codex_kwargs("codex-next", "turn"), + messages=[self.REASONING_ASK, self.COMPACTION_SUMMARY, self.SIMPLE_ASK], + ) + assert next_ask.model == "gpt-4o-mini" + assert next_ask.routing_decision["cause"] != "user_turn_continuation" + @pytest.mark.asyncio async def test_session_affinity_wins_when_both_knobs_are_on(self, mock_router_instance, user_turn_config): """With session_affinity also on, the pin short-circuits new asks too and keeps its own From 34c10829b1f2aba5ee44887a413bac12d77e2bc0 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 27 Oct 2026 13:33:47 -0700 Subject: [PATCH 04/83] fix(ui): percent-encode request id in spend log detail URL (#48630) * test: cover request ids with URL special characters in log detail fetches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): percent-encode request id in spend log detail URL Co-authored-by: Charlie Freeman <243496084+cf3901646@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Charlie Freeman <243496084+cf3901646@users.noreply.github.com> --- litellm/proxy/client/cli/commands/debug.py | 3 ++- .../proxy/client/cli/test_debug_commands.py | 18 ++++++++++++++++ .../src/components/networking.test.ts | 21 +++++++++++++++++++ .../src/components/networking.tsx | 4 ++-- 4 files changed, 43 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/client/cli/commands/debug.py b/litellm/proxy/client/cli/commands/debug.py index 8728a9323b7..89b93e2bbb8 100644 --- a/litellm/proxy/client/cli/commands/debug.py +++ b/litellm/proxy/client/cli/commands/debug.py @@ -15,6 +15,7 @@ from datetime import datetime, timezone from pathlib import Path from types import MappingProxyType from typing import Final +from urllib.parse import quote import click import requests @@ -177,7 +178,7 @@ class SpendLogsFetcher: return DebugFailure(f"Unexpected {uri} response: {e}") def payload(self, request_id: str) -> RequestResponsePayload | None | DebugFailure: - uri: Final = f"/spend/logs/ui/{request_id}" + uri: Final = f"/spend/logs/ui/{quote(request_id, safe='')}" raw: Final = self._get(uri) if isinstance(raw, DebugFailure): return raw diff --git a/tests/unit/proxy/client/cli/test_debug_commands.py b/tests/unit/proxy/client/cli/test_debug_commands.py index 419fd821cfb..311edde0af9 100644 --- a/tests/unit/proxy/client/cli/test_debug_commands.py +++ b/tests/unit/proxy/client/cli/test_debug_commands.py @@ -1,6 +1,8 @@ import json import os +import re import time +from urllib.parse import unquote import pytest import requests @@ -114,6 +116,22 @@ def test_recent_bodies_fetches_latest_turns_even_when_successful(): assert _called_paths() == ["/spend/logs/session/ui", "/spend/logs/ui/req-ok"] +@pytest.mark.parametrize("request_id", ["req#a1", "req?a=1", "req%41b", "req a&b"]) +@responses.activate +def test_bodies_are_fetched_for_request_ids_with_url_special_characters(request_id): + def payload_for_stored_id(request): + stored_id = unquote(request.path_url.split("?")[0].removeprefix("/spend/logs/ui/")) + found = stored_id == request_id + return 200, {}, json.dumps({"response": {"id": "msg_special"}} if found else None) + + _mock_proxy([{**OK_ROW, "request_id": request_id}], {}) + responses.add_callback(responses.GET, re.compile(f"{PROXY}/spend/logs/ui/.*"), callback=payload_for_stored_id) + result = CliRunner().invoke(cli, ["debug", "claude", "--session-id", SESSION, "--no-save"]) + + assert result.exit_code == 0, result.output + assert "msg_special" in result.output + + @responses.activate def test_bodies_are_truncated_to_max_chars(): _mock_proxy([OK_ROW], {"req-ok": {"proxy_server_request": {"body": "a" * 5000}, "response": None}}) diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 8b0f0389772..874cdb52fad 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -602,6 +602,27 @@ describe("uiSpendLogsCall exclude_internal_health_checks serialization", () => { }); }); +describe("uiSpendLogDetailsCall", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it.each(["chatcmpl-123#a1", "chatcmpl-123?a=1", "chatcmpl-123%41b", "chatcmpl-123 a&b"])( + "requests the details of %s with the full request id and start date", + async (logId) => { + const mockFetch = vi.fn().mockResolvedValue({ ok: true, json: vi.fn().mockResolvedValue({ id: logId }) }); + vi.stubGlobal("fetch", mockFetch); + + await Networking.uiSpendLogDetailsCall("token", logId, "2026-01-01 00:00:00"); + + const url = new URL(String(mockFetch.mock.calls.at(-1)?.[0]), "http://example.com"); + expect(decodeURIComponent(url.pathname.split("/spend/logs/ui/")[1])).toBe(logId); + expect(url.hash).toBe(""); + expect([...url.searchParams.entries()]).toEqual([["start_date", "2026-01-01 00:00:00"]]); + }, + ); +}); + describe("sessionSpendLogsCall", () => { const originalFetch = global.fetch; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 38d150f1d63..486ec2be068 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -4709,8 +4709,8 @@ export const uiSpendLogDetailsCall = async (accessToken: string, logId: string, try { // Construct base URL let url = proxyBaseUrl - ? `${proxyBaseUrl}/spend/logs/ui/${logId}?start_date=${encodeURIComponent(start_date)}` - : `/spend/logs/ui/${logId}?start_date=${encodeURIComponent(start_date)}`; + ? `${proxyBaseUrl}/spend/logs/ui/${encodeURIComponent(logId)}?start_date=${encodeURIComponent(start_date)}` + : `/spend/logs/ui/${encodeURIComponent(logId)}?start_date=${encodeURIComponent(start_date)}`; const response = await fetch(url, { method: "GET", From 52562c47f83be9df28ecf6ca9f9fb620286d04b9 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 27 Oct 2026 13:53:23 -0700 Subject: [PATCH 05/83] refactor(proxy-extras): drop the spend-log api_key special case from the row check remedy (#48636) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_proxy_extras/migration_data_checks.py | 8 -------- .../database/migrations/test_request_log_indexes.py | 9 ++++++--- .../litellm_proxy_extras/test_migration_data_checks.py | 2 +- 3 files changed, 7 insertions(+), 12 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py index a2d3728349a..7b01e64dd73 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py @@ -42,12 +42,6 @@ _OCTET_LENGTH_TYPES: Final = frozenset({"text", "varchar", "bpchar", "bytea", "n MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR: Final = "LITELLM_MIGRATION_DATA_CHECK_TIMEOUT" DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS: Final = 600.0 REFUSAL: Final = "Existing rows would make a pending schema change fail:" -LEGACY_SPEND_LOG_KEY_REMEDY: Final = ( - "Rows written before v1.99.0 hold the raw key. Hash them in place (one-way, idempotent):\n" - ' UPDATE "LiteLLM_SpendLogs" SET api_key = encode(sha256(convert_to(' - "regexp_replace(api_key, '^[Bb][Ee][Aa][Rr][Ee][Rr] ', ''), 'UTF8')), 'hex')\n" - " WHERE api_key <> '' AND api_key !~ '^[0-9a-fA-F]{64}$' AND api_key !~ '^hashed-jwt-[0-9A-Fa-f]{64}$';" -) Kind = Literal["too_long", "duplicates", "nulls", "any_row", "orphans", "check_violations"] @@ -571,8 +565,6 @@ def find_blockers( def _remedy(check: Check) -> str: - if check.kind == "too_long" and check.table == "LiteLLM_SpendLogs" and "api_key" in check.columns: - return LEGACY_SPEND_LOG_KEY_REMEDY joined: Final = ", ".join(check.columns) match check.kind: case "too_long": diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index bf37ca80665..1de2090b487 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -956,8 +956,11 @@ def _legacy_bearer_token() -> str: return "Bearer " + secrets.token_hex(1600) -def _hash_sql_from(refusal: str) -> str: - return refusal[refusal.index('UPDATE "LiteLLM_SpendLogs"') : refusal.index("{64}$';") + len("{64}$';")] +_HASH_LEGACY_KEYS: Final = ( + 'UPDATE "LiteLLM_SpendLogs" SET api_key = encode(sha256(convert_to(' + "regexp_replace(api_key, '^[Bb][Ee][Aa][Rr][Ee][Rr] ', ''), 'UTF8')), 'hex') " + "WHERE api_key <> '' AND api_key !~ '^[0-9a-fA-F]{64}$' AND api_key !~ '^hashed-jwt-[0-9A-Fa-f]{64}$'" +) @requires_db @@ -985,7 +988,7 @@ def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_key assert API_KEY_INDEX not in _index_validity(scratch_database, "api_key_startTime_idx") with psycopg.connect(scratch_database, autocommit=True) as conn: - conn.execute(_hash_sql_from(refusal)) # pyright: ignore[reportArgumentType] # the refusal's own SQL + conn.execute(_HASH_LEGACY_KEYS) assert ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver) is True assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True} assert len(_ledger(scratch_database)) > len(ledger_before) diff --git a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py index 7897a388474..f3b787beda1 100644 --- a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py +++ b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py @@ -125,7 +125,7 @@ def test_the_report_names_the_rows_the_statement_and_the_fix_and_says_nothing_wa "(api_key up to 3207 bytes)" ) in text assert 'blocks: CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" (request-log index)' in text - assert 'UPDATE "LiteLLM_SpendLogs" SET api_key = encode(sha256(' in text + assert 'fix: Shorten or hash the long "LiteLLM_SpendLogs" (api_key, startTime) values' in text assert text.startswith("Existing rows would make a pending schema change fail:") assert "No schema change was made" in text From e0d1e64591dc0f4243a17014ce216943ffa0d186 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 27 Oct 2026 17:04:30 -0400 Subject: [PATCH 06/83] feat(router): respect caller output budgets in complexity routing (#48632) Co-authored-by: stuti612 --- litellm/router.py | 96 ++++---- .../complexity_router/README.md | 14 ++ .../complexity_router/complexity_router.py | 63 ++++-- .../complexity_router/config.py | 5 +- .../common_utils/test_cache_aware_routing.py | 18 +- .../router_strategy/test_complexity_router.py | 206 +++++++++++++++++- .../add_model/ComplexityRouterConfig.tsx | 2 +- .../add_model/ResponseFormatControls.tsx | 56 ++++- .../build_complexity_router_config.test.ts | 20 ++ .../build_complexity_router_config.ts | 8 +- ...d_updated_complexity_router_config.test.ts | 23 ++ ...dit_auto_router_modal.integration.test.tsx | 22 +- .../hydrate_complexity_router_config.ts | 3 +- .../src/lib/autorouter_presets.test.ts | 7 + .../src/lib/autorouter_presets.ts | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 16 files changed, 448 insertions(+), 100 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 0d4045e6ca8..c8153b596d6 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1122,6 +1122,27 @@ def as_output_cap(value: object) -> int | None: return cap if cap >= 0 else None +def stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str, int] | None: + """None means unstamped; an empty mapping means the original caller supplied no limit.""" + stamped: Final = next( + ( + bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) + for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket + ), + None, + ) + if not isinstance(stamped, dict): + return None + return MappingProxyType( + { + carrier: cap + for carrier, value in stamped.items() + if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None + } + ) + + def _exact_base_model_max_input_tokens(base_model: str) -> int | None: try: model_info: Final = litellm.get_model_info(model=base_model) @@ -15235,9 +15256,7 @@ class Router: request_kwargs: dict, responses_call: bool, ) -> bool: - """Apply a routing strategy's per-tier litellm_params on top of the request and report - whether they pinned an output ceiling, so the caller can hand the request its own ceiling - back on a routing pass that pins none.""" + """Apply per-tier litellm_params and report whether they pinned an output ceiling.""" if not tier_litellm_params: return False reject_api_key_secret_reference(tier_litellm_params) @@ -15250,34 +15269,26 @@ class Router: return not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(surface_tier_params) @staticmethod - def _restore_client_ceiling_no_tier_pins(request_kwargs: MutableMapping[str, object]) -> None: - """A model-group fallback re-enters routing with the kwargs an earlier auto-router pass - already rewrote, so a ceiling sized for that pass's tier would ride onto a group no tier - chose. When this pass pins none, hand the request back exactly the carriers the caller - sent, which the first pinning pass stamped. The stamp lives in a metadata bucket a - caller can also write, so the proxy strips the key at ingestion and this read takes - nothing but the three ceiling carriers as integers: no other key ever reaches kwargs.""" - stamped: Final = next( - ( - bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) - for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) - if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket - ), - None, - ) - if not isinstance(stamped, dict): + def _restore_client_output_ceiling(request_kwargs: MutableMapping[str, object]) -> None: + """Restore caller intent before routing, so neither a previous tier pin nor alias defaults + become the next attempt's budget. The proxy strips caller-written stamps at ingestion.""" + callers_ceiling: Final = stamped_client_ceiling(request_kwargs) + if callers_ceiling is None: return - callers_ceiling: Final = MappingProxyType( - { - carrier: cap - for carrier, value in stamped.items() - if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None - } - ) for carrier in OUTPUT_TOKEN_CEILING_PARAMS: request_kwargs.pop(carrier, None) request_kwargs.update(callers_ceiling) + @staticmethod + def _stamp_client_output_ceiling( + request_kwargs: dict[str, object], # mutable-ok: preserve request-owned metadata across routing attempts + ) -> None: + _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) + metadata_bucket.setdefault( + CLIENT_OUTPUT_CEILING_METADATA_KEY, + {carrier: request_kwargs[carrier] for carrier in OUTPUT_TOKEN_CEILING_PARAMS if carrier in request_kwargs}, + ) + @staticmethod def _drop_client_carriers_a_tier_pin_supersedes( request_kwargs: dict[str, object], @@ -15294,15 +15305,7 @@ class Router: one setting under three names, and a provider handed two of them either rejects the request or picks one by iteration order.""" if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_litellm_params): - _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) - metadata_bucket.setdefault( - CLIENT_OUTPUT_CEILING_METADATA_KEY, - { - carrier: request_kwargs[carrier] - for carrier in OUTPUT_TOKEN_CEILING_PARAMS - if carrier in request_kwargs - }, - ) + Router._stamp_client_output_ceiling(request_kwargs) for carrier in OUTPUT_TOKEN_CEILING_PARAMS: request_kwargs.pop(carrier, None) if "reasoning_effort" not in tier_litellm_params: @@ -15374,14 +15377,12 @@ class Router: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages record_pre_routing_selection(request_kwargs, model) - tier_pins_ceiling: Final = self._pin_tier_params_onto_request( + self._pin_tier_params_onto_request( model=model, tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, request_kwargs=request_kwargs, responses_call=responses_call, ) - if not tier_pins_ceiling: - self._restore_client_ceiling_no_tier_pins(request_kwargs) ######################################################### # Resolve the strategy and logger AFTER the pre-routing hook, since @@ -15539,14 +15540,12 @@ class Router: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages record_pre_routing_selection(request_kwargs, model) - tier_pins_ceiling: Final = self._pin_tier_params_onto_request( + self._pin_tier_params_onto_request( model=model, tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, request_kwargs=request_kwargs, responses_call=responses_call, ) - if not tier_pins_ceiling: - self._restore_client_ceiling_no_tier_pins(request_kwargs) # 2. Get healthy deployments healthy_deployments: Final = await self.async_get_healthy_deployments( @@ -15896,6 +15895,7 @@ class Router: spend metadata is stamped before routing and the response carries the tier group the strategy picked. """ + self._restore_client_output_ceiling(request_kwargs) self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) requested_registered_model_name: Final = self.get_model_from_alias(model=model) or model registered_model_name: Final = await self._resolve_claude_code_session_router( @@ -16062,7 +16062,8 @@ class Router: # them) until `_update_kwargs_with_deployment` drops any the selected # deployment sets itself (its own `aws_region_name` beats the marker's). # Per-tier `litellm_params` on the hook response are deliberate overrides - # the caller applies on top, so those keys are never forwarded here. + # the caller applies on top, so those keys and equivalent output-limit + # names are never tracked as alias defaults here. marker_params: Final = ( self._forwardable_alias_marker_params( model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs @@ -16073,9 +16074,16 @@ class Router: tier_param_keys: Final = ( tuple(pre_routing_hook_response.litellm_params or ()) if pre_routing_hook_response is not None else () ) - newly_forwarded: Final = tuple( - (key, value) for key, value in marker_params if key not in request_kwargs and key not in tier_param_keys + overridden_marker_keys: Final = frozenset(tier_param_keys) | ( + OUTPUT_TOKEN_CEILING_PARAMS if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_param_keys) else frozenset() ) + newly_forwarded: Final = tuple( + (key, value) + for key, value in marker_params + if key not in request_kwargs and key not in overridden_marker_keys + ) + if pre_routing_hook_response is not None: + self._stamp_client_output_ceiling(request_kwargs) request_kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, None) request_kwargs.update(newly_forwarded) if newly_forwarded: diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 7567a0ef37e..08d39841b61 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -107,6 +107,20 @@ Session affinity, user-turn classification, adaptive routing, routing plugins, c When cache costs change the model, the routing decision reports `cause: prompt_cache_cost`. Its signals include the original model, classification cause and both estimated costs +### Output token limits + +Set `max_tokens_from_tier_model` under `complexity_router_config` to choose how a routed request's output limit is determined: + +| Setting | Effective output limit | +| --- | --- | +| `true` (default) | Replace the caller's limit with the selected model group's ceiling | +| `"clamp"` | Use the smaller of the caller's positive limit and the selected model group's ceiling | +| `false` | Keep the caller's limit | + +Clamp reads `max_tokens`, `max_completion_tokens` and `max_output_tokens`. If several are supplied, it uses the smallest positive value. With no positive caller limit, it uses the model ceiling. For example, a caller requesting 100 tokens keeps that limit on a model with a 32,768-token ceiling; a request for 100,000 is capped at 32,768. Rerouting uses the original caller budget, and cache-aware cost estimates use the same effective limit. + +All modes preserve explicit output limits in a tier's own `litellm_params`. An unknown ceiling leaves the caller's value unchanged. In the dashboard, select **Respect caller budget** under **Advanced settings > Compatibility > Output token limit** to enable clamp. + ### Capability forecasting Set `classifier_type: capability` to use diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 12968bf7d4a..f5fd4ae9c98 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -23,6 +23,7 @@ import random import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from functools import partial from itertools import accumulate, chain, islice, takewhile from threading import Lock from types import MappingProxyType @@ -161,6 +162,17 @@ def _tier_name(tier: ComplexityTier | str) -> str: return tier.value if isinstance(tier, ComplexityTier) else tier +def _callers_output_cap(request_kwargs: Mapping[str, object]) -> int | None: + """The smallest positive output ceiling the caller sent under any of its three names, read + from the stamp when an earlier routing pass already rewrote the kwargs.""" + from litellm.router import as_output_cap, stamped_client_ceiling + + stamped: Final = stamped_client_ceiling(request_kwargs) + carriers: Final = request_kwargs if stamped is None else stamped + caps: Final = tuple(as_output_cap(carriers.get(carrier)) for carrier in OUTPUT_TOKEN_CEILING_PARAMS) + return min((cap for cap in caps if cap is not None and cap > 0), default=None) + + def _built_in_tier_or_none(tier_name: str) -> ComplexityTier | None: """The built-in tier a `tiers` key names, or None when the key is an operator-defined name.""" return ComplexityTier.__members__.get(tier_name) @@ -3028,7 +3040,9 @@ class ComplexityRouter(CustomLogger): raise ValueError(f"No model configured for tier {tier_key} and no default_model set") - def _litellm_params_for_model(self, tier: ComplexityTier | str | None, model: str) -> Mapping[str, object]: + def _litellm_params_for_model( + self, tier: ComplexityTier | str | None, model: str, request_kwargs: Mapping[str, object] + ) -> Mapping[str, object]: entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) if tier is not None else () entry: Final = next((candidate for candidate in entries if candidate.model_name == model), None) explicit: Final = entry.litellm_params if entry is not None else MappingProxyType({}) @@ -3037,7 +3051,12 @@ class ComplexityRouter(CustomLogger): ceiling: Final = self._group_output_ceiling(model) if ceiling is None: return explicit - return MappingProxyType({**explicit, "max_tokens": ceiling}) + callers_cap: Final = ( + _callers_output_cap(request_kwargs) if self.config.max_tokens_from_tier_model == "clamp" else None + ) + return MappingProxyType( + {**explicit, "max_tokens": ceiling if callers_cap is None else min(callers_cap, ceiling)} + ) @staticmethod def _pick_from_tier_value(model: str | Sequence[str], tier_key: str) -> str: @@ -3768,7 +3787,7 @@ class ComplexityRouter(CustomLogger): response: PreRoutingHookResponse, messages: list[dict[str, object]] | None, # mutable-ok: forwarded verbatim to the list-typed re-pick resolved_messages: Sequence[Mapping[str, object]] | None, - request_kwargs: dict, # mutable-ok: same shape the hook receives + request_kwargs: dict[str, object], # mutable-ok: same shape the hook receives context_fit: _RequestContextFit | None = None, retained_pin: _SessionAffinityPin | None = None, ) -> PreRoutingHookResponse: @@ -3865,7 +3884,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing=bool(decision.get("conversation_continuing", True)) if decision is not None else True, - tier_litellm_params=self._litellm_params_for_model(new_tier, new_model), + tier_litellm_params=self._litellm_params_for_model(new_tier, new_model, request_kwargs), context_escalation_original_tier=( decision.get("context_escalation_original_tier") if decision is not None else None ), @@ -3876,7 +3895,7 @@ class ComplexityRouter(CustomLogger): return HookResponse( model=new_model, messages=response.messages, - litellm_params=self._litellm_params_for_model(new_tier, new_model), + litellm_params=self._litellm_params_for_model(new_tier, new_model, request_kwargs), routing_decision=new_decision, ) @@ -3904,7 +3923,7 @@ class ComplexityRouter(CustomLogger): model_name: str, messages: list[dict[str, Any]] | None, # mutable-ok: forwarded verbatim to the router's own probe input: str | list | None, # mutable-ok: mirrors the owner's own input parameter, which this forwards verbatim - request_kwargs: dict, # mutable-ok: same shape the hook receives + request_kwargs: dict[str, object], # mutable-ok: same shape the hook receives ) -> bool: """Whether the router would find a deployment for this group ON THIS REQUEST. @@ -3962,7 +3981,7 @@ class ComplexityRouter(CustomLogger): messages: list[dict[str, object]] | None, # mutable-ok: forwarded verbatim to the list-typed re-pick input: str | list | None, # mutable-ok: mirrors the owner's own input parameter, which this forwards verbatim resolved_messages: Sequence[Mapping[str, object]] | None, - request_kwargs: dict, # mutable-ok: same shape the hook receives + request_kwargs: dict[str, object], # mutable-ok: same shape the hook receives context_fit: _RequestContextFit | None = None, retained_pin: _SessionAffinityPin | None = None, ) -> PreRoutingHookResponse: @@ -4045,14 +4064,14 @@ class ComplexityRouter(CustomLogger): classifier_model=decision.get("classifier_model"), classifier_cost=decision.get("classifier_cost"), conversation_continuing=bool(decision.get("conversation_continuing", True)), - tier_litellm_params=self._litellm_params_for_model(candidate_tier, new_model), + tier_litellm_params=self._litellm_params_for_model(candidate_tier, new_model, request_kwargs), context_escalation_original_tier=decision.get("context_escalation_original_tier"), previous_decision=decision, ) return response.model_copy( update={ "model": new_model, - "litellm_params": self._litellm_params_for_model(candidate_tier, new_model), + "litellm_params": self._litellm_params_for_model(candidate_tier, new_model, request_kwargs), "routing_decision": new_decision, } ) @@ -4090,14 +4109,14 @@ class ComplexityRouter(CustomLogger): classifier_model=decision.get("classifier_model"), classifier_cost=decision.get("classifier_cost"), conversation_continuing=bool(decision.get("conversation_continuing", True)), - tier_litellm_params=self._litellm_params_for_model(None, default_model), + tier_litellm_params=self._litellm_params_for_model(None, default_model, request_kwargs), context_escalation_original_tier=decision.get("context_escalation_original_tier"), previous_decision=decision, ) return response.model_copy( update={ "model": default_model, - "litellm_params": self._litellm_params_for_model(None, default_model), + "litellm_params": self._litellm_params_for_model(None, default_model, request_kwargs), "routing_decision": default_decision, } ) @@ -4580,7 +4599,9 @@ class ComplexityRouter(CustomLogger): if pin_placement is not None and pin_context_original_tier is not None else pin_source_tier ) - session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model) + session_tier_litellm_params: Final = self._litellm_params_for_model( + routed_pin_tier, routed_model, request_kwargs + ) has_original_messages: Final = messages is not None and len(messages) > 0 return self._with_session_deployment_affinity( await self._gate_response_health( @@ -4678,14 +4699,14 @@ class ComplexityRouter(CustomLogger): choice: Final = await choose_cached_model( router=self.litellm_router_instance, config=self.config, - params_for_model=self._litellm_params_for_model, + params_for_model=partial(self._litellm_params_for_model, request_kwargs=request_kwargs), response=response, request_kwargs=request_kwargs, messages=messages, ) if choice is None or not context_fit.accepts(choice.model): return response - params: Final = self._litellm_params_for_model(choice.tier, choice.model) + params: Final = self._litellm_params_for_model(choice.tier, choice.model, request_kwargs) decision: Final[StandardLoggingRoutingDecision] = { **response.routing_decision, "routed_model": choice.model, @@ -4714,7 +4735,7 @@ class ComplexityRouter(CustomLogger): async def _classify_and_route( self, model: str, - request_kwargs: dict, + request_kwargs: dict[str, object], messages: list[dict[str, object]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, @@ -4772,7 +4793,7 @@ class ComplexityRouter(CustomLogger): ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs ) fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM - default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model) + default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model, request_kwargs) return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, @@ -4810,7 +4831,7 @@ class ComplexityRouter(CustomLogger): _tier_name(plan_floor), routed_model, ) - plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model) + plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model, request_kwargs) return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, @@ -4839,7 +4860,9 @@ class ComplexityRouter(CustomLogger): ) keyword_plan_floored: Final = routed_tier != escalated_tier routed_model = await self._pick_model_for_tier(routed_tier, messages, resolved_messages, request_kwargs) - keyword_tier_litellm_params: Final = self._litellm_params_for_model(routed_tier, routed_model) + keyword_tier_litellm_params: Final = self._litellm_params_for_model( + routed_tier, routed_model, request_kwargs + ) keyword_cause: Final[RoutingDecisionCause] = ( "plan_mode" if keyword_plan_floored @@ -4919,7 +4942,7 @@ class ComplexityRouter(CustomLogger): outcome.signals, fallback_model, ) - fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model) + fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model, request_kwargs) return PreRoutingHookResponse( model=fallback_model, messages=messages if has_original_messages else None, @@ -5001,7 +5024,7 @@ class ComplexityRouter(CustomLogger): routed_model, ) - tier_litellm_params: Final = self._litellm_params_for_model(tier, routed_model) + tier_litellm_params: Final = self._litellm_params_for_model(tier, routed_model, request_kwargs) classifier_model: Final = ( f"{outcome.jev_verdict.provider}/{outcome.jev_verdict.model}" if outcome.cause == "jev_classifier" and outcome.jev_verdict is not None diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index ec960a19f79..020f2d33742 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1462,7 +1462,7 @@ class ComplexityRouterConfig(LiteLLMBaseModel): "wording the built-ins don't cover, or after a client release changes its strings." ), ) - max_tokens_from_tier_model: bool = Field( + max_tokens_from_tier_model: bool | Literal["clamp"] = Field( default=True, description=( "Set max_tokens on every routed request to the output ceiling of the tier model it " @@ -1473,7 +1473,8 @@ class ComplexityRouterConfig(LiteLLMBaseModel): "from each deployment's model_info and then the model cost map; a tier model with a " "deployment whose ceiling is unknown keeps the caller's value. A max_tokens, " "max_completion_tokens or max_output_tokens in the tier's own litellm_params still " - "wins. Set false to forward the caller's value unchanged." + 'wins. Set "clamp" to keep a smaller positive caller value and cap a larger one at ' + "the ceiling. Set false to forward the caller's value unchanged." ), ) route_housekeeping_to_cheapest_tier: bool = Field( diff --git a/tests/unit/proxy/common_utils/test_cache_aware_routing.py b/tests/unit/proxy/common_utils/test_cache_aware_routing.py index 24fc1f22a6a..4eb22b2e34f 100644 --- a/tests/unit/proxy/common_utils/test_cache_aware_routing.py +++ b/tests/unit/proxy/common_utils/test_cache_aware_routing.py @@ -1,5 +1,6 @@ from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from functools import partial from typing import Final import pytest @@ -148,7 +149,7 @@ async def _select( return await select_cached_model( router=router, config=config, - params_for_model=complexity._litellm_params_for_model, + params_for_model=partial(complexity._litellm_params_for_model, request_kwargs=request_kwargs), response=response, body=body, request_kwargs=request_kwargs, @@ -234,7 +235,9 @@ async def test_disabled_setting_does_not_access_prediction_services() -> None: await choose_cached_model( router=_router(), config=config, - params_for_model=ComplexityRouter("smart", _router(), config.model_dump())._litellm_params_for_model, + params_for_model=partial( + ComplexityRouter("smart", _router(), config.model_dump())._litellm_params_for_model, request_kwargs={} + ), response=_response(), request_kwargs={}, messages=None, @@ -348,6 +351,7 @@ def test_repeated_model_in_multiple_tiers_is_only_considered_once() -> None: (True, "custom_endpoint", "cheap"), (True, "compaction", "cheap"), (True, "guardrail", "cheap"), + (True, "clamp", "strong"), ], ) async def test_router_applies_opt_in_and_preserves_failure_semantics( @@ -368,7 +372,9 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( from litellm.router_strategy.complexity_router.context_compaction import initialize_compaction_state config: Final = _config( - cache_aware_routing=enabled, cache_aware_routing_timeout_ms=1 if behavior == "deadline" else 2000 + cache_aware_routing=enabled, + cache_aware_routing_timeout_ms=1 if behavior == "deadline" else 2000, + max_tokens_from_tier_model="clamp" if behavior == "clamp" else True, ) models: Final = _router( strong_output_rate=0.000048 if behavior == "tier_cost" else 0.000015, @@ -423,6 +429,7 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( body: Final = { **_body(), **({"max_tokens": 1000} if behavior in ("tier_cost", "tier_context") else {}), + **({"max_tokens": 100} if behavior == "clamp" else {}), **({"thinking": {"type": "enabled", "budget_tokens": 10000}} if behavior == "unsupported_shape" else {}), } kwargs: Final = { @@ -437,6 +444,7 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( else {} ), **({"guardrails": ["test-guardrail"]} if behavior == "guardrail" else {}), + **({"max_tokens": 100} if behavior == "clamp" else {}), } if expected is None: with pytest.raises(asyncio.CancelledError): @@ -458,6 +466,10 @@ async def test_router_applies_opt_in_and_preserves_failure_semantics( if enabled and behavior == "success": assert requests.qsize() == 4 assert response.routing_decision is not None + if behavior == "clamp": + assert "cache-aware:estimated-cost=0.01950000;original-cost=0.06390000" in ( + response.routing_decision["signals"] or () + ), "both arms must be estimated at the caller's 100 output tokens, not the 1024 default" assert response.routing_decision["cause"] == ( "prompt_cache_cost" if expected == "strong" else "heuristic_scorer" ) diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 2d27a025ffa..78dc2eaf314 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -36,11 +36,12 @@ from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, OUTPUT_TOKEN_CEILING_PARAMS, RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, ) -from litellm.router import as_output_cap +from litellm.router import as_output_cap, stamped_client_ceiling from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -13733,7 +13734,7 @@ def test_tier_model_params_are_used_by_pools_and_savings_baseline(mock_router_in assert router._tier_pools() == {"SIMPLE": ["mini"], "REASONING": ["opus", "abc"]} assert router._hardest_tier_models() == ("opus", "abc") - assert router._litellm_params_for_model(ComplexityTier.REASONING, "opus") == {"reasoning_effort": "xhigh"} + assert router._litellm_params_for_model(ComplexityTier.REASONING, "opus", {}) == {"reasoning_effort": "xhigh"} @pytest.mark.asyncio @@ -16917,9 +16918,10 @@ class TestMaxTokensFromTierModel: @staticmethod def _router( tier_litellm_params: dict | None = None, - max_tokens_from_tier_model: bool | None = None, + max_tokens_from_tier_model: bool | Literal["clamp"] | None = None, simple_deployments: list[dict] | None = None, extra_config: dict | None = None, + marker_params: Mapping[str, object] | None = None, ) -> Router: simple_tier: dict = {"model_name": "small"} if tier_litellm_params: @@ -16934,7 +16936,11 @@ class TestMaxTokensFromTierModel: model_list=[ { "model_name": "smart-router", - "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": config}, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": config, + **(marker_params or {}), + }, }, *(simple_deployments or [TestMaxTokensFromTierModel.SMALL]), { @@ -17016,6 +17022,167 @@ class TestMaxTokensFromTierModel: assert sent["max_tokens"] == 8192 + @pytest.mark.asyncio + @pytest.mark.parametrize( + "prompt, client_kwargs, expected", + [ + (COMPLEX_PROMPT, {"max_tokens": 2000}, ("anthropic/claude-sonnet-5", 2000)), + ("hi", {"max_tokens": 100000}, ("anthropic/claude-haiku-4-5", 8192)), + (COMPLEX_PROMPT, {}, ("anthropic/claude-sonnet-5", 64000)), + (COMPLEX_PROMPT, {"max_tokens": 0}, ("anthropic/claude-sonnet-5", 64000)), + (COMPLEX_PROMPT, {"max_completion_tokens": 2000}, ("anthropic/claude-sonnet-5", 2000)), + (COMPLEX_PROMPT, {"max_tokens": 3000, "max_completion_tokens": 2000}, ("anthropic/claude-sonnet-5", 2000)), + ], + ids=["below-ceiling", "above-ceiling", "omitted", "zero", "other-carrier", "smallest-carrier"], + ) + async def test_clamp_keeps_the_smaller_of_the_client_value_and_the_ceiling( + self, prompt: str, client_kwargs: dict[str, int], expected: tuple[str, int] + ) -> None: + sent: Final = await self._routed(self._router(max_tokens_from_tier_model="clamp"), prompt, **client_kwargs) + + assert (sent["model"], sent["max_tokens"]) == expected + assert not (OUTPUT_TOKEN_CEILING_PARAMS - {"max_tokens"}) & sent.keys() + + @pytest.mark.asyncio + async def test_clamp_on_the_responses_surface_keeps_its_own_name(self) -> None: + router: Final = self._router(max_tokens_from_tier_model="clamp") + + sent: Final = await self._routed_responses(router, self.COMPLEX_PROMPT, max_output_tokens=2000) + + assert (sent["model"], sent["max_output_tokens"]) == ("anthropic/claude-sonnet-5", 2000) + assert "max_tokens" not in sent + + @pytest.mark.asyncio + @pytest.mark.parametrize("tier_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + @pytest.mark.parametrize("marker_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + @pytest.mark.parametrize("surface", ["chat", "responses"]) + @pytest.mark.parametrize("caller_cap", [None, 100]) + async def test_clamp_leaves_the_operators_own_tier_ceiling_in_charge( + self, + tier_carrier: str, + marker_carrier: str, + surface: str, + caller_cap: int | None, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + surface_key: Final = "max_output_tokens" if surface == "responses" else "max_tokens" + caller_key: Final = "max_tokens" if surface == "responses" else "max_output_tokens" + request_caps: Final = {caller_key: caller_cap} if caller_cap is not None else {} + router: Final = self._router( + tier_litellm_params={tier_carrier: 4321}, + max_tokens_from_tier_model="clamp", + marker_params={marker_carrier: 2000}, + simple_deployments=[{ + "model_name": "small", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "test-only", + "api_base": "https://tier.test/v1", + surface_key: 9000, + }, + "model_info": {"max_output_tokens": 16384}, + }], + ) + + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host="tier.test").mock(side_effect=TestHealthFallbackDispatch._http_response) + if surface == "responses": + await router.aresponses(model="smart-router", input="hi", **request_caps) + else: + await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "hi"}], **request_caps + ) + + assert len(upstream.calls) == 1 + sent: Final = json.loads(upstream.calls[0].request.content) + wire_key: Final = "max_output_tokens" if surface == "responses" else "max_completion_tokens" + assert {key: value for key, value in sent.items() if key in OUTPUT_TOKEN_CEILING_PARAMS} == {wire_key: 4321} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "callers_original, expected", + [({"max_tokens": 2000}, 2000), ({}, 64000)], + ids=["caller-sent-a-value", "caller-sent-nothing"], + ) + async def test_clamp_on_a_rerouted_request_uses_the_callers_original_value( + self, callers_original: dict[str, int], expected: int + ) -> None: + """A retry re-enters routing with kwargs an earlier pass already rewrote; the stamp holds + what the caller actually sent.""" + router: Final = self._router(max_tokens_from_tier_model="clamp") + rewritten: Final[dict[str, object]] = { + "max_tokens": 8192, + "metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: callers_original}, + } + + sent: Final = await self._routed(router, self.COMPLEX_PROMPT, **rewritten) + + assert sent["max_tokens"] == expected + + @pytest.mark.asyncio + @pytest.mark.parametrize("marker_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + @pytest.mark.parametrize("caller_cap", [None, 3000]) + @pytest.mark.parametrize("tier_cap", [None, 4321]) + @pytest.mark.parametrize("known_ceiling", [True, False]) + async def test_rerouting_preserves_caller_budget_without_adopting_marker_defaults( + self, marker_carrier: str, caller_cap: int | None, tier_cap: int | None, known_ceiling: bool + ) -> None: + router: Final = self._router( + max_tokens_from_tier_model="clamp", + marker_params={marker_carrier: 2000}, + tier_litellm_params={"max_tokens": tier_cap} if tier_cap else None, + simple_deployments=None if known_ceiling else [{ + "model_name": "small", + "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, + }], + ) + request: Final[dict[str, object]] = {} if caller_cap is None else {"max_output_tokens": caller_cap} + await router.async_get_available_deployment( + model="smart-router", request_kwargs=request, messages=[{"role": "user", "content": "hi"}] + ) + first: Final = {key: request[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in request} + await router.async_get_available_deployment( + model="smart-router", request_kwargs=request, + messages=[{"role": "user", "content": self.COMPLEX_PROMPT}], + ) + + expected_first: Final = ( + {"max_tokens": tier_cap or caller_cap or 8192} + if known_ceiling or tier_cap else + {marker_carrier: 2000, **({"max_output_tokens": caller_cap} if caller_cap else {})} + ) + assert first == expected_first + assert request["max_tokens"] == (caller_cap or 64000) + assert not (OUTPUT_TOKEN_CEILING_PARAMS - {"max_tokens"}) & request.keys() + + await router.async_get_available_deployment( + model="smart-router", request_kwargs=request, messages=[{"role": "user", "content": "hi"}] + ) + assert {key: request[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in request} == first + + @pytest.mark.asyncio + @pytest.mark.parametrize("mode", [False, "clamp"]) + @pytest.mark.parametrize("marker_carrier", sorted(OUTPUT_TOKEN_CEILING_PARAMS)) + async def test_unpinned_output_limits_still_use_marker_defaults( + self, mode: bool | Literal["clamp"], marker_carrier: str + ) -> None: + router: Final = self._router( + max_tokens_from_tier_model=mode, + marker_params={marker_carrier: 2000}, + simple_deployments=[{ + "model_name": "small", + "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}, + }], + ) + sent: Final = await self._routed(router) + + assert {key: sent[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in sent} == {marker_carrier: 2000} + + def test_an_unknown_mode_is_rejected(self) -> None: + with pytest.raises(ValidationError): + ComplexityRouterConfig(tiers={"SIMPLE": "small"}, max_tokens_from_tier_model="sometimes") + @pytest.mark.asyncio async def test_a_tier_model_with_an_unknown_ceiling_keeps_the_client_value(self): unmapped: dict = {"model_name": "small", "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}} @@ -17112,7 +17279,7 @@ class TestMaxTokensFromTierModel: complexity_router_config={"tiers": {"SIMPLE": "small"}, "default_model": "big"}, ) - assert dict(strategy._litellm_params_for_model(None, "big")) == {"max_tokens": 64000} + assert dict(strategy._litellm_params_for_model(None, "big", {})) == {"max_tokens": 64000} @pytest.mark.asyncio async def test_a_fallback_into_a_plain_group_gets_the_callers_ceiling_back(self): @@ -17228,29 +17395,52 @@ class TestMaxTokensFromTierModel: def test_a_client_cap_is_read_as_an_integer_or_ignored(self, value, expected): assert as_output_cap(value) == expected + @pytest.mark.parametrize( + "request_kwargs, expected", + [ + ({"max_tokens": 8192}, None), + ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {}}}, {}), + ({"litellm_metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_tokens": "2000"}}}, {"max_tokens": 2000}), + ( + {"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_output_tokens": 5, "api_key": "x", "max_tokens": True}}}, + {"max_output_tokens": 5}, + ), + ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: "not-a-dict"}}, None), + ], + ids=["no-stamp", "empty-stamp", "litellm-metadata-bucket", "only-integer-carriers", "malformed-stamp"], + ) + def test_the_stamp_reads_back_only_the_callers_integer_carriers( + self, request_kwargs: dict[str, object], expected: dict[str, int] | None + ) -> None: + stamped: Final = stamped_client_ceiling(request_kwargs) + + assert (None if stamped is None else dict(stamped)) == expected + def test_restoring_the_callers_ceiling_reads_the_stamp_and_replaces_every_carrier(self): stamped: dict = {"max_output_tokens": 500, "metadata": {"_client_output_ceiling": {"max_tokens": 8192}}} - Router._restore_client_ceiling_no_tier_pins(stamped) + Router._restore_client_output_ceiling(stamped) assert {k: v for k, v in stamped.items() if k != "metadata"} == {"max_tokens": 8192} coerced: dict = { "max_tokens": 64000, "metadata": {"_client_output_ceiling": {"max_tokens": "8192", "max_completion_tokens": 100.0}}, } - Router._restore_client_ceiling_no_tier_pins(coerced) + Router._restore_client_output_ceiling(coerced) assert {k: v for k, v in coerced.items() if k != "metadata"} == { "max_tokens": 8192, "max_completion_tokens": 100, } unstamped: dict = {"max_tokens": 64000, "metadata": {}} - Router._restore_client_ceiling_no_tier_pins(unstamped) + Router._restore_client_output_ceiling(unstamped) assert unstamped["max_tokens"] == 64000 @pytest.mark.asyncio async def test_pinning_stamps_the_callers_carriers_once(self): router = self._router() request_kwargs: dict = {"max_completion_tokens": 8192} + router._stamp_client_output_ceiling(request_kwargs) + assert stamped_client_ceiling(request_kwargs) == {"max_completion_tokens": 8192} first = router._pin_tier_params_onto_request( model="big", tier_litellm_params={"max_tokens": 64000}, request_kwargs=request_kwargs, responses_call=False diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 72a08c53267..b59aa051e38 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -452,7 +452,7 @@ export interface ComplexityRouterConfigValue { route_housekeeping_to_cheapest_tier?: boolean; housekeeping_patterns?: string[]; reminder_markers?: ReminderMarkerPair[]; - max_tokens_from_tier_model?: boolean; + max_tokens_from_tier_model?: boolean | "clamp"; classifier_plugin_timeout_ms?: number; } diff --git a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx index 6e7b1489a4c..f596ecfdc62 100644 --- a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx @@ -1,7 +1,28 @@ import { Switch } from "@/components/ui/switch"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import React from "react"; import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +const OUTPUT_TOKEN_POLICIES = [ + { + value: true, + label: "Use model ceiling", + description: "Replace the caller's output limit with the selected model's ceiling. This is the default.", + }, + { + value: "clamp", + label: "Respect caller budget", + description: + "Use the smaller of the caller's output limit and the model ceiling. Use the ceiling when no limit is sent.", + }, + { + value: false, + label: "Pass through caller limit", + description: "Forward the caller's output limit unchanged.", + }, +] as const; + const ResponseFormatControls: React.FC<{ value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; @@ -18,18 +39,31 @@ const ResponseFormatControls: React.FC<{ Return the resolved underlying model name in responses instead of the autorouter alias. -
- onChange({ ...value, max_tokens_from_tier_model: enabled })} - aria-label="Cap max_tokens at the tier model's output ceiling" - /> - Cap max_tokens at the tier model's output ceiling +
+ Output token limit + { + if (typeof mode === "boolean" || mode === "clamp") { + onChange({ ...value, max_tokens_from_tier_model: mode }); + } + }} + aria-label="Output token limit" + > + {OUTPUT_TOKEN_POLICIES.map((policy) => ( + + ))} + + + Explicit tier output limits take precedence. Unknown model ceilings keep the caller's limit. +
- - Replace the caller's max_tokens with the routed tier model's output ceiling so one client value fits - every tier. Off forwards the caller's value unchanged. - ); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 89d419d0487..b25baa0d1e4 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -1315,6 +1315,17 @@ describe("buildComplexityRouterConfig with an edited tier set", () => { expect(payload.fallback_tier).toBe("CASUAL"); }); + it("keeps clamp and explicit output limits on custom tiers", () => { + const payload = build({ + maxTokensFromTierModel: "clamp", + tierModelParams: { sec: { "o1-preview": { max_completion_tokens: 700 } } }, + }); + expect(payload.max_tokens_from_tier_model).toBe("clamp"); + expect(payload.tier_model_configs?.SECURITY_REVIEW).toEqual([ + { model_name: "o1-preview", litellm_params: { max_completion_tokens: 700 } }, + ]); + }); + it("forces the LLM classifier even when the form still holds heuristic, which the backend rejects", () => { expect(build({ classifierType: "heuristic" }).classifier_type).toBe("llm"); }); @@ -1595,6 +1606,15 @@ describe("classifier vision wire payload", () => { }); describe("advanced complexity router fields", () => { + it.each([undefined, true, false, "clamp"] as const)("serializes output token policy %s", (mode) => { + const payload = buildComplexityRouterConfig({ ...baseParams, maxTokensFromTierModel: mode }); + if (mode === false || mode === "clamp") { + expect(payload.max_tokens_from_tier_model).toBe(mode); + } else { + expect(payload).not.toHaveProperty("max_tokens_from_tier_model"); + } + }); + it("normalizes lists, reminder markers, and explicit false values", () => { const payload = buildComplexityRouterConfig({ ...baseParams, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 09bf0ea41a9..511a1653aae 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -226,7 +226,7 @@ export interface BuildComplexityRouterConfigParams { routeHousekeepingToCheapestTier?: boolean; housekeepingPatterns?: string[]; reminderMarkers?: ReminderMarkerPair[]; - maxTokensFromTierModel?: boolean; + maxTokensFromTierModel?: ComplexityRouterConfigValue["max_tokens_from_tier_model"]; classifierPluginTimeoutMs?: number; } @@ -310,7 +310,7 @@ export interface ComplexityRouterConfigPayload { route_housekeeping_to_cheapest_tier?: boolean; housekeeping_patterns?: string[]; reminder_markers?: ReminderMarkerPair[]; - max_tokens_from_tier_model?: boolean; + max_tokens_from_tier_model?: ComplexityRouterConfigValue["max_tokens_from_tier_model"]; classifier_plugin_timeout_ms?: number; } @@ -827,7 +827,9 @@ export const buildComplexityRouterConfig = ({ ...cleanedLists, ...(routeHousekeepingToCheapestTier === false && { route_housekeeping_to_cheapest_tier: false }), ...(cleanedReminderMarkers && cleanedReminderMarkers.length > 0 && { reminder_markers: cleanedReminderMarkers }), - ...(maxTokensFromTierModel === false && { max_tokens_from_tier_model: false }), + ...((maxTokensFromTierModel === false || maxTokensFromTierModel === "clamp") && { + max_tokens_from_tier_model: maxTokensFromTierModel, + }), ...(hasValidCustomClassifierTimeout && { classifier_plugin_timeout_ms: classifierPluginTimeoutMs }), ...scorerKnobs, }; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 9445e18913e..6238e30ce07 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -47,6 +47,29 @@ const hydratedState: KeywordMatchingState = { matchThreshold: 0.72, }; +describe("output token policy edits", () => { + it.each([undefined, true, false, "clamp"] as const)("preserves and edits stored mode %s", (mode) => { + const stored = { + ...STORED, + ...(mode !== undefined && { max_tokens_from_tier_model: mode }), + tier_model_configs: { SIMPLE: [{ model_name: "gpt-4o-mini", litellm_params: { max_tokens: 700 } }] }, + }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + expect(hydrated.max_tokens_from_tier_model).toBe(mode); + const unchanged = buildUpdatedComplexityRouterConfig(stored, { ...hydrated, return_raw_model_name: true }); + expect(unchanged.max_tokens_from_tier_model).toBe(mode); + expect("max_tokens_from_tier_model" in unchanged).toBe(mode !== undefined); + for (const nextMode of [true, false, "clamp"] as const) { + const saved = buildUpdatedComplexityRouterConfig(stored, { + ...hydrated, + max_tokens_from_tier_model: nextMode, + }); + expect(saved.max_tokens_from_tier_model).toBe(nextMode === true && mode !== true ? undefined : nextMode); + expect(saved.tier_model_configs).toEqual(stored.tier_model_configs); + } + }); +}); + describe("buildUpdatedComplexityRouterConfig keyword matching", () => { it.each([undefined, false, true])("preserves cache settings through edit and save: %s", (enabled) => { const stored = { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx index 72d557d1f8f..e8c9288cfc5 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx @@ -472,9 +472,17 @@ describe("EditAutoRouterModal advanced field round trips", () => { modelPatchUpdateCall.mockClear(); }); - it("hydrates housekeeping and reminder fields, then omits the default max-token value after editing", async () => { + it("reopens clamp and saves the selected output policy without changing other advanced settings", async () => { const user = userEvent.setup(); - renderAdvancedModal(); + renderAdvancedModal({ + modelData: { + ...MODEL_DATA, + litellm_params: { + ...MODEL_DATA.litellm_params, + complexity_router_config: { ...storedAdvancedConfig, max_tokens_from_tier_model: "clamp" }, + }, + }, + }); openAutoRouterAdvanced("Housekeeping Routing"); expect(screen.getByRole("switch", { name: "Route housekeeping calls to the cheapest tier" })).not.toBeChecked(); @@ -485,13 +493,17 @@ describe("EditAutoRouterModal advanced field round trips", () => { expect(screen.getByLabelText("Closing tag")).toHaveValue(""); openAutoRouterAdvanced("Response Format"); - const maxTokensSwitch = screen.getByRole("switch", { name: "Cap max_tokens at the tier model's output ceiling" }); - await user.click(maxTokensSwitch); + expect(screen.getByRole("radio", { name: /^Respect caller budget/ })).toBeChecked(); + await user.click(screen.getByRole("radio", { name: /^Pass through caller limit/ })); + expect(screen.getByRole("radio", { name: /^Pass through caller limit/ })).toBeChecked(); + await user.click(screen.getByRole("radio", { name: /^Use model ceiling/ })); + expect(screen.getByRole("radio", { name: /^Use model ceiling/ })).toBeChecked(); + await user.click(screen.getByRole("radio", { name: /^Respect caller budget/ })); await waitFor(() => expect(screen.getByRole("button", { name: /save changes/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalledOnce()); - expect(savedConfig()).not.toHaveProperty("max_tokens_from_tier_model"); + expect(savedConfig().max_tokens_from_tier_model).toBe("clamp"); expect(savedConfig()).toMatchObject({ route_housekeeping_to_cheapest_tier: false, housekeeping_patterns: ["conversation title"], diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts index 233b5aed922..75fd6e15bc1 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts @@ -184,7 +184,8 @@ export const hydrateComplexityRouterConfig = ( ? parsedConfig.reminder_markers.filter(isReminderMarkerPair) : undefined, max_tokens_from_tier_model: - typeof parsedConfig.max_tokens_from_tier_model === "boolean" + typeof parsedConfig.max_tokens_from_tier_model === "boolean" || + parsedConfig.max_tokens_from_tier_model === "clamp" ? parsedConfig.max_tokens_from_tier_model : undefined, classifier_plugin_timeout_ms: diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index b398780d14e..f4e45e473df 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -877,6 +877,13 @@ describe("autorouter_presets", () => { }); describe("buildPresetPrefill", () => { + it.each([undefined, true, false, "clamp"] as const)("preserves a preset's output token policy %s", (mode) => { + const preset = getPresetByKey("anthropic_family")!; + const config = { ...preset.complexity_router_config, max_tokens_from_tier_model: mode }; + const prefill = buildPresetPrefill(config, groupsOnly(getRequiredModelsInPreset(preset))); + expect(prefill.complexityRouterConfig.max_tokens_from_tier_model).toBe(mode); + }); + it("preserves Laya settings and drops inactive classifier settings when prefilling", () => { const config = { tiers: { SIMPLE: ["fast"], MEDIUM: [], COMPLEX: [], REASONING: [] }, diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index 0a83a06050b..fd400a8163a 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -364,6 +364,7 @@ export const buildPresetPrefill = ( tier_distance_penalty: config.tier_distance_penalty, adaptive_eligible: config.adaptive_eligible, return_raw_model_name: config.return_raw_model_name, + max_tokens_from_tier_model: config.max_tokens_from_tier_model, dimension_weights: hydrateDimensionWeights(config.dimension_weights), custom_dimensions: hydrateCustomDimensions(config.custom_dimensions), tier_boundaries: hydrateTierBoundaries(config.tier_boundaries), diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d03f4b4cffd..e3f80b37498 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -42620,10 +42620,10 @@ export interface components { match_threshold: number; /** * Max Tokens From Tier Model - * @description Set max_tokens on every routed request to the output ceiling of the tier model it lands on, replacing whatever the caller sent. A caller behind an auto-router cannot pick one value that fits every tier: the smallest tier's ceiling starves a bigger tier's thinking budget, and a bigger tier's ceiling is rejected by the smallest. The ceiling is the smallest max_output_tokens across the tier model's deployments, read from each deployment's model_info and then the model cost map; a tier model with a deployment whose ceiling is unknown keeps the caller's value. A max_tokens, max_completion_tokens or max_output_tokens in the tier's own litellm_params still wins. Set false to forward the caller's value unchanged. + * @description Set max_tokens on every routed request to the output ceiling of the tier model it lands on, replacing whatever the caller sent. A caller behind an auto-router cannot pick one value that fits every tier: the smallest tier's ceiling starves a bigger tier's thinking budget, and a bigger tier's ceiling is rejected by the smallest. The ceiling is the smallest max_output_tokens across the tier model's deployments, read from each deployment's model_info and then the model cost map; a tier model with a deployment whose ceiling is unknown keeps the caller's value. A max_tokens, max_completion_tokens or max_output_tokens in the tier's own litellm_params still wins. Set "clamp" to keep a smaller positive caller value and cap a larger one at the ceiling. Set false to forward the caller's value unchanged. * @default true */ - max_tokens_from_tier_model: boolean; + max_tokens_from_tier_model: boolean | "clamp"; /** * Modality Pin Override * @description Let modality_routing replace a kept session-affinity pin on the turns that carry an image. Without this, a session pinned to a text-only model fails every image turn with a provider 400, since the pin is exempt from the modality gate. When enabled, such a turn routes to a capable model for that request only and the stored pin is left untouched, so the next text turn replays the session's own model; the override is reported as cause modality_pin_override and is never itself pinned. Inert unless modality_routing is also enabled. From c23cff7e459aa83373209e273179481227155005 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 27 Oct 2026 17:30:48 -0400 Subject: [PATCH 07/83] fix(mcp): enforce deny-first tool policy consistently (#48625) --- .../proxy/mcp_server/mcp_server_manager.py | 25 ++++- litellm/proxy/mcp_server/operations.py | 27 ++--- litellm/proxy/mcp_server/utils.py | 9 ++ .../mcp_server/test_mcp_server_manager.py | 102 ++++++++++++------ .../test_mcp_server_tool_calls_and_headers.py | 35 +++--- 5 files changed, 129 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index 8e613a8a6e8..f7146ac9428 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -1705,6 +1705,24 @@ def warn_on_shared_identifier_prefixes(servers: Iterable[MCPServer]) -> None: ) +def _warn_overlapping_tool_policies(server: MCPServer) -> None: + overlapping_tools: Final = tuple( + entry + for entry in server.allowed_tools or () + if any( + match_known_tool_name(name, server, server.disallowed_tools or ()) is not None + for name in (entry, strip_known_server_prefix(entry, server)) + ) + ) + if overlapping_tools: + verbose_logger.warning( + "MCP server %s has entries in both allowed_tools and disallowed_tools: %s. " + "disallowed_tools takes precedence; matching tools will be blocked.", + server.name, + ", ".join(overlapping_tools), + ) + + def _warn_legacy_delegate_auth_if_applicable(server: MCPServer, *, source: str) -> None: """Direct legacy delegated OAuth configurations to the admitted replacement.""" if server.auth_type != MCPAuth.oauth2: @@ -2854,6 +2872,7 @@ class MCPServerManager: oauth_identity_binding=server_config.get("oauth_identity_binding", None), ) self.assign_unique_short_prefix(new_server) + _warn_overlapping_tool_policies(new_server) _warn_legacy_delegate_auth_if_applicable(new_server, source="config") _warn_config_id_jag_server_outruns_sso(new_server) self.invalidate_server_definition_caches(server_id) @@ -5943,12 +5962,10 @@ class MCPServerManager: cut a second boundary out of a native name that opens with the server prefix. """ from litellm.proxy.mcp_server.utils import ( - server_applies_tool_allowlist, + is_tool_allowed_for_server, ) - if server_applies_tool_allowlist(server): - return match_known_tool_name(tool_name, server, server.allowed_tools or ()) is not None - return match_known_tool_name(tool_name, server, server.disallowed_tools or ()) is None + return is_tool_allowed_for_server(tool_name, server) def validate_allowed_params(self, tool_name: str, arguments: _ToolArguments, server: MCPServer) -> None: """ diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index 763e3578a96..4c2cbe2dfbf 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -800,28 +800,17 @@ def filter_tools_by_allowed_tools( Filtered list of tools """ from litellm.proxy.mcp_server.utils import ( + is_tool_allowed_for_server, server_applies_tool_allowlist, ) - tools_to_return = tools - - # Filter by allowed_tools (whitelist) - if server_applies_tool_allowlist(mcp_server): - if not mcp_server.allowed_tools: - return [] - tools_to_return = [ - tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools, mcp_server) - ] - - # Filter by disallowed_tools (blacklist) - if mcp_server.disallowed_tools: - tools_to_return = [ - tool - for tool in tools_to_return - if not _tool_name_matches(tool.name, mcp_server.disallowed_tools, mcp_server) - ] - - return tools_to_return + if not server_applies_tool_allowlist(mcp_server) and not mcp_server.disallowed_tools: + return tools + return [ + tool + for tool in tools + if is_tool_allowed_for_server(strip_known_server_prefix(tool.name, mcp_server), mcp_server) + ] def apply_display_name_overrides( diff --git a/litellm/proxy/mcp_server/utils.py b/litellm/proxy/mcp_server/utils.py index ee8f3ea3b30..3f72804ff35 100644 --- a/litellm/proxy/mcp_server/utils.py +++ b/litellm/proxy/mcp_server/utils.py @@ -266,6 +266,15 @@ def server_applies_tool_allowlist(mcp_server: object) -> bool: return is_server_tool_allowlist_enforced(mcp_server) or bool(allowed_tools) +def is_tool_allowed_for_server(tool_name: str, server: MCPServer) -> bool: + """Evaluate server tool policy for a bare name, with explicit denials taking precedence.""" + if match_known_tool_name(tool_name, server, server.disallowed_tools or ()) is not None: + return False + return not server_applies_tool_allowlist(server) or ( + match_known_tool_name(tool_name, server, server.allowed_tools or ()) is not None + ) + + def validate_and_normalize_mcp_server_payload(payload: McpServerPayloadLike) -> None: """ Validate and normalize MCP server payload fields (server_name, alias, and diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py index c76326c51f6..c42579dc7f7 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py @@ -71,6 +71,7 @@ from litellm.proxy.mcp_server.legacy_callbacks import ElicitationCallback from litellm.proxy.mcp_server.mcp_context import active_mcp_request_ctx_var from litellm.proxy.mcp_server.mcp_server_manager import ( ListedToolsCaller, + MCPServerConfig, MCPServerManager, _deserialize_json_dict, _deserialize_json_list, @@ -1076,6 +1077,50 @@ class TestMCPServerManager: assert server.alias == "friendly_alias" assert server.server_name == "validserver" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "allowed_entry,denied_entry,expect_warning", + [ + ("read_status", "read_status", True), + ("policy-read_status", "read_status", True), + ("read_status", "policy-read_status", True), + ("policy-read_status", "policy-read_status", True), + ("read_status", "READ_STATUS", False), + ("read_status", "delete_records", False), + ], + ) + async def test_config_tool_policy_overlap_warns_without_changing_lists( + self, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + allowed_entry: str, + denied_entry: str, + expect_warning: bool, + ) -> None: + monkeypatch.delenv("LITELLM_ENABLE_MCP_STDIO", raising=False) + manager: Final = MCPServerManager() + config: Final[dict[str, MCPServerConfig]] = { + "policy_server": { + "alias": "policy", + "transport": MCPTransport.stdio, + "command": "python", + "allowed_tools": [allowed_entry], + "disallowed_tools": [denied_entry], + } + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await manager.load_servers_from_config(config) + + server: Final = next(iter(manager.config_mcp_servers.values())) + warnings: Final = tuple( + message for message in caplog.messages if "both allowed_tools and disallowed_tools" in message + ) + assert len(warnings) == int(expect_warning) + assert all("will be blocked" in message for message in warnings) + assert server.allowed_tools == [allowed_entry] + assert server.disallowed_tools == [denied_entry] + def _oauth2_config(self, **overrides): base = { "url": "https://example.com/mcp", @@ -6431,53 +6476,42 @@ class TestMCPServerManager: ) @pytest.mark.asyncio - async def test_pre_call_tool_check_allowed_tools_takes_precedence(self): - """Test that allowed_tools list takes precedence over disallowed_tools list""" - manager = MCPServerManager() - - # Create server with both allowed_tools and disallowed_tools - # Note: The logic in check_allowed_or_banned_tools prioritizes allowed_tools - server = MCPServer( + @pytest.mark.parametrize("tool_name,denied", [("tool1", False), ("tool2", True), ("tool3", True)]) + async def test_pre_call_tool_check_disallowed_tools_takes_precedence( + self, tool_name: str, denied: bool + ) -> None: + manager: Final = MCPServerManager() + server: Final = MCPServer( server_id="test-server", name="test-server", transport=MCPTransport.stdio, allowed_tools=["tool1", "tool2"], - disallowed_tools=["tool2", "tool3"], # tool2 is in both lists + disallowed_tools=["tool2", "tool3"], ) - - # Mock dependencies - set object_permission and object_permission_id to None - # so permission checks return None (no restrictions) user_api_key_auth: Final = UserAPIKeyAuth() - proxy_logging_obj = _mock_proxy_logging() - - # Mock the async methods that pre_call_tool_check calls - proxy_logging_obj.create_mcp_request_object_from_kwargs = MagicMock(return_value={}) - proxy_logging_obj.convert_mcp_to_llm_format = MagicMock(return_value={}) - proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) - - # tool2 should be allowed since it's in allowed_tools (takes precedence) - await manager.pre_call_tool_check( - name="tool2", - arguments={"param": "value"}, - server_name="test-server", - user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, - server=server, - ) - - # tool3 should be blocked since it's not in allowed_tools - with pytest.raises(HTTPException) as exc_info: + if not denied: await manager.pre_call_tool_check( - name="tool3", - arguments={"param": "value"}, + name=tool_name, + arguments={}, server_name="test-server", user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=None, + server=server, + ) + return + + with pytest.raises(HTTPException) as exc_info: + await manager.pre_call_tool_check( + name=tool_name, + arguments={}, + server_name="test-server", + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=None, server=server, ) assert exc_info.value.status_code == 403 - assert "Tool tool3 is not allowed for server test-server" in exc_info.value.detail["error"] + assert f"Tool {tool_name} is not allowed for server test-server" in exc_info.value.detail["error"] async def test_get_tools_from_server_add_prefix(self): """Verify _get_tools_from_server respects add_prefix True/False.""" diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py index aa61deaf554..18907228b00 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -10941,7 +10941,7 @@ class TestListFiltersHonorThePrefixBoundary: assert _tool_name_matches("eiG-read_wiki_contents", ["deepwiki_cfg-read_wiki_contents"], server) - def test_discovery_hides_exactly_what_dispatch_refuses(self, monkeypatch): + def test_discovery_hides_exactly_what_dispatch_refuses(self, monkeypatch: pytest.MonkeyPatch) -> None: """Both halves of one decision, driven through both production paths. A spelling the blocklist enforces but the filter misses leaves a blocked @@ -10963,7 +10963,10 @@ class TestListFiltersHonorThePrefixBoundary: monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") - def _server(**overrides): + def _server( + allowed_tools: list[str] | None = None, + disallowed_tools: list[str] | None = None, + ) -> MCPServer: return MCPServer( server_id=self.SERVER_ID, name="deepwiki_prod", @@ -10972,27 +10975,35 @@ class TestListFiltersHonorThePrefixBoundary: short_prefix="eiG", url="http://127.0.0.1:5115/mcp", transport=MCPTransport.http, - **overrides, + allowed_tools=allowed_tools, + disallowed_tools=disallowed_tools, ) - manager = MCPServerManager() + manager: Final = MCPServerManager() manager._create_prefixed_tools( [MCPTool(name="read_wiki_contents", description="", inputSchema={"type": "object"})], _server(), ) - registered = sorted(manager.tool_name_to_mcp_server_name_mapping) + registered: Final = sorted(manager.tool_name_to_mcp_server_name_mapping) assert len(registered) > 1 - published = MCPTool(name="eiG-read_wiki_contents", description="", inputSchema={"type": "object"}) + published: Final = MCPTool(name="eiG-read_wiki_contents", description="", inputSchema={"type": "object"}) for spelling in registered: - for entry, expected in ((spelling, True), (spelling.upper(), False)): - server = _server(disallowed_tools=[entry]) + policies: Final[tuple[tuple[list[str] | None, list[str], bool], ...]] = ( + (None, [spelling], True), + (None, [spelling.upper()], False), + (["read_wiki_contents"], [spelling], True), + ([spelling], ["read_wiki_contents"], True), + ([spelling], [spelling.upper()], False), + ) + for allowed_tools, disallowed_tools, expected in policies: + server: Final = _server(allowed_tools, disallowed_tools) - refused = not manager.check_allowed_or_banned_tools("read_wiki_contents", server) - hidden = filter_tools_by_allowed_tools([published], server) == [] + refused: Final = not manager.check_allowed_or_banned_tools("read_wiki_contents", server) + hidden: Final = filter_tools_by_allowed_tools([published], server) == [] - assert refused == hidden, entry - assert refused is expected, entry + assert refused == hidden, (allowed_tools, disallowed_tools) + assert refused is expected, (allowed_tools, disallowed_tools) @pytest.mark.asyncio @pytest.mark.parametrize( From ad4b35d791bd79b142f25a022e9f8a77230f0798 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 27 Oct 2026 17:46:15 -0400 Subject: [PATCH 08/83] fix(mcp): accept YAML OpenAPI documents (#48622) --- .../mcp_server/openapi_to_mcp_generator.py | 101 +++++++++++- .../mcp_server/test_mcp_server_manager.py | 37 ++++- .../mcp_server/test_openapi_spec_path_url.py | 152 +++++++++++++++++- 3 files changed, 284 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py index 96d178dbeb1..0fa904cb582 100644 --- a/litellm/proxy/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/mcp_server/openapi_to_mcp_generator.py @@ -5,15 +5,29 @@ This module is used to generate MCP tools from OpenAPI specs. import asyncio import contextvars import json +import math import os import re from collections.abc import Mapping, Sequence from pathlib import PurePosixPath -from typing import Final, Literal, Protocol, TypeAlias, TypedDict, overload +from typing import ( + ClassVar, + Final, + Literal, + Protocol, + TypeAlias, + TypedDict, + cast, # noqa: TID251 # existing OpenAPI Protocol fields have no runtime schema + overload, +) from urllib.parse import quote import httpx +import yaml +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, Required +from yaml.events import AliasEvent, ScalarEvent +from yaml.nodes import MappingNode, Node, ScalarNode from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError from litellm.proxy.mcp_server.exceptions import ( @@ -30,6 +44,18 @@ _OPENAPI_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]") OPENAPI_TOOL_NAME_MAX_LEN: Final = 128 _OPENAPI_TOOL_NAME_MAX_LEN: Final = OPENAPI_TOOL_NAME_MAX_LEN +_OPENAPI_DOCUMENT: Final = TypeAdapter(dict[str, object]) +_MAX_YAML_INT_LENGTH: Final = 1024 +# YAML 1.2 core scalars; isolate these from PyYAML's global YAML 1.1 resolver. +_YAML_SCALARS: Final = { + "tag:yaml.org,2002:null": re.compile(r"^(?:null|Null|NULL|~|)$"), + "tag:yaml.org,2002:bool": re.compile(r"^(?:true|True|TRUE|false|False|FALSE)$"), + "tag:yaml.org,2002:int": re.compile(r"^(?:[-+]?[0-9]+|0o[0-7]+|0x[0-9a-fA-F]+)$"), + "tag:yaml.org,2002:float": re.compile( + r"^(?:[-+]?(?:\.[0-9]+|[0-9]+(?:\.[0-9]*)?)(?:[eE][-+]?[0-9]+)?" + r"|[-+]?\.(?:inf|Inf|INF)|\.(?:nan|NaN|NAN))$" + ), +} def sanitize_openapi_tool_name(raw_name: str) -> str: @@ -199,6 +225,70 @@ def load_openapi_spec(filepath: str) -> _OpenAPISpec: return asyncio.run(load_openapi_spec_async(filepath)) +class _OpenAPISafeLoader(yaml.SafeLoader): + yaml_implicit_resolvers: ClassVar[Mapping[str | None, Sequence[tuple[str, re.Pattern[str]]]]] = { # pyright: ignore[reportIncompatibleVariableOverride] # PyYAML resolver tables are class state + None: list(_YAML_SCALARS.items()) + } + + def compose_node(self, parent: Node | None, index: int) -> Node | None: + event: Final[object] = self.peek_event() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # incomplete PyYAML event stubs + if isinstance(event, AliasEvent): + raise ValueError("YAML aliases are not supported in OpenAPI specs") # noqa: TRY004 # callers handle malformed documents as ValueError + node: Final = super().compose_node(parent, index) + # PyYAML otherwise implicitly resolves even an explicit nonspecific ! tag. + if isinstance(event, ScalarEvent) and event.tag == "!" and isinstance(node, ScalarNode): + node.tag = "tag:yaml.org,2002:str" + return node + + def construct_object(self, node: Node, deep: bool = False) -> object: + pattern: Final = _YAML_SCALARS.get(node.tag) + if pattern is not None: + if not isinstance(node, ScalarNode): + raise ValueError("Invalid YAML OpenAPI spec") + scalar: Final[str] = node.value + if node.tag == "tag:yaml.org,2002:int" and len(scalar) > _MAX_YAML_INT_LENGTH: + raise ValueError("YAML integer is too long in OpenAPI spec") + if pattern.fullmatch(scalar) is None: + raise ValueError("Invalid YAML OpenAPI spec") + if node.tag == "tag:yaml.org,2002:int": + # YAML 1.2 decimal values with leading zeroes are not octal. + return int(scalar, 0 if scalar.startswith(("0o", "0x")) else 10) + if node.tag == "tag:yaml.org,2002:float": + value: Final = self.construct_yaml_float(node) + if not math.isfinite(value): + raise ValueError("Invalid YAML OpenAPI spec") + return value + elif node.tag not in {"tag:yaml.org,2002:str", "tag:yaml.org,2002:seq", "tag:yaml.org,2002:map"} or ( + node.tag == "tag:yaml.org,2002:str" and not isinstance(node, ScalarNode) + ): + raise ValueError("Invalid YAML OpenAPI spec") + if isinstance(node, MappingNode): + pairs: Final[Sequence[tuple[Node, Node]]] = node.value + if any(not isinstance(key, ScalarNode) or key.tag != "tag:yaml.org,2002:str" for key, _ in pairs): + raise ValueError("Invalid YAML OpenAPI spec") + return super().construct_object( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # incomplete PyYAML stubs + node, deep + ) + + +def _validate_openapi_spec(value: object) -> _OpenAPISpec: + try: + document: Final = _OPENAPI_DOCUMENT.validate_python(value) + except ValidationError as exc: + raise ValueError("OpenAPI spec must be a JSON or YAML object") from exc + return cast(_OpenAPISpec, document) # cast-ok: preserve existing lazy OpenAPI Protocol validation + + +def _parse_openapi_spec(content: str | bytes) -> _OpenAPISpec: + try: + return _validate_openapi_spec(json.loads(content)) + except json.JSONDecodeError: + try: + return _validate_openapi_spec(yaml.load(content, Loader=_OpenAPISafeLoader)) + except (yaml.YAMLError, RecursionError) as exc: + raise ValueError("Invalid YAML OpenAPI spec") from exc + + async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None) -> _OpenAPISpec: if filepath.startswith("http://") or filepath.startswith("https://"): client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) @@ -208,14 +298,17 @@ async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None else await async_safe_get(client, filepath, max_response_bytes=max_bytes) ) r.raise_for_status() - return r.json() + return await asyncio.to_thread(_parse_openapi_spec, r.content) # fallback: local file - # Local filesystem path + return await asyncio.to_thread(_load_openapi_file, filepath) + + +def _load_openapi_file(filepath: str) -> _OpenAPISpec: if not os.path.exists(filepath): raise FileNotFoundError(f"OpenAPI spec not found at {filepath}") with open(filepath, "r", encoding="utf-8") as f: - return json.load(f) + return _parse_openapi_spec(f.read()) def get_base_url(spec: _OpenAPISpec, spec_path: str | None = None) -> str: diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py index c42579dc7f7..ac79c10869b 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py @@ -5405,7 +5405,42 @@ class TestMCPServerManager: ( httpx.Response(200, text="secret invalid JSON body"), "unhealthy", - "OpenAPI specification could not be loaded (JSONDecodeError)", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="paths: [secret invalid YAML body"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="paths: {}\nx: \x01"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: &a []\ny: *a"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!bool invalid"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!int ''"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!float ''"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", + ), + ( + httpx.Response(200, text="x: !!timestamp invalid"), + "unhealthy", + "OpenAPI specification could not be loaded (ValueError)", ), ], ) diff --git a/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py b/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py index aac85160eba..0f29a3bf20d 100644 --- a/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py +++ b/tests/unit/proxy/mcp_server/test_openapi_spec_path_url.py @@ -1,9 +1,15 @@ from __future__ import annotations -from typing import Any, Dict +import asyncio +import json +import threading +from pathlib import Path +from typing import Any, Dict, Final import httpx import pytest +import yaml +from respx import MockRouter from litellm.proxy.mcp_server import openapi_to_mcp_generator as gen @@ -98,3 +104,147 @@ def test_load_openapi_spec_supports_local_file_path( spec = gen.load_openapi_spec(str(p)) assert spec == expected + + +@pytest.mark.parametrize("source", ["url", "file"]) +@pytest.mark.parametrize("suffix", [".yaml", ".json", ""]) +@pytest.mark.parametrize("document_format", ["json", "yaml"]) +@pytest.mark.parametrize("use_async", [False, True]) +def test_openapi_formats_preserve_document_across_sources( + tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + source: str, suffix: str, document_format: str, use_async: bool, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + expected: Final = { + "openapi": "3.0.0", + "paths": {"/ping": {"get": { + "operationId": "on", + "parameters": [ + {"name": "on", "in": "query", "schema": {"type": "string"}}, + {"name": "off", "in": "query", "schema": {"type": "string"}}, + ], + }}}, + "x-values": [100.0, "on", "off", "yes", "no", "2026-01-01", "1:20", True, False, None, 12, 10, 10], + } + content: Final = ( + json.dumps(expected).replace("100.0", "1e2") + if document_format == "json" + else "openapi: 3.0.0\npaths:\n /ping:\n get:\n operationId: on\n" + " parameters:\n" + " - {name: on, in: query, schema: {type: string}}\n" + " - {name: off, in: query, schema: {type: string}}\n" + "x-values: [1e2, on, off, yes, no, 2026-01-01, 1:20, TRUE, False, null, 012, 0o12, 0xA]\n" + ) + path: Final = tmp_path / ("spec" + suffix) + path.write_text(content, encoding="utf-8") + target: Final = str(path) if source == "file" else "https://93.184.216.34/spec" + suffix + if source == "url": + respx_mock.get(target).respond(200, text=content, headers={"content-type": "text/plain"}) + + spec: Final = asyncio.run(gen.load_openapi_spec_async(target)) if use_async else gen.load_openapi_spec(target) + + assert spec == expected + operation: Final = spec["paths"]["/ping"]["get"] + assert gen.sanitize_openapi_tool_name(operation["operationId"]) == "on" + assert gen.build_input_schema(operation) == { + "type": "object", + "properties": {"on": {"type": "string", "description": ""}, "off": {"type": "string", "description": ""}}, + "required": [], + } + + +_LONG_YAML_INTEGER: Final = "1" * 1025 + + +@pytest.mark.parametrize("content,message", [ + ("", "OpenAPI spec must be a JSON or YAML object"), + ("null", "OpenAPI spec must be a JSON or YAML object"), + ("[]", "OpenAPI spec must be a JSON or YAML object"), + ("- item", "OpenAPI spec must be a JSON or YAML object"), + ("plain text", "OpenAPI spec must be a JSON or YAML object"), + ("paths: [", "Invalid YAML OpenAPI spec"), + ("paths: {}\nx: \x01", "Invalid YAML OpenAPI spec"), + ("paths: " + "[" * 600 + "0" + "]" * 600, "Invalid YAML OpenAPI spec"), + ("x: !!python/object/apply:builtins.str [unsafe]", "Invalid YAML OpenAPI spec"), + ("x: &a []\ny: *a", "YAML aliases are not supported in OpenAPI specs"), + ("x: " + _LONG_YAML_INTEGER, "YAML integer is too long in OpenAPI spec"), + ("x: !!int " + _LONG_YAML_INTEGER, "YAML integer is too long in OpenAPI spec"), + ("x: !!int {=: '" + _LONG_YAML_INTEGER + "'}", "Invalid YAML OpenAPI spec"), + ("x: !!int [1]", "Invalid YAML OpenAPI spec"), + ("x: !!bool invalid", "Invalid YAML OpenAPI spec"), + ("x: !!bool []", "Invalid YAML OpenAPI spec"), + ("x: !!str []", "Invalid YAML OpenAPI spec"), + ("x: !!int ''", "Invalid YAML OpenAPI spec"), + ("x: !!float ''", "Invalid YAML OpenAPI spec"), + ("x: !!float {=: 1.2}", "Invalid YAML OpenAPI spec"), + ("x: !!timestamp invalid", "Invalid YAML OpenAPI spec"), + ("x: !!timestamp 2026-01-01", "Invalid YAML OpenAPI spec"), + ("x: !!set {a: null}", "Invalid YAML OpenAPI spec"), + ("x: !!binary YQ==", "Invalid YAML OpenAPI spec"), + ("x: {1: value}", "Invalid YAML OpenAPI spec"), + ("x: {[key]: value}", "Invalid YAML OpenAPI spec"), + ("x: .inf", "Invalid YAML OpenAPI spec"), + ("x: .nan", "Invalid YAML OpenAPI spec"), + ("x: 1e999", "Invalid YAML OpenAPI spec"), +], ids=[ + "empty", "null", "json-array", "yaml-sequence", "scalar", "syntax", "reader", "nesting", + "python-tag", "alias", "implicit-int", "explicit-int", "mapping-int", "sequence-int", + "invalid-bool", "sequence-bool", "sequence-str", "empty-int", "empty-float", "mapping-float", "invalid-timestamp", + "timestamp", "set", "binary", "non-string-key", "complex-key", "infinity", "nan", "float-overflow", +]) +def test_openapi_loader_rejects_invalid_or_unsafe_documents( + tmp_path: Path, content: str, message: str, +) -> None: + path: Final = tmp_path / "spec.yaml" + path.write_text(content, encoding="utf-8") + + with pytest.raises(ValueError, match="^" + message + "$"): + gen.load_openapi_spec(str(path)) + + +def test_openapi_yaml_integer_boundary_preserves_quoted_strings(tmp_path: Path) -> None: + integer: Final = "1" * 1024 + path: Final = tmp_path / "spec.yaml" + path.write_text("x-integer: " + integer + "\nx-string: '" + integer + "1'", encoding="utf-8") + + assert gen.load_openapi_spec(str(path)) == {"x-integer": int(integer), "x-string": integer + "1"} + + +def test_openapi_json_integer_preserved_even_with_yaml_extension(tmp_path: Path) -> None: + path: Final = tmp_path / "spec.yaml" + path.write_text('{"x-integer":' + _LONG_YAML_INTEGER + "}", encoding="utf-8") + + assert gen.load_openapi_spec(str(path)) == {"x-integer": int(_LONG_YAML_INTEGER)} + + +@pytest.mark.parametrize("value", ["! " + _LONG_YAML_INTEGER, "! '" + _LONG_YAML_INTEGER + "'"], ids=["plain", "quoted"]) +def test_openapi_nonspecific_tag_disables_scalar_resolution(tmp_path: Path, value: str) -> None: + path: Final = tmp_path / "spec.yaml" + path.write_text("x: " + value, encoding="utf-8") + + assert gen.load_openapi_spec(str(path)) == {"x": _LONG_YAML_INTEGER} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", ["url", "file"]) +async def test_openapi_yaml_parses_off_the_event_loop( + tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, source: str, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + event_loop_thread: Final = threading.get_ident() + parser_threads: Final[list[int]] = [] + original_load: Final = yaml.load + + def observed_load(stream: str | bytes, Loader: type[yaml.SafeLoader]) -> object: + parser_threads.append(threading.get_ident()) + return original_load(stream, Loader=Loader) + + monkeypatch.setattr(yaml, "load", observed_load) + path: Final = tmp_path / "spec.yaml" + path.write_text("paths: {}", encoding="utf-8") + target: Final = str(path) if source == "file" else "https://93.184.216.34/spec.yaml" + if source == "url": + respx_mock.get(target).respond(200, text="paths: {}") + + assert await gen.load_openapi_spec_async(target) == {"paths": {}} + assert len(parser_threads) == 1 and parser_threads[0] != event_loop_thread From 2a9c63018242f014b5a7416a1e7b065144d10b1e Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 27 Oct 2026 17:47:04 -0400 Subject: [PATCH 09/83] fix(proxy): release concurrency slots when synthetic streams complete (#48628) --- litellm/proxy/common_request_processing.py | 7 +- litellm/proxy/utils.py | 21 +--- .../agent_endpoints/test_a2a_endpoints.py | 3 +- .../test_streaming_model_restamp.py | 53 ++++---- .../hooks/test_parallel_request_limiter_v3.py | 18 ++- tests/unit/proxy/test_budget_reservation.py | 3 +- .../proxy/test_common_request_processing.py | 117 ++++++++++++++++++ 7 files changed, 168 insertions(+), 54 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3b9c17baa02..740d8742e5a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4355,6 +4355,8 @@ class ProxyBaseLLMRequestProcessing: proxy_logging_obj: ProxyLogging | None = None, ) -> None: with anyio.CancelScope(shield=True): + if stream_completed and proxy_logging_obj is not None and user_api_key_dict is not None: + await proxy_logging_obj.arelease_max_parallel_requests(user_api_key_dict) should_record_client_disconnect: Final = client_disconnected or (not stream_completed) recorded_client_disconnect = False if should_record_client_disconnect: @@ -4369,10 +4371,7 @@ class ProxyBaseLLMRequestProcessing: # A disconnect-time success event (the deferred-guardrail flush # above, or the partial-spend billing below) releases the # request's max_parallel_requests slot through the limiter's - # own success callback. Release the slot explicitly only when - # no such event fires, so exactly one release happens; two - # concurrent releases would race and double-decrement under the - # limiter's in-memory fallback. + # own success callback. Release explicitly when no event fires. success_event_owns_slot_release = deferred_stream_logging_armed if not deferred_stream_logging_armed: success_event_owns_slot_release = await _bill_partial_streamed_spend_on_disconnect( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d3202368d76..452434e8c81 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4942,31 +4942,22 @@ class ProxyLogging: logging_obj._deferred_stream_complete_args = None return True - async def arelease_max_parallel_requests_on_disconnect( + async def arelease_max_parallel_requests( self, user_api_key_dict: UserAPIKeyAuth, ) -> None: """ - Release the api-key max_parallel_requests slot when a streaming - response is cancelled mid-flight (client disconnect) and no logging - callback fired for it. Neither the success nor failure callback runs on - the resulting CancelledError / GeneratorExit, so the pre-call +1 would - otherwise leak. - - Awaited from the shielded streaming cleanup rather than scheduled - fire-and-forget, so the caller can make it the single owner of the - release: when a disconnect-time success event does fire (partial-spend - billing or a deferred-guardrail flush), that event's own limiter - callback releases the slot and this is not called at all. Two - concurrent releases of the same acquisition would otherwise race and - double-decrement under the limiter's in-memory fallback. + Release an ended request's v3 concurrency acquisition independently of logging. + The request stash serializes release with callbacks and clears the acquisition, + so repeated cleanup cannot release another request's slot. """ limiter: Final = self.get_proxy_hook("parallel_request_limiter") if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) - _arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests_on_disconnect + arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests + _arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests async def arelease_responses_websocket_slot(self, user_api_key_dict: UserAPIKeyAuth) -> None: """Release the slot a Responses WebSocket took at admission. Only call this once that admission succeeded: diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index a2448694dd7..a7d7ac5627c 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -1058,6 +1058,7 @@ async def test_subscribe_to_task_calls_pre_call_hook(): """tasks/resubscribe must run pre_call_hook so guardrails configured on the agent are enforced before streaming begins.""" from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging agent = _make_agent_mock() mock_request = _make_request_mock("tasks/resubscribe", {"id": "task-1"}) @@ -1083,7 +1084,7 @@ async def test_subscribe_to_task_calls_pre_call_hook(): async for chunk in response: yield chunk - mock_proxy_logging = MagicMock() + mock_proxy_logging: Final = MagicMock(spec=ProxyLogging) mock_proxy_logging.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data ) diff --git a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py index b7bc670c7f8..85bb8304b95 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py +++ b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -3,15 +3,18 @@ Tests for restamping the public model on Anthropic Messages streaming chunks. """ import json -from unittest.mock import AsyncMock, MagicMock +from collections.abc import AsyncIterator, Sequence import pytest +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( AnthropicStreamModelRestamper, restamp_anthropic_stream_chunk_model, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging def _message_start_frame(model: str, line_end: str = "\n") -> bytes: @@ -22,15 +25,9 @@ def _message_start_frame(model: str, line_end: str = "\n") -> bytes: return f"event: message_start{line_end}data: {json.dumps(payload)}{line_end}{line_end}".encode() -def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock: - async def _iterator_hook(**_kwargs): - for frame in frames: - yield frame - - proxy_logging_obj = MagicMock() - proxy_logging_obj.async_post_call_streaming_iterator_hook = _iterator_hook - proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) - return proxy_logging_obj +async def _stream_frames(frames: Sequence[bytes]) -> AsyncIterator[bytes]: + for frame in frames: + yield frame def _model_from_frame(frame: bytes | str) -> str: @@ -76,15 +73,14 @@ def test_leaves_chunks_without_a_model_untouched(chunk): async def test_sse_generator_publishes_requested_model_on_message_start(): """The message_start event reports the requested model, not the provider's.""" delta_frame = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' - proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001"), delta_frame]) chunks = [ chunk async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=MagicMock(), - user_api_key_dict=MagicMock(), + response=_stream_frames((_message_start_frame("claude-haiku-4-5-20251001"), delta_frame)), + user_api_key_dict=UserAPIKeyAuth(), request_data={"model": "claude-auto-1"}, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), restamp_model="claude-auto-1", ) ] @@ -95,15 +91,13 @@ async def test_sse_generator_publishes_requested_model_on_message_start(): @pytest.mark.asyncio async def test_sse_generator_keeps_provider_model_when_restamping_is_off(): - proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001")]) - chunks = [ chunk async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=MagicMock(), - user_api_key_dict=MagicMock(), + response=_stream_frames((_message_start_frame("claude-haiku-4-5-20251001"),)), + user_api_key_dict=UserAPIKeyAuth(), request_data={"model": "claude-auto-1"}, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), ) ] @@ -175,15 +169,14 @@ def test_dict_message_start_disarms_after_restamp(): @pytest.mark.asyncio async def test_sse_generator_restamps_message_start_split_across_chunks(): frame = _message_start_frame("claude-haiku-4-5-20251001") - proxy_logging_obj = _proxy_logging_obj_streaming([frame[:30], frame[30:]]) chunks = [ chunk async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=MagicMock(), - user_api_key_dict=MagicMock(), + response=_stream_frames((frame[:30], frame[30:])), + user_api_key_dict=UserAPIKeyAuth(), request_data={"model": "claude-auto-1"}, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), restamp_model="claude-auto-1", ) ] @@ -250,15 +243,14 @@ def test_flush_disarms_the_restamper(): @pytest.mark.asyncio async def test_sse_generator_flushes_held_tail_at_end_of_stream(): unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2] - proxy_logging_obj = _proxy_logging_obj_streaming([unterminated]) chunks = [ chunk async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=MagicMock(), - user_api_key_dict=MagicMock(), + response=_stream_frames((unterminated,)), + user_api_key_dict=UserAPIKeyAuth(), request_data={"model": "claude-auto-1"}, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), restamp_model="claude-auto-1", ) ] @@ -271,15 +263,14 @@ async def test_sse_generator_flushes_held_tail_at_end_of_stream(): async def test_sse_generator_restamps_crlf_stream(): frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n' - proxy_logging_obj = _proxy_logging_obj_streaming([frame, delta]) chunks = [ chunk async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=MagicMock(), - user_api_key_dict=MagicMock(), + response=_stream_frames((frame, delta)), + user_api_key_dict=UserAPIKeyAuth(), request_data={"model": "claude-auto-1"}, - proxy_logging_obj=proxy_logging_obj, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), restamp_model="claude-auto-1", ) ] diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index b8dba5e1e68..cd0d786f841 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4378,9 +4378,12 @@ async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_hea @pytest.mark.asyncio -async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None: +@pytest.mark.parametrize("release_from", ["post_call", "stream_cleanup"]) +async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails(release_from: str) -> None: from unittest.mock import AsyncMock + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + api_key = hash_token("sk-concurrent-release") local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler( @@ -4409,9 +4412,20 @@ async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redi counter_keys=[parallel_key], ) data = {"litellm_call_id": call_id} + logger: Final = ProxyLogging(user_api_key_cache=local_cache) + logger.proxy_hook_mapping["parallel_request_limiter"] = handler post_call_task = asyncio.create_task( - handler.async_post_call_success_hook( + ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, + request_data=data, + response=None, + stream_completed=True, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=logger, + ) + if release_from == "stream_cleanup" + else handler.async_post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=ResponsesAPIResponse( diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 42b41a2e3f2..535c3da8c3a 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -4,6 +4,7 @@ import time from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -3233,7 +3234,7 @@ async def test_streaming_slow_path_processes_and_yields_chunk(spend_counter_stat async def one_chunk(user_api_key_dict, response, request_data): yield {"content": "hi"} - streaming_logging_obj = MagicMock() + streaming_logging_obj: Final = MagicMock(spec=ProxyLogging) streaming_logging_obj.async_post_call_streaming_iterator_hook = one_chunk streaming_logging_obj.async_post_call_streaming_hook = AsyncMock( side_effect=lambda **kwargs: kwargs["response"] diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 5c99e5a01ea..52816dc50a8 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -1,4 +1,5 @@ import asyncio +import contextvars import copy import datetime import json @@ -12055,6 +12056,122 @@ class TestResponseCostFromLoggingObjForCacheHits: assert cost == "" +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", ["messages", "openai"]) +@pytest.mark.parametrize("keepalive", ["disabled", "fast", "pinged"]) +async def test_completed_synthetic_stream_reuses_capacity_without_freeing_another_request( + surface: Literal["messages", "openai"], + keepalive: Literal["disabled", "fast", "pinged"], + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache, hash_token + from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse + + def clock() -> datetime.datetime: + return datetime.datetime(2026, 1, 1) + + cache: Final = DualCache() + limiter: Final = PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(cache), time_provider=clock + ) + logger: Final = ProxyLogging(user_api_key_cache=cache) + logger.proxy_hook_mapping["parallel_request_limiter"] = limiter + auth: Final = ProxyUserAPIKeyAuth(api_key=hash_token("sk-synthetic-stream"), max_parallel_requests=2) + counter: Final = f"{{api_key:{auth.api_key}}}:max_parallel_requests" + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", logger) + + async def admit(call_id: str) -> None: + await limiter.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": "test-model", "litellm_call_id": call_id}, + call_type="anthropic_messages", + ) + + async def request(call_id: str) -> None: + upstream_answers: Final = asyncio.Event() + if keepalive != "pinged": + upstream_answers.set() + + async def produce() -> StreamingResponse: + await admit(call_id) + with pytest.raises(HTTPException) as rejected: + await contextvars.Context().run(asyncio.create_task, admit("overlapping")) + assert rejected.value.status_code == 429 + await upstream_answers.wait() + body: Final = AnthropicMessagesResponse( + id=call_id, + type="message", + role="assistant", + model="test-model", + content=[{"type": "text", "text": "Search completed"}], + stop_reason="end_turn", + stop_sequence=None, + usage={"input_tokens": 0, "output_tokens": 0}, + ) + source: Final = FakeAnthropicMessagesStreamIterator(body) + data: Final = {"model": "test-model", "litellm_call_id": call_id} + stream: Final = ( + ProxyBaseLLMRequestProcessing.async_sse_data_generator(source, auth, data, logger) + if surface == "messages" + else proxy_server.async_data_generator(source, auth, data) + ) + return StreamingResponse(stream) + + response: Final = await open_sse_before_first_byte( + produce(), None if keepalive == "disabled" else 0.001 if keepalive == "pinged" else 60.0 + ) + first: Final = await anext(response.body_iterator) + if keepalive == "pinged": + assert first == SSE_COMMENT_PING_BYTES + upstream_answers.set() + frames: Final = (first, *tuple([frame async for frame in response.body_iterator])) + wire: Final = b"".join(frame.encode() if isinstance(frame, str) else frame for frame in frames) + assert b'"type": "message_stop"' in wire + + await contextvars.Context().run(asyncio.create_task, admit("still-streaming")) + for index in range(4): + await contextvars.Context().run(asyncio.create_task, request(f"completed-{index}")) + occupied: Final = await cache.async_get_cache(counter) + assert isinstance(occupied, dict) and len(occupied) == 1, occupied + + +@pytest.mark.asyncio +async def test_completed_stream_cleanup_does_not_release_an_unowned_legacy_slot() -> None: + from litellm.proxy.hooks.parallel_request_limiter import PROXY_MaxParallelRequestsHandler + from litellm.proxy.utils import InternalUsageCache, hash_token + + def clock() -> datetime.datetime: + return datetime.datetime(2026, 1, 1) + + cache: Final = DualCache() + limiter: Final = PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache), clock=clock) + logger: Final = ProxyLogging(user_api_key_cache=cache) + logger.proxy_hook_mapping["parallel_request_limiter"] = limiter + auth: Final = ProxyUserAPIKeyAuth(api_key=hash_token("sk-legacy-stream"), max_parallel_requests=1) + await limiter.async_pre_call_hook(auth, cache, {}, "acompletion") + + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=None, + request_data={}, + response=None, + stream_completed=True, + user_api_key_dict=auth, + proxy_logging_obj=logger, + ) + + with pytest.raises(HTTPException) as rejected: + await limiter.async_pre_call_hook(auth, cache, {}, "acompletion") + assert rejected.value.status_code == 429 + + def test_maybe_get_model_id_reads_the_chat_metadata_the_router_updated_in_place(): processing_obj: Final = ProxyBaseLLMRequestProcessing(data={"metadata": {"model_info": {"id": "declared-text"}}}) From a53ea5bd7784c51e4c0b726efe55de44c56c8b5d Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 27 Oct 2026 14:49:07 -0700 Subject: [PATCH 10/83] fix(proxy): keep every startup line JSON when json_logs is on (#47035) * fix(proxy): keep every startup line JSON when json_logs is on Boot-time messages the proxy printed straight to stdout (banner, feedback box, model list, cache and callback notices, policy table, migration notices, CLI launch notes) and the first records of uvicorn- or granian-spawned workers came out as plain text while json_logs was on, breaking log pipelines that parse every line as JSON. - print_startup_message() prints unchanged in plain mode and, with json_logs on, emits one INFO record on the proxy logger (ANSI stripped) that skips the logger and handler level checks like the print it replaces; show_banner() routes to it in JSON mode only - the CLI exports JSON_LOGS=true when the config sets json_logs so spawned workers log JSON from their first record - the --enforce_prisma_migration_check deprecation notice prints after the config is read so a config-only json_logs covers it Resolves LIT-1624 * fix(proxy): keep show_banner() SDK-free, pick the JSON banner in initialize() Since #46803 the `lite` shell calls show_banner() against an uninitialized litellm package, and reading litellm.json_logs inside it initialized the whole SDK at the banner (0.14 s to 0.83 s for a bare `lite` here, the seconds that PR removed on small VMs). show_banner() is back to its main version; the proxy's initialize() now emits the JSON banner record itself and calls show_banner() only in plain mode. The banner unit tests move to the initialize() call site. * fix(proxy): keep the JSON banner when an operator --log_config names the proxy logger uvicorn applies --log_config after proxy_cli has switched JSON logging on and before initialize() switches it on again through load_config(). A logging config that names "LiteLLM Proxy" removes that logger's handlers in the window between the two, which is where initialize() emits the banner and the feedback box, so at one worker both went nowhere; two-worker children import litellm after uvicorn's logging setup and were not affected. print_startup_message() now falls back to a one-off JSON handler, on the stream the JSON handler uses, when the proxy logger has no handler of its own. Live, json_logs on plus a log_config naming the proxy logger, one worker: main prints the banner and the feedback box as plaintext, e2303a3925 prints neither, this head prints both as one JSON record each. New unit test fails on e2303a3925; new integration cell test_json_logs_with_operator_log_config_keeps_every_line_json passes at one and two workers. * fix(proxy): keep a long startup listing whole in its JSON record print_startup_message ran its record through the proxy logger's filters, so in JSON mode the required-migration SQL list and the policy list were middle-truncated at MAX_STRING_LENGTH_STDOUT_LOG (4096 by default) with the litellm_truncated marker, where the prints they replace showed them whole. That cap bounds echoed request payloads; a startup message is operator text, so the record now gets the uncapped processing of a DEBUG line: secret redaction and base64 collapse still apply, the text cap does not. Red on 41f92b745a, green here: - tests/unit/test_logging.py: a listing longer than the cap comes out whole; a credential in a startup message is still redacted - tests/integration/configuration/test_json_logs_startup.py: a 60-policy config boots with json_logs on and the policy listing record (9762 chars) carries every policy with no truncation marker (4096 chars, marker, on the unfixed tree) * fix(proxy): keep the --enforce_prisma_migration_check notice on --health and --test This PR moved the deprecation notice after the config is read, so that a json_logs set only in the config applies to it; --health and --test return before that point and stopped printing it, a visible change main did not have. The notice is now also emitted before those two returns, through the same helper, so a deployment running `litellm --health --enforce_prisma_migration_check` still learns the flag is dead, as on main. A normal server start emits it once, after the config, as in the previous commit. Red on 2d7a2d6844, green here: tests/unit/proxy/test_proxy_cli.py asserts the notice on --health and on --test. * fix(proxy): keep the JSON startup records on stdout under an operator handler on the proxy logger An operator --log_config that gives "LiteLLM Proxy" a handler of its own, a NullHandler that silences it or a plain WARNING-level handler on stderr, is in force in the window where the banner, the feedback box and the cache notice are emitted: uvicorn applies it after the CLI process switched JSON on and before initialize() re-installs the JSON handler through load_config(). 41f92b745a sent the startup records through the proxy logger's handlers and fell back to a JSON stdout handler only when there were none, so under such a config the records landed in the operator's handler, as plaintext on stderr, or vanished, where the prints they replace reached stdout whatever the logging config said. print_startup_message() now writes every startup record through the proxy's own JSON handler on stdout, filters included, and never through the proxy logger's handlers; the operator's handler gets nothing from a startup message, as it got nothing from the print. The normal JSON handler is the same class, formatter and filters, so the records a pipeline reads are unchanged when no operator config is in play. Tests: tests/unit/test_logging.py gets the two handler shapes (red on f6f5957d16: the record went to stderr as plaintext or nowhere); tests/integration/configuration/ test_json_logs_startup.py gets a silencing operator --log_config cell at 1 and 2 workers (red on f6f5957d16: banner, feedback box and cache notice missing; green here). * fix(proxy): print the --enforce_prisma_migration_check notice when the config fails to load f6f5957d16 printed the deprecation notice after the config is read, so a json_logs set only in the config makes it a JSON record. A boot whose config file is missing or does not parse then exited before that point and never printed it, where main printed it before the error (10-20 gauntlet finding). The notice is printed where main printed it, unless a config file can still switch json_logs on: then it is held until the config is read and printed right after the json_logs switch, and a click close callback prints it on the way out if the boot fails first. functools.cache keeps it to one print either way. Tests: tests/unit/proxy/test_proxy_cli.py covers a failed config read (red on 75043aed83), a plain config printing it once and a config-only json_logs making it one JSON record; tests/integration/configuration/test_json_logs_startup.py adds the JSON cell and a missing and an unparseable config cell (red on 75043aed83, green here), through refused_boot_log, which now takes extra CLI arguments. * fix(proxy): keep startup records as printed instead of running secret redaction In JSON mode the startup messages ran secret redaction, which the prints they replace never had. Its sk- key pattern has no left boundary, so a model named risk-assessment-gpt4o was logged as riREDACTED in the model list (10-21 gauntlet finding). The records now skip redaction like the prints did; the processed record still counts as scrubbed, so the handler's secret filter and the formatter leave it alone. The stdout cap stays off and base64 collapse stays on, as before. Tests: tests/unit/test_logging.py checks the model name comes out whole (red on 7d17f126e5), replacing the test that asserted startup records were redacted; tests/integration/configuration/test_json_logs_startup.py adds a boot with a risk-assessment- model whose model list record carries the name (red on 7d17f126e5). * refactor(proxy): meet the AGENTS.md rules Greptile flagged on the startup-message change The policy listing builds its rows as tuples in two helpers instead of appending to a list; plain output is unchanged (same bytes as before across 3000 random policy and attachment configs). The explanatory comments added with this change are gone, and the docstring of print_startup_message now says the record is whole and unredacted. Tests: the migration-diff JSON test no longer runs the fake prisma CLI; it checks the skipped-diff notice in-process, and the fixture extension it needed is reverted. The json_logs export and deprecation-notice CLI tests run the real proxy module and a real config file instead of a mocked ProxyConfig and proxy module, and the --health/--test notice tests no longer patch our helpers. Each still fails when the export, the hold or call_on_close is removed. New test parameters are typed. --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- litellm/_logging.py | 44 ++ litellm/proxy/db/check_migration.py | 19 +- litellm/proxy/policy_engine/init_policies.py | 111 ++--- litellm/proxy/proxy_cli.py | 113 +++-- litellm/proxy/proxy_server.py | 70 +-- tests/integration/_support/process.py | 2 + .../configuration/test_json_logs_startup.py | 397 ++++++++++++++++++ tests/unit/proxy/db/test_check_migration.py | 36 ++ .../proxy/policy_engine/test_init_policies.py | 116 +++++ tests/unit/proxy/test_proxy_cli.py | 138 ++++++ ...test_proxy_server_endpoints_and_startup.py | 62 ++- tests/unit/test_logging.py | 220 ++++++++++ 12 files changed, 1188 insertions(+), 140 deletions(-) create mode 100644 tests/integration/configuration/test_json_logs_startup.py create mode 100644 tests/unit/proxy/policy_engine/test_init_policies.py diff --git a/litellm/_logging.py b/litellm/_logging.py index d08254db3fd..8974f85f517 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -1,6 +1,7 @@ import ast import contextvars import functools +import inspect import itertools import logging import os @@ -897,6 +898,49 @@ verbose_proxy_logger.addFilter(_diagnostic_filter) verbose_proxy_stdout_logger.addFilter(_diagnostic_filter) verbose_logger.addFilter(_diagnostic_filter) +_ANSI_ESCAPE: Final = re.compile(r"\x1b\[[0-9;]*m") + + +def print_startup_message(message: str) -> None: + """Write a boot-time message meant for the operator's terminal. + + Plain mode prints it unchanged, colors and spacing included. With JSON logs on, every line of + stdout is read by a log pipeline, so the message goes out as one INFO record of the proxy logger + instead, on stdout like the print whatever handlers the logging config gave that logger, with the + ANSI color codes and surrounding blank lines removed, located at the caller and emitted at any log + level, whole and unredacted as the print showed it; a blank spacer produces no record. + """ + if not litellm.json_logs: + print(message) # noqa: T201 # plain mode: this is the operator-terminal print it replaces + return + text: Final = _ANSI_ESCAPE.sub("", message).strip("\n") + if not text: + return + frame: Final = inspect.currentframe() + caller: Final = frame.f_back if frame is not None else None + record: Final = verbose_proxy_logger.makeRecord( + verbose_proxy_logger.name, + logging.INFO, + caller.f_code.co_filename if caller is not None else __file__, + caller.f_lineno if caller is not None else 0, + text, + (), + None, + func=caller.f_code.co_name if caller is not None else None, + ) + if not _process_record(record, base64_limit=_get_max_base64_length_stdout_log(), text_limit=0, redact=False): + return + _startup_json_handler().handle(record) + + +def _startup_json_handler() -> logging.Handler: + """The proxy's JSON handler on stdout, filters included, for one startup message.""" + handler: Final = LevelRoutingStreamHandler() + handler.setFormatter(JsonFormatter()) + handler.addFilter(_secret_filter) + handler.addFilter(_correlation_filter) + return handler + def _suppress_loggers(): """Suppress noisy loggers at INFO level""" diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index 6e2a06c96e1..83cc4501466 100644 --- a/litellm/proxy/db/check_migration.py +++ b/litellm/proxy/db/check_migration.py @@ -4,7 +4,7 @@ import os import subprocess from typing import Final -from litellm._logging import verbose_logger +from litellm._logging import print_startup_message, verbose_logger def extract_sql_commands(diff_output: str) -> list[str]: @@ -63,8 +63,8 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: run_prisma, ) except ImportError as e: - print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages - f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}" + print_startup_message( + f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}", ) return False, [] @@ -89,22 +89,21 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: sql_commands: Final = extract_sql_commands(result.stdout) if sql_commands: - print("Changes to DB Schema detected") # noqa: T201 - print("Required SQL commands:") # noqa: T201 - for command in sql_commands: - print(command) # noqa: T201 + print_startup_message( + "\n".join(("Changes to DB Schema detected", "Required SQL commands:", *sql_commands)), + ) return True, sql_commands else: return False, [] except subprocess.TimeoutExpired: - print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + print_startup_message( f"Timed out after {timeout}s generating the migration diff. " - f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer." + f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer.", ) return False, [] except subprocess.CalledProcessError as e: error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}" - print(error_message) # noqa: T201 + print_startup_message(error_message) return False, [] diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py index d75e651e6c7..59d15b0ba0a 100644 --- a/litellm/proxy/policy_engine/init_policies.py +++ b/litellm/proxy/policy_engine/init_policies.py @@ -6,10 +6,12 @@ Configuration structure: - policy_attachments: Define WHERE policies apply (teams, keys, models) """ +import sys from collections.abc import Mapping, Sequence +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Optional -from litellm._logging import verbose_proxy_logger +from litellm._logging import print_startup_message, verbose_proxy_logger from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_validator import PolicyValidator @@ -25,6 +27,42 @@ _yellow_color_code: Final = "\033[93m" _reset_color_code: Final = "\033[0m" +def _policy_lines(policy_name: str, policy_data: Mapping[str, object]) -> tuple[str, ...]: + guardrails: Final = policy_data.get("guardrails", {}) + inherit: Final = policy_data.get("inherit") + condition: Final = policy_data.get("condition") + description: Final = policy_data.get("description") + + guardrails_add: Final = guardrails.get("add", []) if isinstance(guardrails, dict) else [] + guardrails_remove: Final = guardrails.get("remove", []) if isinstance(guardrails, dict) else [] + model_condition: Final = condition.get("model") if isinstance(condition, dict) else None + inherit_str: Final = f" (inherits: {inherit})" if inherit else "" + + return ( + f"{_blue_color_code} - {policy_name}{inherit_str}{_reset_color_code}", + *((f" description: {description}",) if description else ()), + *((f" guardrails.add: {guardrails_add}",) if guardrails_add else ()), + *((f" guardrails.remove: {guardrails_remove}",) if guardrails_remove else ()), + *((f" condition.model: {model_condition}",) if model_condition else ()), + ) + + +def _attachment_line(attachment: Mapping[str, object]) -> str: + policy: Final = attachment.get("policy", "unknown") + scope: Final = attachment.get("scope") + teams: Final = attachment.get("teams") + keys: Final = attachment.get("keys") + models: Final = attachment.get("models") + + scope_parts: Final = ( + *(("scope=* (global)",) if scope == "*" else ()), + *((f"teams={teams}",) if teams else ()), + *((f"keys={keys}",) if keys else ()), + *((f"models={models}",) if models else ()), + ) + return f" - {policy} -> {', '.join(scope_parts) if scope_parts else 'all'}" + + def _print_policies_on_startup( policies_config: Mapping[str, Mapping[str, object]], policy_attachments_config: Sequence[Mapping[str, object]] | None = None, @@ -32,67 +70,32 @@ def _print_policies_on_startup( """ Print loaded policies to console on startup (similar to model list). """ - import sys - - print( # noqa: T201 + header: Final = ( f"{_green_color_code}\nLiteLLM Policy Engine: Loaded {len(policies_config)} policies{_reset_color_code}\n" ) - sys.stdout.flush() - - for policy_name, policy_data in policies_config.items(): - guardrails = policy_data.get("guardrails", {}) - inherit = policy_data.get("inherit") - condition = policy_data.get("condition") - description = policy_data.get("description") - - guardrails_add = guardrails.get("add", []) if isinstance(guardrails, dict) else [] - guardrails_remove = guardrails.get("remove", []) if isinstance(guardrails, dict) else [] - inherit_str = f" (inherits: {inherit})" if inherit else "" - - print( # noqa: T201 - f"{_blue_color_code} - {policy_name}{inherit_str}{_reset_color_code}" + policy_lines: Final = tuple( + chain.from_iterable( + _policy_lines(policy_name, policy_data) for policy_name, policy_data in policies_config.items() ) - if description: - print(f" description: {description}") # noqa: T201 - if guardrails_add: - print(f" guardrails.add: {guardrails_add}") # noqa: T201 - if guardrails_remove: - print(f" guardrails.remove: {guardrails_remove}") # noqa: T201 - if condition: - model_condition = condition.get("model") if isinstance(condition, dict) else None - if model_condition: - print(f" condition.model: {model_condition}") # noqa: T201 + ) - # Print attachments if policy_attachments_config: - print( # noqa: T201 - f"\n{_yellow_color_code}Policy Attachments: {len(policy_attachments_config)} attachment(s){_reset_color_code}" + print_startup_message( + "\n".join( + ( + header, + *policy_lines, + f"\n{_yellow_color_code}Policy Attachments: {len(policy_attachments_config)} attachment(s){_reset_color_code}", + *(_attachment_line(attachment) for attachment in policy_attachments_config), + "", + ) + ) ) - for attachment in policy_attachments_config: - policy = attachment.get("policy", "unknown") - scope = attachment.get("scope") - teams = attachment.get("teams") - keys = attachment.get("keys") - models = attachment.get("models") - - scope_parts = [] - if scope == "*": - scope_parts.append("scope=* (global)") - if teams: - scope_parts.append(f"teams={teams}") - if keys: - scope_parts.append(f"keys={keys}") - if models: - scope_parts.append(f"models={models}") - scope_str = ", ".join(scope_parts) if scope_parts else "all" - - print(f" - {policy} -> {scope_str}") # noqa: T201 else: - print( # noqa: T201 - f"\n{_yellow_color_code}Warning: No policy_attachments configured. Policies will not be applied to any requests.{_reset_color_code}" + print_startup_message("\n".join((header, *policy_lines))) + print_startup_message( + f"\n{_yellow_color_code}Warning: No policy_attachments configured. Policies will not be applied to any requests.{_reset_color_code}\n", ) - - print() # noqa: T201 sys.stdout.flush() diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index fc33ab1bd0b..909f6de40cd 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1,4 +1,5 @@ # ruff: noqa: T201 +import functools import importlib import json import os @@ -7,7 +8,7 @@ import re import subprocess import sys import urllib.parse as urlparse -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from pathlib import Path from typing import TYPE_CHECKING, Any, Final @@ -18,6 +19,7 @@ from dotenv import load_dotenv from pydantic import ConfigDict import litellm +from litellm._logging import print_startup_message from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY from litellm.proxy.db.pgbouncer import ( PgBouncerError, @@ -202,6 +204,29 @@ def deprecated_v2_flag_passed_on_cli() -> bool: return ctx.get_parameter_source("use_v2_migration_resolver") is ParameterSource.COMMANDLINE +def _print_enforce_prisma_migration_check_notice() -> None: + print_startup_message( + "\033[1;33mLiteLLM Proxy: --enforce_prisma_migration_check is " + "deprecated and has no effect, because the proxy always exits " + "when database setup fails at startup. You can safely remove it.\033[0m", + ) + + +def _print_or_hold_enforce_prisma_migration_check_notice( + *, reads_config_file: bool, returns_before_config: bool +) -> Callable[[], None] | None: + """Print the notice now, or, while a config file can still switch json_logs on, return a printer + to call once it is read; a boot that fails before that prints the held notice on the way out.""" + if not reads_config_file or litellm.json_logs or returns_before_config: + _print_enforce_prisma_migration_check_notice() + return None + print_once: Final = functools.cache(_print_enforce_prisma_migration_check_notice) + ctx: Final = click.get_current_context(silent=True) + if ctx is not None: + ctx.call_on_close(print_once) + return print_once + + class ProxyInitializationHelpers: @staticmethod def echo_litellm_version(): @@ -335,7 +360,7 @@ class ProxyInitializationHelpers: "timeout_graceful_shutdown": GracefulShutdownManager.get_timeout(), } if log_config is not None: - print(f"Using log_config: {log_config}") + print_startup_message(f"Using log_config: {log_config}") uvicorn_args["log_config"] = log_config elif litellm.json_logs: # Use JSON log config for uvicorn to ensure all logs (including exceptions) are JSON @@ -348,10 +373,10 @@ class ProxyInitializationHelpers: if "timeout_worker_healthcheck" in inspect.signature(uvicorn.Config.__init__).parameters: uvicorn_args["timeout_worker_healthcheck"] = timeout_worker_healthcheck else: - print( + print_startup_message( f"\033[1;33mLiteLLM Proxy: --timeout_worker_healthcheck " f"requires uvicorn>=0.37.0, but installed uvicorn=={uvicorn.__version__}. " - f"Ignoring the flag.\033[0m" + f"Ignoring the flag.\033[0m", ) return uvicorn_args @@ -371,18 +396,18 @@ class ProxyInitializationHelpers: import uvicorn if max_requests_before_restart is None: - print( + print_startup_message( "\033[1;33mLiteLLM Proxy: --max_requests_before_restart_jitter " - "has no effect without --max_requests_before_restart\033[0m\n" + "has no effect without --max_requests_before_restart\033[0m\n", ) return if "limit_max_requests_jitter" in inspect.signature(uvicorn.Config.__init__).parameters: uvicorn_args["limit_max_requests_jitter"] = jitter else: - print( + print_startup_message( f"\033[1;33mLiteLLM Proxy: --max_requests_before_restart_jitter " f"requires uvicorn>=0.41.0, but installed uvicorn=={uvicorn.__version__}. " - f"Ignoring the flag.\033[0m" + f"Ignoring the flag.\033[0m", ) _apply_uvicorn_max_requests_jitter = apply_uvicorn_max_requests_jitter @@ -506,12 +531,12 @@ class ProxyInitializationHelpers: from hypercorn.asyncio import serve from hypercorn.config import Config - print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Hypercorn\033[0m\n") + print_startup_message(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Hypercorn\033[0m\n") config: Final = Config() config.bind = [f"{host}:{port}"] if ssl_certfile_path is not None and ssl_keyfile_path is not None: - print( + print_startup_message( f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" ) config.certfile = ssl_certfile_path @@ -544,14 +569,14 @@ class ProxyInitializationHelpers: from granian import Granian from granian.constants import Interfaces - print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Granian\033[0m\n") + print_startup_message(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Granian\033[0m\n") if max_requests_before_restart is not None: - print( + print_startup_message( "\033[1;33mLiteLLM: --max_requests_before_restart is not supported by Granian " - "(Granian uses workers_lifetime in seconds, not a per-request limit).\033[0m\n" + "(Granian uses workers_lifetime in seconds, not a per-request limit).\033[0m\n", ) if ciphers is not None: - print("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") + print_startup_message("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") kwargs: Final[dict[str, object]] = { "target": "litellm.proxy.proxy_server:app", @@ -564,7 +589,7 @@ class ProxyInitializationHelpers: if granian_runtime_threads is not None: kwargs["runtime_threads"] = granian_runtime_threads if ssl_certfile_path is not None and ssl_keyfile_path is not None: - print( + print_startup_message( f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" ) kwargs["ssl_cert"] = Path(ssl_certfile_path) @@ -625,13 +650,15 @@ class ProxyInitializationHelpers: \n """ ) - print() - print( + print_startup_message("") + print_startup_message( '\033[1;34mLiteLLM: Test your local proxy with: "litellm --test" This runs an openai.ChatCompletion request to your proxy [In a new terminal tab]\033[0m\n' ) - print(f"\033[1;34mLiteLLM: Curl Command Test for your local proxy\n {curl_command} \033[0m\n") - print("\033[1;34mDocs: https://docs.litellm.ai/docs/simple_proxy\033[0m\n") - print(f"\033[1;34mSee all Router/Swagger docs on http://0.0.0.0:{port} \033[0m\n") + print_startup_message( + f"\033[1;34mLiteLLM: Curl Command Test for your local proxy\n {curl_command} \033[0m\n" + ) + print_startup_message("\033[1;34mDocs: https://docs.litellm.ai/docs/simple_proxy\033[0m\n") + print_startup_message(f"\033[1;34mSee all Router/Swagger docs on http://0.0.0.0:{port} \033[0m\n") def load_config(self): # note: This Loads the gunicorn config - has nothing to do with LiteLLM Proxy config @@ -651,7 +678,9 @@ class ProxyInitializationHelpers: # gunicorn app function return self.application - print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} with {num_workers} workers\033[0m\n") + print_startup_message( + f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} with {num_workers} workers\033[0m\n" + ) graceful_shutdown_timeout: Final = GracefulShutdownManager.get_timeout() gunicorn_options: Final = { "bind": f"{host}:{port}", @@ -669,9 +698,9 @@ class ProxyInitializationHelpers: gunicorn_options["max_requests"] = max_requests_before_restart if max_requests_before_restart_jitter is not None: if max_requests_before_restart is None: - print( + print_startup_message( "\033[1;33mLiteLLM Proxy: --max_requests_before_restart_jitter " - "has no effect without --max_requests_before_restart\033[0m\n" + "has no effect without --max_requests_before_restart\033[0m\n", ) else: gunicorn_options["max_requests_jitter"] = max_requests_before_restart_jitter @@ -686,7 +715,7 @@ class ProxyInitializationHelpers: gunicorn_options["child_exit"] = child_exit if ssl_certfile_path is not None and ssl_keyfile_path is not None: - print( + print_startup_message( f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" ) gunicorn_options["certfile"] = ssl_certfile_path @@ -710,7 +739,7 @@ class ProxyInitializationHelpers: with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - print(f""" + print_startup_message(f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` """) @@ -779,7 +808,7 @@ class ProxyInitializationHelpers: os.makedirs(multiproc_dir, exist_ok=True) wipe_directory(multiproc_dir) action: Final = "Using existing" if configured_dir else "Auto-created" - print(f"LiteLLM: {action} PROMETHEUS_MULTIPROC_DIR={multiproc_dir}") + print_startup_message(f"LiteLLM: {action} PROMETHEUS_MULTIPROC_DIR={multiproc_dir}") return multiproc_dir _maybe_setup_prometheus_multiproc_dir = maybe_setup_prometheus_multiproc_dir @@ -1199,12 +1228,13 @@ def run_server( return if check_migrations_against_existing_data is True: sys.exit(ProxyInitializationHelpers.run_migration_data_check(use_migrate=not use_prisma_db_push)) - if enforce_prisma_migration_check: - print( - "\033[1;33mLiteLLM Proxy: --enforce_prisma_migration_check is " - "deprecated and has no effect, because the proxy always exits " - "when database setup fails at startup. You can safely remove it.\033[0m" + held_prisma_notice: Final = ( + _print_or_hold_enforce_prisma_migration_check_notice( + reads_config_file=config is not None, returns_before_config=health is True or test is True ) + if enforce_prisma_migration_check + else None + ) if model and "ollama" in model and api_base is None: ProxyInitializationHelpers.run_ollama_serve() if health is True: @@ -1320,8 +1350,11 @@ def run_server( import litellm litellm.json_logs = True + os.environ["JSON_LOGS"] = "true" litellm.turn_on_json() + if held_prisma_notice is not None: + held_prisma_notice() ### GENERAL SETTINGS ### general_settings = _config.get("general_settings", {}) if general_settings is None: @@ -1488,7 +1521,7 @@ def run_server( lifetime_params, ) ) - print( + print_startup_message( postgres_connection_budget_message( writer_limit=_url_query_value(os.getenv("DATABASE_URL"), "connection_limit") or f"{db_connection_pool_limit}", @@ -1520,19 +1553,19 @@ def run_server( env_value=os.getenv("USE_V2_MIGRATION_RESOLVER"), ) if deprecated_v2_flag_passed_on_cli() and use_v2_resolver: - print( + print_startup_message( "\033[1;33mLiteLLM Proxy: --use_v2_migration_resolver is " "deprecated and has no effect, because the v2 migration " "resolver is now the default. You can safely remove it. To " "opt back into the legacy v1 resolver, pass " - "--use_legacy_migration_resolver.\033[0m" + "--use_legacy_migration_resolver.\033[0m", ) if not use_v2_resolver: - print( + print_startup_message( "\033[1;33mLiteLLM Proxy: Using the legacy (v1) migration " "resolver. It performs the diff-and-force recovery that can " "cause schema thrashing during rolling deploys where two " - "LiteLLM versions contend for the same DB.\033[0m" + "LiteLLM versions contend for the same DB.\033[0m", ) try: migrated: Final = PrismaManager.setup_database( @@ -1600,7 +1633,7 @@ def run_server( # Skip server startup if requested (after all setup is done) if skip_server_startup: - print("LiteLLM: Setup complete. Skipping server startup as requested.") + print_startup_message("LiteLLM: Setup complete. Skipping server startup as requested.") return # Auto-create PROMETHEUS_MULTIPROC_DIR for multi-worker setups @@ -1619,7 +1652,7 @@ def run_server( ) except MetricsServerStartupError as error: raise click.ClickException(str(error)) from error - print( + print_startup_message( f"\033[1;32mLiteLLM: Serving Prometheus metrics on {host}:{prometheus_metrics_port}/metrics " f"(pid {metrics_process.pid})\033[0m" ) @@ -1642,7 +1675,7 @@ def run_server( run_gunicorn=run_gunicorn is True, ) if single_worker_warning is not None: - print(f"\033[1;33m{single_worker_warning}\033[0m\n") + print_startup_message(f"\033[1;33m{single_worker_warning}\033[0m\n") if run_gunicorn is False and run_hypercorn is False and run_granian is False: if limit_concurrency is not None: uvicorn_args["limit_concurrency"] = limit_concurrency @@ -1653,7 +1686,7 @@ def run_server( jitter=max_requests_before_restart_jitter, ) if ssl_certfile_path is not None and ssl_keyfile_path is not None: - print( + print_startup_message( f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" ) uvicorn_args["ssl_keyfile"] = ssl_keyfile_path diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bc79becf7bb..3b36ea5afce 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -273,25 +273,27 @@ def generate_feedback_box(): # Select a random message message: Final = random.choice(list_of_messages) - print() # noqa: T201 - print("\033[1;37m" + "#" + "-" * box_width + "#\033[0m") # noqa: T201 - print("\033[1;37m" + "#" + " " * box_width + "#\033[0m") # noqa: T201 - print("\033[1;37m" + f"# {message:^59} #\033[0m") # noqa: T201 - print( # noqa: T201 - "\033[1;37m" + "# {:^59} #\033[0m".format("https://github.com/BerriAI/litellm/issues/new") + print_startup_message( + "\n".join( + ( + "", + "\033[1;37m" + "#" + "-" * box_width + "#\033[0m", + "\033[1;37m" + "#" + " " * box_width + "#\033[0m", + "\033[1;37m" + f"# {message:^59} #\033[0m", + "\033[1;37m" + "# {:^59} #\033[0m".format("https://github.com/BerriAI/litellm/issues/new"), + "\033[1;37m" + "#" + " " * box_width + "#\033[0m", + "\033[1;37m" + "#" + "-" * box_width + "#\033[0m", + "", + " Thank you for using LiteLLM! - Krrish & Ishaan", + "", + "", + "", + "\033[1;31mGive Feedback / Get Help: https://github.com/BerriAI/litellm/issues/new\033[0m", + "", + "", + ) + ) ) - print("\033[1;37m" + "#" + " " * box_width + "#\033[0m") # noqa: T201 - print("\033[1;37m" + "#" + "-" * box_width + "#\033[0m") # noqa: T201 - print() # noqa: T201 - print(" Thank you for using LiteLLM! - Krrish & Ishaan") # noqa: T201 - print() # noqa: T201 - print() # noqa: T201 - print() # noqa: T201 - print( # noqa: T201 - "\033[1;31mGive Feedback / Get Help: https://github.com/BerriAI/litellm/issues/new\033[0m" - ) - print() # noqa: T201 - print() # noqa: T201 import contextlib @@ -304,6 +306,7 @@ import litellm._redis from litellm import Router from litellm._internal_context import service_target, with_service_target from litellm._logging import ( + print_startup_message, redact_string, verbose_proxy_logger, verbose_router_logger, @@ -6412,17 +6415,13 @@ class ProxyConfig: search_tools_parsed: Final[list[SearchToolTypedDict]] = [] - print( # noqa: T201 - "\033[32mLiteLLM: Proxy initialized with Search Tools:\033[0m" - ) + print_startup_message("\033[32mLiteLLM: Proxy initialized with Search Tools:\033[0m") for search_tool in search_tools_raw: # Display loaded search tool search_tool_name = search_tool.get("search_tool_name", "") search_provider = search_tool.get("litellm_params", {}).get("search_provider", "") - print( # noqa: T201 - f"\033[32m {search_tool_name} ({search_provider})\033[0m" - ) + print_startup_message(f"\033[32m {search_tool_name} ({search_provider})\033[0m") # Handle os.environ/ variables in litellm_params litellm_params = search_tool.get("litellm_params", {}) @@ -6598,7 +6597,7 @@ class ProxyConfig: reset_color_code: Final = "\033[0m" for key, value in litellm_settings.items(): if key == "cache" and value is True: - print(f"{blue_color_code}\nSetting Cache on Proxy") # noqa: T201 + print_startup_message(f"{blue_color_code}\nSetting Cache on Proxy") from litellm.caching.caching import Cache cache_params = {} @@ -6789,7 +6788,7 @@ class ProxyConfig: if PrometheusLogger is not None: verbose_proxy_logger.debug("mounting metrics endpoint") PrometheusLogger._mount_metrics_endpoint() - print( # noqa: T201 + print_startup_message( f"{blue_color_code} Initialized Success Callbacks - {litellm.success_callback} {reset_color_code}" ) elif key == "failure_callback": @@ -6808,7 +6807,7 @@ class ProxyConfig: # these are litellm callbacks - "langfuse", "sentry", "wandb" else: litellm.logging_callback_manager.add_litellm_failure_callback(callback) - print( # noqa: T201 + print_startup_message( f"{blue_color_code} Initialized Failure Callbacks - {litellm.failure_callback} {reset_color_code}" ) elif key == "audit_log_callbacks": @@ -6833,7 +6832,7 @@ class ProxyConfig: _store_audit_logs = litellm_settings.get("store_audit_logs", litellm.store_audit_logs) if is_audit_logging_enabled(store_audit_logs=_store_audit_logs): - print( # noqa: T201 + print_startup_message( f"{blue_color_code} Initialized Audit Log Callbacks - {litellm.audit_log_callbacks} {reset_color_code}" ) else: @@ -7232,9 +7231,7 @@ class ProxyConfig: if model_list: router_params["model_list"] = model_list validate_auto_router_capability_limits(model_list, limit=_license_check.auto_router_capability_limit()) - print( # noqa: T201 - "\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m" - ) + print_startup_message("\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m") for model in model_list: ### LOAD FROM os.environ/ ### for k, v in model["litellm_params"].items(): @@ -7251,7 +7248,7 @@ class ProxyConfig: complexity_router_config=complexity_router_config, config_file_path=config_file_path, ) - print(f"\033[32m {model.get('model_name', '')}\033[0m") # noqa: T201 + print_startup_message(f"\033[32m {model.get('model_name', '')}\033[0m") litellm_model_name = model["litellm_params"]["model"] litellm_model_api_base = model["litellm_params"].get("api_base", None) if "ollama" in litellm_model_name and litellm_model_api_base is None: @@ -9672,9 +9669,12 @@ async def initialize( master_key, \ user_custom_auth, \ prisma_client - from litellm.proxy.common_utils.banner import show_banner + from litellm.proxy.common_utils.banner import LITELLM_BANNER, show_banner - show_banner() + if litellm.json_logs: + print_startup_message(LITELLM_BANNER) + else: + show_banner() if os.getenv("LITELLM_DONT_SHOW_FEEDBACK_BOX", "").lower() != "true": generate_feedback_box() user_model = model @@ -11930,7 +11930,7 @@ class ProxyStartupEvent: and proxy_logging_obj.slack_alerting_instance.alerting is not None and prisma_client is not None ): - print("Alerting: Initializing Weekly/Monthly Spend Reports") # noqa: T201 + print_startup_message("Alerting: Initializing Weekly/Monthly Spend Reports") spend_report_frequency: Final[str] = general_settings.get("spend_report_frequency", "7d") or "7d" days: Final = int(spend_report_frequency[:-1]) diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 052c52982b6..262114ee32b 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -987,6 +987,7 @@ def refused_boot_log( overrides: Mapping[str, str], *, config: Path | None = None, + extra_arguments: tuple[str, ...] = (), ) -> str: """Start the proxy and return its log once it exits non-zero instead of becoming ready.""" root: Final = _proxy_root() @@ -1004,6 +1005,7 @@ def refused_boot_log( "--num_workers", "1", *DB_PUSH, + *extra_arguments, ) launch: Final = _launch(command, root, environment, output) try: diff --git a/tests/integration/configuration/test_json_logs_startup.py b/tests/integration/configuration/test_json_logs_startup.py new file mode 100644 index 00000000000..de139b7986f --- /dev/null +++ b/tests/integration/configuration/test_json_logs_startup.py @@ -0,0 +1,397 @@ +"""JSON logging keeps every line the proxy writes during boot, serving and shutdown machine-parseable. + +A deployment that turns on JSON logs ships the proxy's stdout to a log pipeline, and every line that is +not a JSON object breaks that pipeline's parser. This contract boots an owned proxy with JSON logs on, +serves a chat request on every configured model plus one rejected request, stops the proxy and +rejects any line of its output that is not one JSON record, while checking that the +startup information those lines used to carry is still present as JSON records. JSON mode is switched +on through the environment and through the config, at one and two uvicorn workers, at the INFO, the +default and the WARNING log level, under a concurrent burst that lands on both workers, under an operator +--log_config that names the proxy logger or gives it a handler of its own, and with a policy list longer +than the stdout cap. The model list keeps a model name the secret patterns would match. The deprecated +--enforce_prisma_migration_check notice is one JSON record under a config-only json_logs, and a boot whose +config cannot be loaded still prints it before failing. +""" + +from __future__ import annotations + +import re +import uuid +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import pytest +import yaml +from pydantic import JsonValue, TypeAdapter, ValidationError + +from tests.integration._support.client import Gateway +from tests.integration._support.process import owned_proxy_process, refused_boot_log + +JSON_RECORD: Final = TypeAdapter(dict[str, JsonValue]) +PLAIN_LOG_LINE: Final = re.compile(r"^\d\d:\d\d:\d\d - LiteLLM") +PRISMA_ENGINE_TEXT: Final = re.compile(r"query-engine [0-9a-f]{40}") +STARTED_WORKER: Final = "Started server process" +MODEL_LIST_HEADER: Final = "Proxy initialized with Config, Set models:" +CACHE_LINE: Final = "Setting Cache on Proxy" +CONNECTION_BUDGET_LINE: Final = "Postgres connection budget" +REDIS_AUTH_CACHE_LINE: Final = "enable_redis_auth_cache" +BANNER_FRAGMENT: Final = "██╗" +FEEDBACK_BOX_LINE: Final = "Give Feedback / Get Help" +WORKER_COUNTS: Final = (1, 2) +UNKNOWN_MODEL_STATUS: Final = 400 +BURST_SIZE: Final = 16 +CACHE_SETTINGS: Final[Mapping[str, JsonValue]] = {"cache": True, "cache_params": {"type": "redis"}} +STDOUT_CAP: Final = 4096 +POLICY_COUNT: Final = 60 +POLICY_LIST_HEADER: Final = "LiteLLM Policy Engine: Loaded" +TRUNCATION_MARKER: Final = "litellm_truncated" +DEPRECATED_FLAG: Final = "--enforce_prisma_migration_check" +DEPRECATED_FLAG_NOTICE: Final = "--enforce_prisma_migration_check is deprecated and has no effect" +UNPARSEABLE_CONFIG: Final = "model_list: [\n - bad: yaml: here\n" +OPERATOR_LOG_CONFIG: Final[Mapping[str, JsonValue]] = { + "version": 1, + "disable_existing_loggers": False, + "formatters": {"json": {"()": "litellm._logging.JsonFormatter"}}, + "handlers": {"stdout": {"class": "logging.StreamHandler", "stream": "ext://sys.stdout", "formatter": "json"}}, + "loggers": {"LiteLLM Proxy": {"level": "INFO"}}, + "root": {"level": "INFO", "handlers": ["stdout"]}, +} +OPERATOR_SILENCING_LOG_CONFIG: Final[Mapping[str, JsonValue]] = { + "version": 1, + "disable_existing_loggers": False, + "formatters": {"json": {"()": "litellm._logging.JsonFormatter"}}, + "handlers": { + "stdout": {"class": "logging.StreamHandler", "stream": "ext://sys.stdout", "formatter": "json"}, + "silence": {"class": "logging.NullHandler"}, + }, + "loggers": {"LiteLLM Proxy": {"level": "INFO", "handlers": ["silence"], "propagate": False}}, + "root": {"level": "INFO", "handlers": ["stdout"]}, +} + + +def _models() -> tuple[str, str]: + stamp: Final = uuid.uuid4().hex + return (f"json-logs-first-{stamp}", f"json-logs-second-{stamp}") + + +def _policies() -> dict[str, JsonValue]: + stamp: Final = uuid.uuid4().hex + return { + f"startup-listing-policy-{n:02d}-{stamp}": { + "description": f"policy {n:02d} of a listing that runs past the stdout cap meant for request payloads" + } + for n in range(POLICY_COUNT) + } + + +def _write_config( + directory: Path, + upstream_url: str, + models: tuple[str, ...], + litellm_settings: Mapping[str, JsonValue], + extra_sections: Mapping[str, JsonValue] | None = None, +) -> Path: + config: Final = directory / f"json_logs_{uuid.uuid4().hex}.yaml" + config.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": model, + "litellm_params": { + "model": f"openai/{model}", + "api_base": f"{upstream_url}/v1", + "api_key": "integration-provider-key", + }, + } + for model in models + ], + "litellm_settings": dict(litellm_settings), + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_prisma_schema_update": True, + }, + **(extra_sections or {}), + }, + sort_keys=False, + ) + ) + return config + + +def _record(line: str) -> dict[str, JsonValue] | None: + stripped: Final = line.strip() + if not stripped.startswith("{"): + return None + try: + return JSON_RECORD.validate_json(stripped) + except ValidationError: + return None + + +def _lines(log: str) -> tuple[str, ...]: + return tuple(line for line in PRISMA_ENGINE_TEXT.sub("", log).splitlines() if line.strip()) + + +def _messages(log: str) -> tuple[str, ...]: + records: Final = tuple(record for record in (_record(line) for line in _lines(log)) if record is not None) + return tuple(str(record["message"]) for record in records if "message" in record) + + +def _non_json_lines(log: str) -> tuple[str, ...]: + return tuple(line for line in _lines(log) if _record(line) is None) + + +def _boot_serve_and_stop( + gateway: Gateway, + directory: Path, + config: Path, + overrides: Mapping[str, str], + workers: int, + models: tuple[str, ...], + extra_arguments: tuple[str, ...] = (), +) -> str: + with owned_proxy_process( + gateway, + directory, + overrides, + config=config, + workers=workers, + database_setup=(), + extra_arguments=extra_arguments, + ) as owned: + _serve(owned.gateway, models) + log: Final = owned.log + return log.read_text() + + +def _serve(gateway: Gateway, models: tuple[str, ...]) -> None: + """One chat per model and one rejected request, so the log holds the request-time lines of served + and failed requests next to the boot lines.""" + for model in models: + reply: Final = gateway.chat(model, text=f"json logs {uuid.uuid4().hex}") + assert reply["model"] == model, reply + rejected: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": f"missing-{uuid.uuid4().hex}", "messages": [{"role": "user", "content": "json logs"}]}, + ) + assert rejected.status_code == UNKNOWN_MODEL_STATUS, rejected.text + + +def _assert_every_line_is_json( + log: str, workers: int, models: tuple[str, ...], *, redis_auth_cache_line: bool = True +) -> None: + plaintext: Final = _non_json_lines(log) + assert plaintext == (), "plaintext lines with JSON logs on:\n" + "\n".join(plaintext) + messages: Final = _messages(log) + assert sum(STARTED_WORKER in message for message in messages) == workers, messages + assert any(MODEL_LIST_HEADER in message for message in messages), messages + for model in models: + assert any(model in message for message in messages), (model, messages) + assert any(CACHE_LINE in message for message in messages), messages + if redis_auth_cache_line: + assert any(REDIS_AUTH_CACHE_LINE in message for message in messages), messages + assert any(CONNECTION_BUDGET_LINE in message for message in messages), messages + assert sum(BANNER_FRAGMENT in message for message in messages) == workers, messages + assert sum(FEEDBACK_BOX_LINE in message for message in messages) == workers, messages + + +@pytest.mark.parametrize("workers", WORKER_COUNTS) +def test_json_logs_env_keeps_every_line_json(gateway: Gateway, tmp_path: Path, workers: int) -> None: + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, CACHE_SETTINGS) + log: Final = _boot_serve_and_stop( + gateway, tmp_path, config, {"JSON_LOGS": "true", "LITELLM_LOG": "INFO"}, workers, models + ) + _assert_every_line_is_json(log, workers, models) + + +@pytest.mark.parametrize("workers", WORKER_COUNTS) +def test_json_logs_config_keeps_every_line_json(gateway: Gateway, tmp_path: Path, workers: int) -> None: + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, workers, models) + _assert_every_line_is_json(log, workers, models) + + +@pytest.mark.parametrize("workers", WORKER_COUNTS) +def test_json_logs_with_operator_log_config_keeps_every_line_json( + gateway: Gateway, tmp_path: Path, workers: int +) -> None: + """uvicorn applies an operator `--log_config` after the CLI process switched JSON on and before + `initialize()` switches it on again through the config; a config that names the proxy logger leaves it + without handlers in that window, where the banner and the feedback box are emitted. They still arrive, + as JSON, one per worker, and the operator's own formatter keeps uvicorn's lines JSON. The Redis + auth-cache line is a regular logger line written in that same window; this log_config removed the + proxy logger's handlers and left propagation off, so it is lost at one worker on main and here alike.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + log_config: Final = tmp_path / "operator_log_config.yaml" + log_config.write_text(yaml.safe_dump(dict(OPERATOR_LOG_CONFIG))) + log: Final = _boot_serve_and_stop( + gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, workers, models, ("--log_config", str(log_config)) + ) + _assert_every_line_is_json(log, workers, models, redis_auth_cache_line=False) + + +@pytest.mark.parametrize("workers", WORKER_COUNTS) +def test_json_logs_with_operator_handler_on_the_proxy_logger_keeps_the_startup_notices( + gateway: Gateway, tmp_path: Path, workers: int +) -> None: + """An operator `--log_config` that gives the proxy logger a handler of its own, here a NullHandler that + silences it, is in force in the window where the banner and the feedback box are emitted. The prints + they replace wrote to stdout whatever the logging config said, so both still arrive there as JSON, one + per worker. The Redis auth-cache line is a regular logger line written in that same window and goes + to the operator's handler, so it is silenced on main and here alike.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + log_config: Final = tmp_path / "operator_silencing_log_config.yaml" + log_config.write_text(yaml.safe_dump(dict(OPERATOR_SILENCING_LOG_CONFIG))) + log: Final = _boot_serve_and_stop( + gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, workers, models, ("--log_config", str(log_config)) + ) + _assert_every_line_is_json(log, workers, models, redis_auth_cache_line=False) + + +def test_json_logs_keeps_a_startup_listing_longer_than_the_stdout_cap_whole(gateway: Gateway, tmp_path: Path) -> None: + """A startup listing longer than MAX_STRING_LENGTH_STDOUT_LOG reaches the pipeline whole. The policy list + stands in for every such listing (the required migration SQL on a drifted database is the other one): + the print it replaces showed it whole, and the stdout cap exists for echoed request payloads, so the JSON + record must not be middle-truncated.""" + models: Final = _models() + policies: Final = _policies() + config: Final = _write_config( + tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}, {"policies": policies} + ) + + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, 1, models) + + _assert_every_line_is_json(log, 1, models) + listings: Final = tuple(message for message in _messages(log) if POLICY_LIST_HEADER in message) + assert len(listings) == 1, listings + assert len(listings[0]) > STDOUT_CAP, len(listings[0]) + assert TRUNCATION_MARKER not in listings[0], listings[0] + for policy in policies: + assert policy in listings[0], (policy, listings[0]) + + +def test_json_logs_model_list_keeps_a_model_name_the_secret_patterns_would_match( + gateway: Gateway, tmp_path: Path +) -> None: + """The model list printed every model name as configured. The sk- key pattern matches inside a name such as + `risk-assessment-...`, so the JSON record must carry that name as the print did, not `riREDACTED`.""" + models: Final = (f"risk-assessment-{uuid.uuid4().hex}",) + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, 1, models) + + _assert_every_line_is_json(log, 1, models) + assert f" {models[0]}" in _messages(log), _messages(log) + + +def test_json_logs_default_level_keeps_every_line_json(gateway: Gateway, tmp_path: Path) -> None: + """Without LITELLM_LOG the proxy loggers sit at WARNING; the startup notices a plain-mode operator + always sees still arrive, as JSON, and nothing arrives as plaintext.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {}, 1, models) + plaintext: Final = _non_json_lines(log) + assert plaintext == (), "plaintext lines with JSON logs on:\n" + "\n".join(plaintext) + messages: Final = _messages(log) + assert sum(STARTED_WORKER in message for message in messages) == 1, messages + for fragment in ( + BANNER_FRAGMENT, + FEEDBACK_BOX_LINE, + MODEL_LIST_HEADER, + CACHE_LINE, + CONNECTION_BUDGET_LINE, + *models, + ): + assert any(fragment in message for message in messages), (fragment, messages) + + +def test_json_logs_warning_level_keeps_startup_notices(gateway: Gateway, tmp_path: Path) -> None: + """`LITELLM_LOG=WARNING` drops the INFO logger lines; the startup notices a plain-mode operator sees at + every level still arrive, as JSON, and nothing arrives as plaintext.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "WARNING"}, 1, models) + plaintext: Final = _non_json_lines(log) + assert plaintext == (), "plaintext lines with JSON logs on:\n" + "\n".join(plaintext) + messages: Final = _messages(log) + assert not any(REDIS_AUTH_CACHE_LINE in message for message in messages), messages + for fragment in ( + BANNER_FRAGMENT, + FEEDBACK_BOX_LINE, + MODEL_LIST_HEADER, + CACHE_LINE, + CONNECTION_BUDGET_LINE, + *models, + ): + assert any(fragment in message for message in messages), (fragment, messages) + + +def test_json_logs_burst_on_two_workers_keeps_every_line_json(gateway: Gateway, tmp_path: Path) -> None: + """Concurrent requests spread over both workers interleave their log writes; every line must still be + one complete JSON record.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {**CACHE_SETTINGS, "json_logs": True}) + with owned_proxy_process( + gateway, tmp_path, {"LITELLM_LOG": "INFO"}, config=config, workers=2, database_setup=() + ) as owned: + with ThreadPoolExecutor(max_workers=BURST_SIZE) as pool: + replies: Final = tuple( + pool.map(lambda index: owned.gateway.chat(models[index % 2], text=f"burst {index}"), range(BURST_SIZE)) + ) + assert tuple(reply["model"] for reply in replies) == tuple(models[index % 2] for index in range(BURST_SIZE)) + log: Final = owned.log + text: Final = log.read_text() + plaintext: Final = _non_json_lines(text) + assert plaintext == (), "plaintext lines with JSON logs on:\n" + "\n".join(plaintext) + assert sum(STARTED_WORKER in message for message in _messages(text)) == 2, _messages(text) + + +def test_json_logs_config_makes_the_deprecated_flag_notice_one_json_record(gateway: Gateway, tmp_path: Path) -> None: + """The notice waits until the config is read, so a json_logs set only there makes it a JSON record.""" + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, {"json_logs": True}) + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, 1, models, (DEPRECATED_FLAG,)) + plaintext: Final = _non_json_lines(log) + assert plaintext == (), "plaintext lines with JSON logs on:\n" + "\n".join(plaintext) + assert sum(DEPRECATED_FLAG_NOTICE in message for message in _messages(log)) == 1, _messages(log) + + +@pytest.mark.parametrize("config_text", (None, UNPARSEABLE_CONFIG), ids=("missing", "unparseable")) +def test_a_config_that_fails_to_load_still_prints_the_deprecated_flag_notice( + gateway: Gateway, tmp_path: Path, config_text: str | None +) -> None: + """A boot with the deprecated flag and a config file that does not exist or does not parse exits + non-zero, and it still prints the notice once before the error, as it did before the notice waited + for the config.""" + config: Final = tmp_path / f"unloadable_{uuid.uuid4().hex}.yaml" + if config_text is not None: + config.write_text(config_text) + log: Final = refused_boot_log(gateway, tmp_path, {}, config=config, extra_arguments=(DEPRECATED_FLAG,)) + assert log.count(DEPRECATED_FLAG_NOTICE) == 1, log + + +def test_plain_logs_control_prints_startup_text(gateway: Gateway, tmp_path: Path) -> None: + models: Final = _models() + config: Final = _write_config(tmp_path, gateway.upstream_url, models, CACHE_SETTINGS) + log: Final = _boot_serve_and_stop(gateway, tmp_path, config, {"LITELLM_LOG": "INFO"}, 1, models) + plaintext: Final = _non_json_lines(log) + assert any(BANNER_FRAGMENT in line for line in plaintext), log + assert any(FEEDBACK_BOX_LINE in line for line in plaintext), log + assert any(CACHE_LINE in line for line in plaintext), log + assert any(MODEL_LIST_HEADER in line for line in plaintext), log + for model in models: + assert any(line.strip().endswith(model) or line.strip().endswith(f"{model}\x1b[0m") for line in plaintext), ( + model, + log, + ) + assert any(PLAIN_LOG_LINE.match(line) and REDIS_AUTH_CACHE_LINE in line for line in plaintext), log + assert any(line.startswith(f"LiteLLM Proxy: {CONNECTION_BUDGET_LINE}") for line in plaintext), log diff --git a/tests/unit/proxy/db/test_check_migration.py b/tests/unit/proxy/db/test_check_migration.py index 74a6841cb23..b5685676941 100644 --- a/tests/unit/proxy/db/test_check_migration.py +++ b/tests/unit/proxy/db/test_check_migration.py @@ -1,6 +1,12 @@ +import json +import logging +import sys +from typing import Final import pytest +import litellm +from litellm._logging import turn_on_json, verbose_logger, verbose_proxy_logger def test_check_migration_out_of_sync(mocker): @@ -69,3 +75,33 @@ def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(m monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) + + +def test_a_skipped_migration_diff_is_one_json_record_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """ + With json_logs on, the notice that the migration diff was skipped is one INFO record on stdout, + also while the loggers sit at WARNING, as they do at boot without LITELLM_LOG. + """ + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + verbose_logger.setLevel(logging.WARNING) + verbose_proxy_logger.setLevel(logging.WARNING) + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert (record["level"], record["message"]) == ( + "INFO", + "Skipping the migration diff: litellm-proxy-extras has no Prisma runner. " + "Error: import of litellm_proxy_extras.prisma_toolchain halted; None in sys.modules", + ) diff --git a/tests/unit/proxy/policy_engine/test_init_policies.py b/tests/unit/proxy/policy_engine/test_init_policies.py new file mode 100644 index 00000000000..b9b5f278252 --- /dev/null +++ b/tests/unit/proxy/policy_engine/test_init_policies.py @@ -0,0 +1,116 @@ +"""The policy table printed at boot keeps its terminal form in plain mode and is JSON when json_logs is on.""" + +import json +import sys +from collections.abc import Mapping, Sequence +from typing import Final + +import pytest + +import litellm +from litellm._logging import turn_on_json +from litellm.proxy.policy_engine.init_policies import _print_policies_on_startup + +POLICIES: Final[Mapping[str, Mapping[str, object]]] = { + "pii-policy": { + "inherit": "base", + "description": "Block PII", + "guardrails": {"add": ["pii-guard"], "remove": ["noop"]}, + "condition": {"model": "gpt-4o"}, + } +} +ATTACHMENTS: Final[Sequence[Mapping[str, object]]] = ({"policy": "pii-policy", "scope": "*", "teams": ["team-a"]},) +POLICY_ROWS: Final = ( + " - pii-policy (inherits: base)", + " description: Block PII", + " guardrails.add: ['pii-guard']", + " guardrails.remove: ['noop']", + " condition.model: gpt-4o", +) +PLAIN_POLICY_TABLE: Final = ( + "\x1b[92m\nLiteLLM Policy Engine: Loaded 1 policies\x1b[0m\n\n" + "\x1b[94m - pii-policy (inherits: base)\x1b[0m\n" + " description: Block PII\n" + " guardrails.add: ['pii-guard']\n" + " guardrails.remove: ['noop']\n" + " condition.model: gpt-4o\n" +) +NO_ATTACHMENTS_WARNING: Final = ( + "Warning: No policy_attachments configured. Policies will not be applied to any requests." +) + + +@pytest.mark.parametrize( + ("attachments", "expected"), + [ + ( + ATTACHMENTS, + PLAIN_POLICY_TABLE + + "\n\x1b[93mPolicy Attachments: 1 attachment(s)\x1b[0m\n" + + " - pii-policy -> scope=* (global), teams=['team-a']\n\n", + ), + (None, PLAIN_POLICY_TABLE + f"\n\x1b[93m{NO_ATTACHMENTS_WARNING}\x1b[0m\n\n"), + ], + ids=["with-attachments", "without-attachments"], +) +def test_policy_table_prints_unchanged_in_plain_mode( + monkeypatch: pytest.MonkeyPatch, + capfd: pytest.CaptureFixture[str], + attachments: Sequence[Mapping[str, object]] | None, + expected: str, +) -> None: + """Operators without JSON logs keep the colored policy table byte for byte.""" + monkeypatch.setattr(litellm, "json_logs", False) + + _print_policies_on_startup(POLICIES, attachments) + + assert capfd.readouterr() == (expected, "") + + +def test_policy_table_is_one_json_record_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """With JSON logs on the whole table, attachments included, is one INFO record without color codes.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + + _print_policies_on_startup(POLICIES, ATTACHMENTS) + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert record["level"] == "INFO" + assert record["message"] == "\n".join( + ( + "LiteLLM Policy Engine: Loaded 1 policies", + "", + *POLICY_ROWS, + "", + "Policy Attachments: 1 attachment(s)", + " - pii-policy -> scope=* (global), teams=['team-a']", + ) + ) + + +def test_missing_attachments_warning_is_a_second_json_record( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """Without attachments the table and the yellow warning are two INFO records on stdout, where the two + prints put them.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + + _print_policies_on_startup(POLICIES, None) + + out, err = capfd.readouterr() + assert err == "" + records: Final = [json.loads(raw) for raw in out.splitlines() if raw.strip()] + assert [record["message"] for record in records] == [ + "\n".join(("LiteLLM Policy Engine: Loaded 1 policies", "", *POLICY_ROWS)), + NO_ATTACHMENTS_WARNING, + ] + assert {record["level"] for record in records} == {"INFO"} diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 89a35923be0..d77731aa83b 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -1,4 +1,5 @@ import inspect +import json import logging import os from contextlib import nullcontext @@ -9,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import click import fastapi +import httpx import pytest @@ -23,6 +25,8 @@ import yaml from uvicorn.config import LOOP_FACTORIES from uvicorn.importer import import_from_string +import litellm +from litellm._logging import turn_on_json from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS from litellm.proxy import proxy_cli from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server @@ -2875,6 +2879,26 @@ class TestRunServerDbSetup: mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) assert ("--enforce_prisma_migration_check is deprecated and has no effect" in capsys.readouterr().out) is warned + def test_the_retired_enforce_prisma_migration_check_opt_in_warns_on_health( + self, capsys: pytest.CaptureFixture[str] + ) -> None: + """`--health` returns before the config is read, where a server start emits the deprecation line; + it still prints it, as it did when the notice came before that return.""" + with patch("httpx.get", return_value=httpx.Response(200, json={})) as mock_get: + run_server.main(["--health", "--enforce_prisma_migration_check"], standalone_mode=False) + + mock_get.assert_called_once_with(url="http://0.0.0.0:4000/health") + assert "--enforce_prisma_migration_check is deprecated and has no effect" in capsys.readouterr().out + + def test_the_retired_enforce_prisma_migration_check_opt_in_warns_on_test( + self, capsys: pytest.CaptureFixture[str] + ) -> None: + """`--test` returns before the config is read too, and prints the deprecation line before its request.""" + with pytest.raises(ValueError, match="Invalid test value"): + run_server.main(["--test", "--enforce_prisma_migration_check"], standalone_mode=False) + + assert "--enforce_prisma_migration_check is deprecated and has no effect" in capsys.readouterr().out + @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database", return_value=True) @@ -3935,3 +3959,117 @@ class TestValidateConfigFlag: assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "config OK" in result.output mock_popen.assert_not_called() + + +def _boot_from_config_file(config_path: Path, *flags: str) -> str | None: + env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "JSON_LOGS")} + with patch.dict(os.environ, env, clear=True): + run_server.main(["--config", str(config_path), "--skip_server_startup", *flags], standalone_mode=False) + return os.environ.get("JSON_LOGS") + + +class TestJsonLogsConfigReachesSpawnedWorkers: + """`litellm_settings.json_logs: true` turns JSON on inside the CLI process, but the workers that + uvicorn or granian spawn are fresh interpreters that pick plain or JSON from the JSON_LOGS + environment variable at import. The CLI has to export it, or every worker logs plain text until + its own load_config reaches the key.""" + + def test_json_logs_true_exports_json_logs_to_spawned_workers( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {"json_logs": True}})) + + assert _boot_from_config_file(config_path) == "true" + + def test_plain_logs_config_leaves_json_logs_unset(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "json_logs", False) + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {}})) + + assert _boot_from_config_file(config_path) is None + + +class TestEnforcePrismaMigrationCheckNoticeAroundTheConfigRead: + """With a config file, the `--enforce_prisma_migration_check` deprecation notice waits until the config + is read, so a `json_logs: true` set only there makes it a JSON record. It is printed once either way, + and a boot whose config read fails still prints it before the error, as it did before the wait.""" + + NOTICE: Final = "--enforce_prisma_migration_check is deprecated and has no effect" + FLAG: Final = "--enforce_prisma_migration_check" + + def test_a_config_read_that_fails_still_prints_the_notice( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + monkeypatch.setattr(litellm, "json_logs", False) + + with pytest.raises(Exception, match="Config file not found"): + _boot_from_config_file(tmp_path / "missing.yaml", self.FLAG) + + assert capsys.readouterr().out.count(self.NOTICE) == 1 + + def test_a_plain_config_prints_the_notice_once( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + monkeypatch.setattr(litellm, "json_logs", False) + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {}})) + + _boot_from_config_file(config_path, self.FLAG) + + assert capsys.readouterr().out.count(self.NOTICE) == 1 + + def test_a_config_only_json_logs_makes_the_notice_one_json_record( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {"json_logs": True}})) + + _boot_from_config_file(config_path, self.FLAG) + + notice_lines: Final = [line for line in capsys.readouterr().out.splitlines() if self.NOTICE in line] + assert len(notice_lines) == 1, notice_lines + assert self.NOTICE in json.loads(notice_lines[0])["message"] + + +def test_granian_flag_notices_are_json_records_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """The yellow notices for the flags Granian ignores (`--max_requests_before_restart`, `--ciphers`) are + printed while the launcher starts; with json_logs on they are INFO records on stdout like every other + line, not plaintext between the records.""" + pytest.importorskip("granian") + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + + with ( + patch("granian.Granian") as mock_granian_cls, + patch("granian.constants.Interfaces", SimpleNamespace(ASGI="asgi")), + ): + ProxyInitializationHelpers._init_granian_server( + host="0.0.0.0", + port=4000, + num_workers=2, + ssl_certfile_path=None, + ssl_keyfile_path=None, + max_requests_before_restart=100, + ciphers="ECDHE", + granian_runtime_threads=None, + ) + + mock_granian_cls.return_value.serve.assert_called_once() + out, err = capfd.readouterr() + assert err == "" + records: Final = [json.loads(raw) for raw in out.splitlines() if raw.strip()] + assert [record["message"] for record in records] == [ + "LiteLLM Proxy: Starting server on 0.0.0.0:4000 using Granian", + "LiteLLM: --max_requests_before_restart is not supported by Granian " + "(Granian uses workers_lifetime in seconds, not a per-request limit).", + "LiteLLM: --ciphers is not applied when using --run_granian.", + ] + assert {record["level"] for record in records} == {"INFO"} diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 8c62b225992..55d1587669a 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -7,6 +7,7 @@ import os import re import socket import subprocess +import sys import time import types from datetime import datetime, timedelta, timezone @@ -28,6 +29,7 @@ from fastapi.testclient import TestClient import litellm import litellm.proxy.proxy_server as proxy_server_module from litellm._internal_context import current_service_target +from litellm._logging import turn_on_json, verbose_proxy_logger from litellm.caching.caching import RedisCache from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cluster_cache import RedisClusterCache @@ -46,6 +48,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.proxy.common_utils.banner import LITELLM_BANNER from tests._master_key import MASTER_KEY PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234" @@ -14266,7 +14269,6 @@ async def test_moderations_response_carries_litellm_call_id_header(): async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplog): """LIT-7836: the /v1/moderations error line must carry the litellm_call_id the client sent, rendered in the message and as a structured log record field.""" - from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyException call_id = "moderations-call-7836" @@ -15586,3 +15588,61 @@ async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is call(api_provider="openai", capture_rate=0.97), call(api_provider="openai", capture_rate=None), ] + + +def test_feedback_box_is_one_json_record_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """The feedback box is a dozen terminal lines; with JSON logs on it has to land as one parseable + record, not a plaintext line per row.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + + proxy_server_module.generate_feedback_box() + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert "Give Feedback / Get Help: https://github.com/BerriAI/litellm/issues/new" in record["message"] + assert "Thank you for using LiteLLM! - Krrish & Ishaan" in record["message"] + assert "\x1b[" not in record["message"] + + +@pytest.mark.asyncio +async def test_initialize_prints_the_banner_between_blank_lines_in_plain_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """Operators without JSON logs see the banner exactly as before, through show_banner()'s flushing + click.echo.""" + monkeypatch.setattr(litellm, "json_logs", False) + monkeypatch.setenv("LITELLM_DONT_SHOW_FEEDBACK_BOX", "true") + + await proxy_server_module.initialize() + + out, err = capfd.readouterr() + assert err == "" + assert out.startswith(f"\n{LITELLM_BANNER}\n\n"), out + + +@pytest.mark.asyncio +async def test_initialize_prints_the_banner_as_one_json_record_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """With JSON logs on, the six banner rows are one INFO record with their indentation kept, so the + first line a worker writes parses like every other.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + monkeypatch.setenv("LITELLM_DONT_SHOW_FEEDBACK_BOX", "true") + turn_on_json() + + await proxy_server_module.initialize() + + out, err = capfd.readouterr() + assert err == "" + records: Final = [json.loads(raw) for raw in out.splitlines() if raw.strip()] + banner_records: Final = [record for record in records if record["message"] == LITELLM_BANNER] + assert len(banner_records) == 1, records + assert banner_records[0]["level"] == "INFO" diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py index cb4d3f82611..36be864a6f7 100644 --- a/tests/unit/test_logging.py +++ b/tests/unit/test_logging.py @@ -34,6 +34,7 @@ from litellm._logging import ( _parse_json_logs_env, _plain_log_format, _stdout_truncation_marker, + print_startup_message, turn_on_json, format_base64_size, session_id_var, @@ -1824,3 +1825,222 @@ def test_print_verbose_formats_args_only_when_set_verbose_is_on(monkeypatch, cap logging_module.print_verbose("key %s value %s", "k1", _StrCountingValue()) assert _StrCountingValue.str_calls == 1 assert capsys.readouterr().out == "key k1 value cached-vector\n" + + +def test_print_startup_message_prints_plain_mode_text_unchanged( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """Operators without JSON logs keep the colored, spaced terminal output exactly as before.""" + monkeypatch.setattr(litellm, "json_logs", False) + + print_startup_message("\n\033[1;32mLiteLLM: Proxy initialized with Config, Set models:\033[0m\n") + + out, err = capfd.readouterr() + assert out == "\n\033[1;32mLiteLLM: Proxy initialized with Config, Set models:\033[0m\n\n" + assert err == "" + + +def test_print_startup_message_is_one_json_record_in_json_mode( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """With JSON logs on, a multi-line colored startup message is one INFO record with the colors and the + surrounding blank lines removed, so a log pipeline parses it like every other line. The logger sits + at WARNING, as it does in the CLI process, and the record still goes out: a print has no level.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + verbose_proxy_logger.setLevel(logging.WARNING) + + print_startup_message( + "\n\033[1;32mLiteLLM: Proxy initialized with Config, Set models:\033[0m\n\033[32m gpt-4o\033[0m\n\n" + ) + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert record["message"] == "LiteLLM: Proxy initialized with Config, Set models:\n gpt-4o" + assert record["level"] == "INFO" + assert record["component"] == "LiteLLM Proxy" + assert str(record["logger"]).startswith("test_logging.py:"), record["logger"] + + +def test_print_startup_message_yellow_notice_is_an_info_record_on_stdout( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """An operator notice printed in yellow (an ignored flag, a deprecation) stays on stdout in JSON mode, + where the print put it, as one INFO record with the color codes removed.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + verbose_proxy_logger.setLevel(logging.WARNING) + + print_startup_message("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert record["message"] == "LiteLLM: --ciphers is not applied when using --run_granian." + assert record["level"] == "INFO" + + +def test_print_startup_message_shows_above_the_json_handler_level( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """`LITELLM_LOG=WARNING` sets the JSON handler to WARNING, which drops the INFO logger lines; the + startup notices still go out, as the prints they replace did at every level.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + verbose_proxy_logger.setLevel(logging.INFO) + for json_handler in verbose_proxy_logger.handlers: + monkeypatch.setattr(json_handler, "level", logging.WARNING) + + print_startup_message("\033[94m\nSetting Cache on Proxy\033[0m") + verbose_proxy_logger.info("enable_redis_auth_cache is not set") + + out, _ = capfd.readouterr() + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert [json.loads(raw)["message"] for raw in lines] == ["Setting Cache on Proxy"] + + +def test_print_startup_message_emits_json_after_an_operator_log_config_removed_the_handlers( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """uvicorn applies an operator `--log_config` after proxy_cli's turn_on_json() and before initialize() + re-installs the JSON handler through load_config(); a config that names "LiteLLM Proxy" leaves the + logger without handlers in that window, where the banner and the feedback box are emitted, and they + still go out as one JSON record each, on stdout.""" + import logging.config + + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + request.addfinalizer(turn_on_json) + logging.config.dictConfig( + {"version": 1, "disable_existing_loggers": False, "loggers": {"LiteLLM Proxy": {"level": "INFO"}}} + ) + assert verbose_proxy_logger.handlers == [] + + print_startup_message("\n\033[1;37m# Give Feedback / Get Help #\033[0m\n") + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert record["message"] == "# Give Feedback / Get Help #" + assert record["level"] == "INFO" + assert record["component"] == "LiteLLM Proxy" + + +OPERATOR_HANDLERS_ON_THE_PROXY_LOGGER: Final = { + "silenced": {"class": "logging.NullHandler"}, + "stderr_warning": { + "class": "logging.StreamHandler", + "stream": "ext://sys.stderr", + "level": "WARNING", + "formatter": "plain", + }, +} + + +@pytest.mark.parametrize("handler", sorted(OPERATOR_HANDLERS_ON_THE_PROXY_LOGGER)) +def test_print_startup_message_stays_on_stdout_under_an_operator_handler_on_the_proxy_logger( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str], handler: str +) -> None: + """An operator `--log_config` that gives "LiteLLM Proxy" a handler of its own, a NullHandler that silences + it or a plain WARNING-level handler on stderr, is in force when the banner and the feedback box are + emitted; the print they replace wrote to stdout whatever the logging config said, so the record still + goes out as one JSON record on stdout and the operator's handler gets nothing.""" + import logging.config + + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + request.addfinalizer(turn_on_json) + logging.config.dictConfig( + { + "version": 1, + "disable_existing_loggers": False, + "formatters": {"plain": {"format": "%(levelname)s %(message)s"}}, + "handlers": {"proxy": OPERATOR_HANDLERS_ON_THE_PROXY_LOGGER[handler]}, + "loggers": {"LiteLLM Proxy": {"level": "INFO", "handlers": ["proxy"], "propagate": False}}, + } + ) + assert len(verbose_proxy_logger.handlers) == 1 and not isinstance( + verbose_proxy_logger.handlers[0], LevelRoutingStreamHandler + ) + + print_startup_message("\n\033[1;37m# Give Feedback / Get Help #\033[0m\n") + + out, err = capfd.readouterr() + assert err == "", err + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert record["message"] == "# Give Feedback / Get Help #" + assert record["level"] == "INFO" + assert record["component"] == "LiteLLM Proxy" + + +def test_print_startup_message_keeps_a_listing_longer_than_the_stdout_cap_whole( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """The required migration SQL and the policy list run past MAX_STRING_LENGTH_STDOUT_LOG on a drifted + deployment. The prints they replace showed them whole, and that cap exists for echoed request payloads, + so the JSON record carries the whole listing too.""" + monkeypatch.setenv("MAX_STRING_LENGTH_STDOUT_LOG", "500") + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + commands: Final = tuple(f'ALTER TABLE "LiteLLM_Table{n}" ADD COLUMN "column_{n}" TEXT;' for n in range(40)) + listing: Final = "\n".join(("Changes to DB Schema detected", "Required SQL commands:", *commands)) + assert len(listing) > 500 + + print_startup_message(listing) + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + record: Final = json.loads(lines[0]) + assert LITELLM_TRUNCATED_PAYLOAD_FIELD not in record["message"] + assert record["message"] == listing + + +def test_print_startup_message_keeps_a_model_name_the_secret_patterns_would_match( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """The model list printed `risk-assessment-gpt4o` as configured; the sk- key pattern matches inside + that name, so the JSON record has to carry the name as the print did, not `riREDACTED`.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + assert secret_redaction.redact_string("risk-assessment-gpt4o") != "risk-assessment-gpt4o" + + print_startup_message("\033[32m risk-assessment-gpt4o\033[0m") + + out, err = capfd.readouterr() + assert err == "" + lines: Final = [raw for raw in out.splitlines() if raw.strip()] + assert len(lines) == 1, lines + assert json.loads(lines[0])["message"] == " risk-assessment-gpt4o" + + +def test_print_startup_message_blank_spacer_emits_no_json_record( + monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str] +) -> None: + """A spacer that only pads the terminal has nothing to say to a log pipeline.""" + monkeypatch.setattr(litellm, "json_logs", True) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + turn_on_json() + verbose_proxy_logger.setLevel(logging.INFO) + + print_startup_message("") + print_startup_message("\n\033[0m\n") + + assert capfd.readouterr() == ("", "") From 8f31c7171dfbd28c53d294121a28c5d232d9137c Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 27 Oct 2026 15:12:53 -0700 Subject: [PATCH 11/83] fix(router): keep a lone deployment warm under router allowed_fails unless the policy names the exception (#48644) * fix(router): keep a lone deployment warm under router allowed_fails unless the policy names the exception Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): give cooldown fixtures a sibling deployment so allowed_fails=0 still forces a cooldown * fix(router): keep benching a lone deployment that has a fallback group or an empty deployment policy * fix(router): accept dictionary fallback targets in the lone deployment check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_utils/cooldown_handlers.py | 44 ++++++- .../routing/test_payment_required_cooldown.py | 74 +++++++++++ .../router_utils/test_cooldown_handlers.py | 119 ++++++++++++++---- .../test_fallback_event_handlers.py | 14 ++- tests/unit/test_router/test_router.py | 14 ++- 5 files changed, 234 insertions(+), 31 deletions(-) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index c362c6899db..330500789a4 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -185,7 +185,14 @@ def _should_cooldown_based_on_deployment_policy( override that safety net, matching `_has_explicit_allowed_fails_policy_for_exception`. """ allowed_fails_from_policy: Final = _resolve_allowed_fails_from_policy(dep_policy, original_exception) - if allowed_fails_from_policy is None and dep_allowed_fails is not None and is_single_deployment_model_group: + if ( + allowed_fails_from_policy is None + and is_single_deployment_model_group + and ( + dep_allowed_fails is not None + or litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None + ) + ): return False allowed_fails_override: Final[int | None] = ( @@ -399,6 +406,34 @@ def _should_run_cooldown_logic( return True +_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, list[str | Mapping[str, object]]] | str, ...]) + + +def _fallback_target_name(target: str | Mapping[str, object]) -> str: + return target if isinstance(target, str) else str(target.get("model", "")) + + +def _fallback_chain_names( + entry: Mapping[str, Sequence[str | Mapping[str, object]]], +) -> Mapping[str, list[str]]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape + return {key: [_fallback_target_name(target) for target in targets] for key, targets in entry.items()} + + +def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool: + from litellm.router_utils.fallback_event_handlers import get_fallback_model_group + + configured: Final[object] = getattr(litellm_router_instance, "fallbacks", None) + if requested_model_group is None or not configured: + return False + chains: Final = tuple( + _fallback_chain_names(entry) + for entry in _FALLBACK_ENTRIES_ADAPTER.validate_python(configured) + if isinstance(entry, Mapping) + ) + fallback_group, _ = get_fallback_model_group(fallbacks=chains, model_group=requested_model_group) + return bool(fallback_group) + + def _should_cooldown_deployment( litellm_router_instance: LitellmRouter, deployment: str, @@ -431,6 +466,9 @@ def _should_cooldown_deployment( litellm_router_instance.routing_group_has_alternatives(requested_model_group) or litellm_router_instance.team_model_has_alternatives(deployment) ) + is_lone_deployment: Final = is_single_deployment_model_group and not _has_fallback_group( + litellm_router_instance, requested_model_group + ) exception_status_int: Final = cast_exception_status_to_int(exception_status) dep_policy, dep_allowed_fails = _get_deployment_cooldown_policy(litellm_router_instance, deployment) @@ -450,7 +488,7 @@ def _should_cooldown_deployment( original_exception, dep_policy, dep_allowed_fails, - is_single_deployment_model_group, + is_lone_deployment, ) ## BASE CASE - single deployment @@ -495,6 +533,8 @@ def _should_cooldown_deployment( return False else: + if is_lone_deployment and litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None: + return False return should_cooldown_based_on_allowed_fails_policy( litellm_router_instance=litellm_router_instance, deployment=deployment, diff --git a/tests/integration/routing/test_payment_required_cooldown.py b/tests/integration/routing/test_payment_required_cooldown.py index bc5cf1a5913..865edf31260 100644 --- a/tests/integration/routing/test_payment_required_cooldown.py +++ b/tests/integration/routing/test_payment_required_cooldown.py @@ -1,5 +1,6 @@ from __future__ import annotations +import itertools import json import re import uuid @@ -390,6 +391,79 @@ def test_a_router_allowed_fails_policy_cools_a_lone_402_deployment_past_its_coun assert markers_received(refusing) == (first.marker, second.marker) +def failing_reply(request: Request) -> Reply: + return model_list_reply() if is_model_info_probe(request) else anthropic_error(500, "upstream down") + + +def assert_failed(outcome: Outcome) -> None: + assert outcome.status == 500, (outcome.status, outcome.detail) + assert "upstream down" in outcome.detail, outcome.detail + + +@pytest.mark.timeout(900) +def test_a_router_allowed_fails_count_leaves_a_lone_deployment_warm_but_benches_one_with_a_sibling_or_fallback( + tmp_path: Path, +) -> None: + lone_group: Final = new_group() + pair_group: Final = new_group() + primary_group: Final = new_group() + backup_group: Final = new_group() + with ExitStack() as stack: + upstream: Final = stack.enter_context(gateway_from_environment()) + cache: Final = stack.enter_context(owned_redis(tmp_path)) + settings: Final = router_settings(cache, allowed_fails=1, fallbacks=[{primary_group: [backup_group]}]) + owned: Final = stack.enter_context( + owned_proxy_process( + upstream, + tmp_path, + {}, + config=cooldown_config(tmp_path, settings), + remove_environment=("PROMETHEUS_MULTIPROC_DIR",), + ) + ) + failing: Final = stack.enter_context(wire_server(failing_reply)) + healthy: Final = stack.enter_context(wire_server(healthy_reply)) + scenario: Final = stack.enter_context(owned.gateway.scenario()) + lone_id: Final = deployment(scenario, lone_group, anthropic_params(failing)) + failed: Final = tuple(raw_chat(owned.gateway, lone_group, new_marker()) for _ in range(4)) + for outcome in failed: + assert_failed(outcome) + assert not cooled(cache, lone_id) + assert markers_received(failing) == markers(failed) + + paired_id: Final = two_deployments(scenario, pair_group, failing, healthy) + history: Final = tuple( + itertools.takewhile( + lambda outcome: not cooled(cache, paired_id), + (raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(40)), + ) + ) + assert cooled(cache, paired_id), [(outcome.status, outcome.detail[:80]) for outcome in history] + assert sum(outcome.status == 500 for outcome in history) == 1 + served: Final = tuple(raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(10)) + assert_served(served, chat_identity) + assert markers_received(healthy) == (*markers(o for o in history if o.status == 200), *markers(served)) + + primary_id: Final = deployment(scenario, primary_group, anthropic_params(failing)) + deployment(scenario, backup_group, openai_params(healthy)) + failing.drain() + healthy.drain() + fell_back: Final = tuple( + itertools.takewhile( + lambda outcome: not cooled(cache, primary_id), + (raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(40)), + ) + ) + assert cooled(cache, primary_id) and fell_back, [(o.status, o.detail[:80]) for o in fell_back] + assert_served(fell_back, chat_identity) + backed: Final = tuple(raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(4)) + assert_served(backed, chat_identity) + received: Final = markers_received(failing) + assert received[: len(fell_back)] == markers(fell_back) + assert len(received) <= len(fell_back) + 1, received + assert {*markers(fell_back), *markers(backed)} <= set(markers_received(healthy)) + + def test_a_deployment_allowed_fails_policy_cools_a_lone_402_deployment(rig: CooldownRig) -> None: group: Final = new_group() policy: Final = {"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}} diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 0c9d7ab9f85..bab6d257579 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -2931,10 +2931,65 @@ async def test_bad_request_error_does_not_cool_down_a_deployment( assert tuple(cooldown_deployments) == () -@pytest.mark.parametrize("use_policy", [False, True], ids=["allowed_fails", "allowed_fails_policy"]) +TIMED_OUT_DEPLOYMENT: Final = { + "model_name": "test-model", + "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}, + "model_info": {"id": "test-deployment"}, +} +BACKUP_DEPLOYMENT: Final = { + "model_name": "backup-model", + "litellm_params": {"model": "openai/backup-model", "api_key": "test-key", "api_base": "https://backup.test/v1"}, + "model_info": {"id": "backup-deployment"}, +} +BACKUP_REPLY: Final = { + "id": "chatcmpl-backup", + "object": "chat.completion", + "created": 1700000000, + "model": "backup-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +async def _lone_call_outcome(router: Router) -> type[Exception]: + try: + await router.acompletion(model="test-model", messages=[{"role": "user", "content": "timeout test"}]) + except (litellm.Timeout, RouterRateLimitError) as error: + return type(error) + pytest.fail("the timed-out deployment answered") + + +BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, RouterRateLimitError, RouterRateLimitError) + + +@pytest.mark.parametrize( + ("router_kwargs", "model_info", "outcomes"), + [ + ({"allowed_fails": 1}, {}, (litellm.Timeout,) * 5), + ({"allowed_fails": 1}, {"allowed_fails_policy": {}}, (litellm.Timeout,) * 5), + ({"allowed_fails_policy": AllowedFailsPolicy(RateLimitErrorAllowedFails=1)}, {}, (litellm.Timeout,) * 5), + ({"allowed_fails_policy": AllowedFailsPolicy(TimeoutErrorAllowedFails=1)}, {}, BENCHED), + ({"allowed_fails": 1, "fallbacks": [{"test-model": ["backup-model"]}]}, {}, BENCHED), + ( + {"allowed_fails": 1, "fallbacks": [{"test-model": [{"model": "backup-model", "temperature": 0}]}]}, + {}, + BENCHED, + ), + ], + ids=[ + "allowed_fails", + "allowed_fails_with_an_empty_deployment_policy", + "policy_names_another_exception", + "policy_names_timeout", + "allowed_fails_with_a_fallback_group", + "allowed_fails_with_a_dict_fallback_target", + ], +) @pytest.mark.asyncio -async def test_allowed_fails_cools_a_single_deployment( - use_policy: bool, +async def test_router_allowed_fails_benches_a_lone_deployment_only_when_the_policy_names_the_exception( + router_kwargs: Mapping[str, object], + model_info: Mapping[str, object], + outcomes: tuple[type[Exception], ...], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -2942,38 +2997,52 @@ async def test_allowed_fails_cools_a_single_deployment( route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=httpx.ReadTimeout("upstream timeout") ) - policy_kwargs: Final = ( - {"allowed_fails_policy": AllowedFailsPolicy(TimeoutErrorAllowedFails=1)} - if use_policy - else {"allowed_fails": 1} + deployment: Final = {**TIMED_OUT_DEPLOYMENT, "model_info": {**TIMED_OUT_DEPLOYMENT["model_info"], **model_info}} + router: Final = Router(model_list=[deployment], cooldown_time=300, num_retries=0, **router_kwargs) + + seen: Final = tuple([await _lone_call_outcome(router) for _ in range(5)]) + + assert seen == outcomes + assert len(route.calls) == outcomes.count(litellm.Timeout) + + +async def _group_call_outcome(router: Router) -> str: + try: + response: Final = await router.acompletion( + model="test-model", messages=[{"role": "user", "content": "timeout test"}] + ) + except litellm.Timeout: + return "timeout" + return response.id + + +@pytest.mark.asyncio +async def test_router_allowed_fails_still_benches_a_deployment_that_has_a_sibling_in_its_group( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + timed_out: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=httpx.ReadTimeout("upstream timeout") ) + respx_mock.post("https://backup.test/v1/chat/completions").mock(return_value=httpx.Response(200, json=BACKUP_REPLY)) router: Final = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}, - "model_info": {"id": "test-deployment"}, - } - ], + model_list=[TIMED_OUT_DEPLOYMENT, {**BACKUP_DEPLOYMENT, "model_name": "test-model"}], cooldown_time=300, num_retries=0, - **policy_kwargs, + allowed_fails=1, ) - for _ in range(2): - with pytest.raises(litellm.Timeout): - await router.acompletion( - model="test-model", - messages=[{"role": "user", "content": "timeout test"}], - ) - + while len(timed_out.calls) < 2: + await _group_call_outcome(router) cooldown_deployments: Final = await async_get_cooldown_deployments( - litellm_router_instance=router, - parent_otel_span=None, + litellm_router_instance=router, parent_otel_span=None ) + served: Final = tuple([await _group_call_outcome(router) for _ in range(5)]) - assert len(route.calls) == 2 assert tuple(cooldown_deployments) == ("test-deployment",) + assert served == ("chatcmpl-backup",) * 5 + assert len(timed_out.calls) == 2 def test_per_user_session_upstream_errors_never_cool_down_the_shared_deployment(single_deployment_router): diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 53375ee0ab4..c5a476c450b 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -1255,7 +1255,12 @@ class TestTriggerCooldownForFailedDeployment: "model_name": "fallback-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "fallback-deployment"}, - } + }, + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "fallback-sibling"}, + }, ], allowed_fails=0, cooldown_time=60, @@ -1345,7 +1350,12 @@ async def test_a_fallback_deployment_whose_stream_stalled_stays_out_of_cooldown( "model_name": "fallback-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "fallback-deployment"}, - } + }, + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "fallback-sibling"}, + }, ], allowed_fails=0, cooldown_time=60, diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index e2ee80b7874..d9b655a0aff 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -12212,7 +12212,12 @@ class TestBackgroundResponseCostPollCooldown: "model_name": "gpt-4.1", "litellm_params": {"model": "openai/gpt-4.1"}, "model_info": {"id": "dep-1"}, - } + }, + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "openai/gpt-4.1"}, + "model_info": {"id": "dep-2"}, + }, ], allowed_fails=0, ) @@ -12340,7 +12345,12 @@ class TestCallerTimeoutCooldown: "model_name": "slow-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "dep-1"}, - } + }, + { + "model_name": "slow-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "dep-2"}, + }, ], allowed_fails=0, cooldown_time=120, From cb70bc560e1bbd8a1d619fb5dead7f2f0c333e74 Mon Sep 17 00:00:00 2001 From: Ambuj Upadhyay Date: Wed, 28 Oct 2026 03:51:11 +0530 Subject: [PATCH 12/83] fix(streaming): dict-usage arm drops prompt_tokens_details/completion_tokens_details its sibling arms preserve (#36678) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(streaming): dict-usage arm no longer drops token detail fields In _dispatch_provider_chunk's openai/azure else-branch, a plain-dict usage was rebuilt into litellm.Usage from only prompt_tokens / completion_tokens / total_tokens. prompt_tokens_details (cached_tokens), completion_tokens_details (reasoning_tokens) and cache token counts were silently dropped — while both sibling arms (Usage passthrough, BaseModel via Usage(**model_dump())) preserve them. Every in-repo chunk producer currently normalizes usage to Usage/BaseModel before this point, so the dict arm is reachable today only via chunks litellm did not construct (e.g. a custom streaming_decoder feeding CustomStreamWrapper). When it does fire on a usage-only final chunk (stream_options.include_usage), the stripped Usage is exactly what the stream consumer receives — the early-return path has no later aggregation to repair it. Fix: pass the full dict through Usage(**usage_dict). Usage.__init__ coerces nested detail dicts into their wrapper types and tolerates extra keys. Side effect: a legitimate 0 count is no longer coerced to None by the old 'or None' expressions. Co-Authored-By: Claude Fable 5 * chore: drop explanatory source comment per repo comment policy Rationale lives in the PR body and the regression tests. Co-Authored-By: Claude Fable 5 * fix(streaming): type the dict-usage unpack for the basedpyright gate Usage(**response_obj["usage"]) unpacks a dict whose value type pyright cannot see, and the delta gate reports that once per typed constructor parameter: five, against the three .get() calls the line replaced. Moving the unpack behind _usage_from_payload keeps the same Usage(**payload) call and is reported once, at the call site, so the rule count ends below its merge base --------- Co-authored-by: Claude Fable 5 --- .../litellm_core_utils/streaming_handler.py | 11 ++- .../test_streaming_handler.py | 99 +++++++++++++++++++ 2 files changed, 105 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 8fb602ebeb0..b64e146ca08 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1589,11 +1589,7 @@ class CustomStreamWrapper: setattr( model_response, "usage", - litellm.Usage( - prompt_tokens=response_obj["usage"].get("prompt_tokens", None) or None, - completion_tokens=response_obj["usage"].get("completion_tokens", None) or None, - total_tokens=response_obj["usage"].get("total_tokens", None) or None, - ), + _usage_from_payload(response_obj["usage"]), ) elif isinstance(response_obj["usage"], Usage): setattr( @@ -2662,6 +2658,11 @@ class CustomStreamWrapper: _strip_sse_data_from_chunk = strip_sse_data_from_chunk +def _usage_from_payload(payload: Mapping[str, Any]) -> Usage: + """Build ``Usage`` from a wire-shaped usage dict, forwarding every key.""" + return litellm.Usage(**payload) + + def _cache_token_count(details: PromptTokensDetailsWrapper | None, keys: tuple[str, ...]) -> int: for key in keys: value = getattr(details, key, None) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index cc1fa1b8803..6d700bae514 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -7844,3 +7844,102 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream(): async with _wrapper_over(stream): pass assert stream.close_calls == 1 + + +def _final_chunk_with_dict_usage(choices) -> ModelResponseStream: + """A stream chunk whose ``usage`` is a plain dict. + + ``ModelResponseStream.__init__`` coerces a dict ``usage`` kwarg into + ``litellm.Usage``, so the dict is attached after construction — the same + shape a custom ``streaming_decoder`` (or any third-party iterator that + bypasses litellm chunk construction) delivers to ``chunk_creator``. + """ + chunk: Final = ModelResponseStream( + id="chatcmpl-dict-usage-test", + created=1754900000, + model="gpt-4o", + object="chat.completion.chunk", + choices=choices, + ) + chunk.usage = { # mutable-ok: wire-shaped usage payload under test + "prompt_tokens": 100, + "completion_tokens": 10, + "total_tokens": 110, + "prompt_tokens_details": {"cached_tokens": 80}, # mutable-ok: wire payload + "completion_tokens_details": {"reasoning_tokens": 5}, # mutable-ok: wire payload + "cache_creation_input_tokens": 42, + } + return chunk + + +def test_dispatch_provider_chunk_dict_usage_preserves_token_details( + logging_obj: Logging, +): + """The openai else-branch of ``_dispatch_provider_chunk`` handles three + usage shapes: ``Usage`` (passed through), ``BaseModel`` + (``Usage(**model_dump())``), and plain ``dict``. The dict arm used to + rebuild ``litellm.Usage`` from only prompt/completion/total tokens, + silently dropping ``prompt_tokens_details`` (cached_tokens), + ``completion_tokens_details`` (reasoning_tokens), and cache token counts + that both sibling arms preserve. + """ + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + stream_options={"include_usage": True}, # mutable-ok: constructor payload + ) + chunk: Final = _final_chunk_with_dict_usage( + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=None)) + ] # mutable-ok: chunk payload + ) + model_response: Final = wrapper.model_response_creator() + wrapper._dispatch_provider_chunk( + chunk=chunk, + model_response=model_response, + completion_obj={"content": ""}, # mutable-ok: dispatch scratch dict + ) + + usage: Final = model_response.usage + assert isinstance(usage, Usage) + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 10 + assert usage.total_tokens == 110 + assert usage.prompt_tokens_details is not None, ( + "dict-usage arm dropped prompt_tokens_details (its Usage/BaseModel sibling arms preserve it)" + ) + assert usage.prompt_tokens_details.cached_tokens == 80 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 5 + assert getattr(usage, "cache_creation_input_tokens", None) == 42 + + +def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details( + logging_obj: Logging, +): + """End-to-end through ``chunk_creator``: a usage-only final chunk (empty + ``choices``, ``stream_options.include_usage`` set) takes the early-return + path, so whatever the dict-usage arm wrote is exactly what the consumer + receives — there is no later aggregation to repair it. + """ + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + stream_options={"include_usage": True}, # mutable-ok: constructor payload + ) + chunk: Final = _final_chunk_with_dict_usage(choices=[]) # mutable-ok: empty choices list + result: Final = wrapper.chunk_creator(chunk=chunk) + + assert result is not None + usage: Final = result.usage + assert isinstance(usage, Usage) + assert usage.prompt_tokens_details is not None, ( + "usage-only early return delivered a Usage stripped of prompt_tokens_details to the stream consumer" + ) + assert usage.prompt_tokens_details.cached_tokens == 80 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 5 From 9b2fc9db40f8f976597388b4450ff002ee0aa732 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 27 Oct 2026 15:26:15 -0700 Subject: [PATCH 13/83] test: remove checks that pin implementation details and add a shadcn testing rule (#48663) * test: remove e2e checks and a CI line count that pin implementation details Deletes the Auto Router template select anchoring spec, the 390px logs pagination variant, and the call id tooltip and clipboard step in the logs spec. Each one tests library or test helper behavior rather than what users see, and each has been red or flaky on main without a product regression Also drops the exact line count on the fan-in requirements export in the integration coverage job. It went from 6 to 10 when uv.lock gained a Python 3.15 fork, while the install was still correct. A package the filter drops still fails the job, either at the hash-checked install or at the next step's import * docs(agents): do not test what shadcn and Base UI already guarantee Adds the boundary to the dashboard AGENTS.md next to the test tiers, and points the Playwright suite's entry in tests/e2e/AGENTS.md at it --- .circleci/config.yml | 1 - tests/e2e/AGENTS.md | 2 +- tests/e2e/ui/tests/logs/logs.spec.ts | 11 --- .../e2e/ui/tests/logs/logsPagination.spec.ts | 75 +++++++++---------- .../autoRouterTemplateSelect.spec.ts | 52 ------------- ui/litellm-dashboard/AGENTS.md | 2 + 6 files changed, 38 insertions(+), 105 deletions(-) delete mode 100644 tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts diff --git a/.circleci/config.yml b/.circleci/config.yml index 1edf22578ff..4ce5ea8b6fe 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1379,7 +1379,6 @@ jobs: uv export --frozen --all-groups --all-extras --no-emit-project --no-emit-workspace --no-header --no-annotate \ | awk '/^[^ ]/ { keep = $0 ~ /^(annotated-types|coverage|pydantic|pydantic-core|typing-extensions|typing-inspection)==/ } keep' \ > /tmp/fan-in-requirements.txt - test "$(grep -c '^[^ ]' /tmp/fan-in-requirements.txt)" -eq 6 uv venv .venv --python 3.12 uv pip install --no-config --python .venv --require-hashes -r /tmp/fan-in-requirements.txt - attach_workspace: diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index ef57135a4d2..cf1e41a59a6 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -49,7 +49,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture. Ships its own driver/builder/publisher, covered by `tests/e2e_harness/claude_code/`. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`). The CLI-driving path stays bespoke -- `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream, see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport. The Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` +- `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream, see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport. The Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json`. Its specs follow the dashboard's rule against testing what shadcn and Base UI already guarantee, in `ui/litellm-dashboard/AGENTS.md` ## MCP suite: real Datadog only diff --git a/tests/e2e/ui/tests/logs/logs.spec.ts b/tests/e2e/ui/tests/logs/logs.spec.ts index 688bfe8e3b4..49400f1509a 100644 --- a/tests/e2e/ui/tests/logs/logs.spec.ts +++ b/tests/e2e/ui/tests/logs/logs.spec.ts @@ -194,17 +194,6 @@ test.describe("Logs page", () => { await expect(row, `no logs row for call id ${callId}`).toHaveCount(1, { timeout: 30_000 }); await expect(row, "the row itself shows only the request id").not.toContainText(callId); - await row.getByText(requestId).hover(); - const tooltip = page.locator("[data-slot='tooltip-content']"); - await expect(tooltip, "hovering the Request ID cell does not list the x-litellm-call-id").toContainText( - `x-litellm-call-id: ${callId}`, - { timeout: 10_000 }, - ); - await tooltip.getByRole("button", { name: "Copy x-litellm-call-id" }).click(); - if (await page.evaluate(() => window.isSecureContext)) { - await expect.poll(() => page.evaluate(() => navigator.clipboard.readText())).toBe(callId); - } - await row.click(); const drawer = page.getByRole("dialog").first(); await expect(drawer.getByText("Request & Response")).toBeVisible({ timeout: 20_000 }); diff --git a/tests/e2e/ui/tests/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts index 838a829a3e2..585eb31c37c 100644 --- a/tests/e2e/ui/tests/logs/logsPagination.spec.ts +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -54,52 +54,47 @@ async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): test.describe("Logs pagination with Aria in the header", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); - for (const viewport of [ - { width: 1440, height: 900 }, - { width: 390, height: 844 }, - ]) { - test(`keeps pagination reachable and Aria available at ${viewport.width}px`, async ({ page, request }) => { - await page.setViewportSize(viewport); - const alias = `e2e-aria-header-${uniqueSuffix()}`; - const key = await createVirtualKey(request, { key_alias: alias }); - try { - const ids = await Promise.all( - Array.from({ length: 26 }, (_, index) => - sendChatCompletion(request, { - model: CHAT_MODEL_A, - prompt: `${alias}-${index}`, - apiKey: key.key, - }), - ), - ); - await Promise.all(ids.map((id) => waitForSpendLog(request, id))); - await openLogs(page); - await setAriaHidden(page, false); - await applyKeyAliasFilter(page, await openFilterDrawer(page), alias); - await setRowsPerPage(page, "25"); - await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); - await visibleTestId(page, "pagination-next").click(); - await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-26 of 26"); + test("keeps pagination reachable and Aria available at 1440px", async ({ page, request }) => { + await page.setViewportSize({ width: 1440, height: 900 }); + const alias = `e2e-aria-header-${uniqueSuffix()}`; + const key = await createVirtualKey(request, { key_alias: alias }); + try { + const ids = await Promise.all( + Array.from({ length: 26 }, (_, index) => + sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `${alias}-${index}`, + apiKey: key.key, + }), + ), + ); + await Promise.all(ids.map((id) => waitForSpendLog(request, id))); + await openLogs(page); + await setAriaHidden(page, false); + await applyKeyAliasFilter(page, await openFilterDrawer(page), alias); + await setRowsPerPage(page, "25"); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); + await visibleTestId(page, "pagination-next").click(); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-26 of 26"); - const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true }); + const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true }); + await aria.click(); + await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); + await page.keyboard.press("Escape"); + await expect(aria).toBeFocused(); + await visibleTestId(page, "pagination-first").click(); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); + for (const destination of [Page.Usage, Page.LlmPlayground, Page.ApiKeys]) { + await navigateToPage(page, destination); await aria.click(); await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); await page.keyboard.press("Escape"); await expect(aria).toBeFocused(); - await visibleTestId(page, "pagination-first").click(); - await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); - for (const destination of [Page.Usage, Page.LlmPlayground, Page.ApiKeys]) { - await navigateToPage(page, destination); - await aria.click(); - await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); - await page.keyboard.press("Escape"); - await expect(aria).toBeFocused(); - } - } finally { - await deleteVirtualKey(request, key.token); } - }); - } + } finally { + await deleteVirtualKey(request, key.token); + } + }); }); test.describe("Logs page session-grouped pagination", () => { diff --git a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts deleted file mode 100644 index 5f05953cc80..00000000000 --- a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts +++ /dev/null @@ -1,52 +0,0 @@ -import { expect, test, type Locator, type Page as PlaywrightPage } from "@playwright/test"; -import { ADMIN_STORAGE_PATH } from "../../constants"; -import { navigateToPage } from "../../helpers/navigation"; -import { Page } from "../../fixtures/pages"; - -/** - * Opens Add Auto Router and returns the Template select's trigger, which is the - * shallowest real page that renders SelectContent with tall multi-line options. - */ -async function openTemplateSelect(page: PlaywrightPage) { - await navigateToPage(page, Page.Models); - await page.getByRole("tab", { name: "Auto-Routers" }).click(); - await page.getByRole("button", { name: "Add Auto Router" }).click(); - - const trigger = page.getByTestId("template-selector"); - await expect(trigger).toBeVisible(); - return trigger; -} - -async function boxes(trigger: Locator, options: Locator) { - const triggerBox = await trigger.boundingBox(); - const optionsBox = await options.boundingBox(); - return triggerBox && optionsBox ? { triggerBox, optionsBox } : null; -} - -const clippedPopup = (page: PlaywrightPage) => page.locator('[data-slot="select-content"]'); - -function pollOptionsCoverTrigger(trigger: Locator, options: Locator) { - return expect.poll(async () => { - const box = await boxes(trigger, options); - return ( - box && - box.optionsBox.y < box.triggerBox.y + box.triggerBox.height && - box.optionsBox.y + box.optionsBox.height > box.triggerBox.y - ); - }); -} - -test.describe("Auto Router template select anchoring", () => { - test.use({ storageState: ADMIN_STORAGE_PATH }); - - test("keeps the trigger uncovered when the options open with no room below it", async ({ page }) => { - await page.setViewportSize({ width: 1280, height: 560 }); - const trigger = await openTemplateSelect(page); - await trigger.scrollIntoViewIfNeeded(); - - await trigger.click(); - await expect(page.getByRole("listbox")).toBeVisible(); - - await pollOptionsCoverTrigger(trigger, clippedPopup(page)).toBe(false); - }); -}); diff --git a/ui/litellm-dashboard/AGENTS.md b/ui/litellm-dashboard/AGENTS.md index b6fe5219022..f5576b1228b 100644 --- a/ui/litellm-dashboard/AGENTS.md +++ b/ui/litellm-dashboard/AGENTS.md @@ -8,6 +8,8 @@ When you fix lint violations that are grandfathered in `eslint-suppressions.json Tests come in three tiers, named by the standard definitions. `Foo.test.tsx` is a unit test: one module, collaborators replaced by doubles, no multi-component tree, and it should run in milliseconds. `Foo.integration.test.tsx` renders a real component tree with real children and only stubs the network boundary. It costs seconds per case, so it earns its place by proving wiring that a unit test cannot reach. Browser-level tests live in `tests/e2e/ui/` as Playwright specs against a live proxy +Do not test what shadcn and Base UI already guarantee, such as a tooltip opening on hover, a popup positioning itself against its trigger, focus trapping and Escape in a dialog, keyboard movement through a menu or select, and open and close animations. Those are the library's contract, and a test of them only fails when we restyle or move the page. Test what we put inside the primitive and what happens because of it, like the text a tooltip carries, the options a select offers, or the request a confirm button sends. Reach that content the most direct way, by a click, the keyboard, or the data the page renders, rather than through a hover race. If we change a primitive's behavior in `src/components/ui/`, its test lives next to that primitive as `*.test.tsx`, never in a page spec that borrows the page as a host + When a component holds logic worth asserting, extract the logic and unit-test it there rather than driving it through a render. `CreateMCPServer` is the worked example: its payload building lives in `createServerPayload.ts` with 46 unit tests that run in single-digit milliseconds, while `CreateMCPServer.integration.test.tsx` keeps only the cases that prove a form field reaches the right payload key. A test that renders a whole modal to assert the shape of one object belongs in the first category, not the second Most of the suite predates this split and is not yet classified, so an unsuffixed `*.test.tsx` is not evidence that a file is really a unit test. Classify what you touch From c4c2fe6ba12c9b76f43e8bf3cde5ce63efb91a2b Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 27 Oct 2026 18:38:07 -0400 Subject: [PATCH 14/83] docs(router): clarify heuristic v2 keyword settings (#48666) --- .../complexity_router/README.md | 27 ++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 08d39841b61..3df1803d264 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -13,7 +13,7 @@ Unlike the semantic `auto_router` which uses embedding-based matching, the `comp ## How It Works -The router scores each request across 7 dimensions: +The default `classifier_type: heuristic` scores each request across 7 dimensions. These weights, keyword lists and score boundaries apply to the weighted heuristic scorer. [`heuristic_v2`](#heuristic-v2) uses a separate success-probability predictor | Dimension | Description | Weight | |-----------|-------------|--------| @@ -254,11 +254,23 @@ model_list: No classifier model call or per-model training data is required. The classifier uses global tier quality, request-type quality, and similar-request cohorts from -the bundled UltraFeedback artifact. It estimates success at every tier, enforces +the bundled UltraFeedback artifact. Cohorts group prompts by request type, length, +code and math markers, multiple-choice formatting and non-ASCII character share. +Each tier's global estimate is `(successes + 1) / (observations + 2)`. Request-type +and cohort estimates combine their successes and observations with the broader +estimate as a prior, so sparse or missing statistics fall back toward broader data. +The predictor estimates success at every tier, enforces monotonic probabilities, and returns the first tier meeting the success threshold, or REASONING if no tier meets it. The existing complexity-router tier pool then selects and dispatches a model from that tier +`heuristic_v2` does not read `reasoning_keywords` or apply `reasoning_override_min_score`, +even when the floor is `0`. The weighted scorer's keyword lists, `dimension_weights`, +`token_thresholds` and `tier_boundaries` do not affect v2's probabilities or tier selection. +This also applies when a chain selects `local_heuristic: heuristic_v2`. Use +`classifier_type: heuristic` for the [reasoning-marker override](#reasoning-override). +The separate `keyword_tier_rules` routing overrides still run before classification + Set `heuristic_v2_success_threshold` to a value from 0 to 1 to override the artifact's threshold. For example, `0.9` requires a predicted success probability of at least 90%. Higher thresholds favor more capable tiers. Omit the setting or @@ -646,7 +658,16 @@ of that threshold, or when no tier meets the threshold. A zero margin still defe ### Reasoning Override -If 2+ reasoning markers are detected in the user message, the request is promoted to the REASONING tier even when the weighted score maps lower, so complex reasoning tasks get the appropriate model. The promotion requires the score to reach `reasoning_override_min_score`, which tracks `tier_boundaries.simple_medium` unless set, so stock phrases on an otherwise trivial prompt cannot buy the top tier. Set it to `0` to promote on the markers alone. +The weighted heuristic scorer promotes its result to REASONING when the user message matches 2+ entries in `reasoning_keywords` and the weighted score reaches `reasoning_override_min_score`. The floor tracks `tier_boundaries.simple_medium` unless set. An explicit `0` lowers the floor to zero, so a negative weighted score still does not qualify. `heuristic_v2` bypasses this scorer and does not apply the override + +For routing based on these reasoning keywords, set the following under your existing `complexity_router_config` and keep your tier mapping + +```yaml +classifier_type: heuristic +reasoning_override_min_score: 0 +``` + +Omit `reasoning_keywords` to use the built-in list, or supply a nonempty list to replace it. These scoring keywords are separate from `keyword_tier_rules`, which can select a tier before either classifier runs ### System Prompt Handling From 969161501acfca278642ea645130754e66dee846 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 27 Oct 2026 15:54:39 -0700 Subject: [PATCH 15/83] feat(ui): add LITELLM_UI_USE_SI_TERMINOLOGY to show SI instead of AI in the Admin UI (#47440) * feat(ui): add LITELLM_UI_USE_SI_TERMINOLOGY to show SI instead of AI in the Admin UI Setting LITELLM_UI_USE_SI_TERMINOLOGY=true on the proxy makes the Admin UI say "SI" where its own copy says "AI" and "Super Intelligence" where it says "Artificial Intelligence". This covers the nav, page headings, settings text, guardrail garden blurbs, the compliance playground prompts and the playground sample prompts. Third-party product, provider and law names keep their spelling (Vertex AI, Cisco AI Defense, EU AI Act, ...), and user data such as guardrail names, gateway URLs and keys is never rewritten. The proxy exposes the flag on /.well-known/litellm-ui-config, a TerminologyProvider at the app root reads it, and components call useTerminology() on their own strings. The context defaults to a pass-through, so nothing changes when the flag is unset. * fix(ui): apply SI terminology to the command palette, page visibility groups and guardrail garden search * test(ui): cover the flag-off wording on the palette, page visibility and garden search tests * fix(ui): apply SI terminology to the Aria usage assistant copy * test(ui): keep the Aria ui-config stub within the inline object lint budget * fix(ui): keep AI wording in prompts sent to models --- .../ui_discovery_endpoints.py | 2 + .../ui_discovery_endpoints.py | 1 + .../test_ui_discovery_endpoints.py | 22 ++++++ .../agents/_components/AgentsPanel.tsx | 7 +- .../_components/EvaluationSettingsModal.tsx | 6 +- .../_components/guardrail_garden.test.tsx | 30 +++++++- .../_components/guardrail_garden.tsx | 6 +- .../_components/guardrail_garden_card.tsx | 4 +- .../_components/guardrail_garden_detail.tsx | 6 +- .../_components/MCPGatewaySessionsTab.tsx | 17 +++-- .../mcp-servers/_components/mcp_connect.tsx | 12 ++-- .../components/chat_ui/RealtimePlayground.tsx | 8 ++- .../components/complianceUI/ComplianceUI.tsx | 4 +- .../_components/ai_suggestion_modal.tsx | 4 +- .../_components/guardrail_selection_modal.tsx | 4 +- .../policies/_components/policy_templates.tsx | 4 +- .../_components/template_parameter_modal.tsx | 4 +- .../UsageViewSelect/UsageViewSelect.tsx | 4 +- .../components/overview/TopAgents.tsx | 4 +- ui/litellm-dashboard/src/app/layout.tsx | 9 ++- .../src/components/AIHub/ModelHubTable.tsx | 4 +- .../CommandPalette.integration.test.tsx | 28 ++++++++ .../CommandPalette/CommandPalette.tsx | 4 +- .../components/CommandPalette/utils.test.ts | 11 +++ .../src/components/CommandPalette/utils.ts | 7 +- .../src/components/DashboardHeader.tsx | 4 +- .../Navbar/ViewSwitcher.integration.test.tsx | 24 +++++-- .../src/components/Navbar/ViewSwitcher.tsx | 9 ++- .../PageVisibilitySettings.test.tsx | 32 ++++++++- .../UISettings/PageVisibilitySettings.tsx | 8 ++- .../AdminSettings/UISettings/UISettings.tsx | 12 +++- .../src/components/leftnav.test.tsx | 19 ++++- .../src/components/leftnav.tsx | 15 ++-- .../liteadmin/LiteAdmin.integration.test.tsx | 69 ++++++++++++++----- .../src/components/liteadmin/LiteAdmin.tsx | 7 +- .../components/model_add/CredentialsPanel.tsx | 4 +- .../components/model_add/CredentialsTable.tsx | 4 +- .../src/components/moyai/MoyaiConnected.tsx | 4 +- .../src/components/networking.tsx | 1 + .../organisms/create_key_button.tsx | 11 ++- .../shared/CreatedKeyDisplay.test.tsx | 8 +++ .../components/shared/CreatedKeyDisplay.tsx | 9 ++- .../templates/KeyEditViewControls.tsx | 51 ++++++++------ .../src/contexts/TerminologyContext.test.tsx | 37 ++++++++++ .../src/contexts/TerminologyContext.tsx | 18 +++++ .../src/data/compliancePrompts.test.ts | 33 +++++++++ .../src/data/compliancePrompts.ts | 11 ++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 ++ .../src/lib/terminology.test.ts | 45 ++++++++++++ ui/litellm-dashboard/src/lib/terminology.ts | 28 ++++++++ 50 files changed, 567 insertions(+), 113 deletions(-) create mode 100644 ui/litellm-dashboard/src/contexts/TerminologyContext.test.tsx create mode 100644 ui/litellm-dashboard/src/contexts/TerminologyContext.tsx create mode 100644 ui/litellm-dashboard/src/data/compliancePrompts.test.ts create mode 100644 ui/litellm-dashboard/src/lib/terminology.test.ts create mode 100644 ui/litellm-dashboard/src/lib/terminology.ts diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 8225984a964..792ef8a3f8c 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -25,6 +25,7 @@ async def get_ui_config(): or general_settings.get("auto_redirect_ui_login_to_sso", False) is True ) admin_ui_disabled: Final = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true" + use_si_terminology: Final = os.getenv("LITELLM_UI_USE_SI_TERMINOLOGY", "false").lower() == "true" hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings) sso_configured: Final = has_user_setup_sso() @@ -49,4 +50,5 @@ async def get_ui_config(): is_control_plane=is_control_plane, workers=proxy_config.worker_registry if is_control_plane else [], mcp_stdio_enabled=is_mcp_stdio_enabled(), + use_si_terminology=use_si_terminology, ) diff --git a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py index e30ad26bfe5..d396c79b844 100644 --- a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -14,4 +14,5 @@ class UiDiscoveryEndpoints(LiteLLMBaseModel): hide_default_credentials_hint: bool = False is_control_plane: bool = False mcp_stdio_enabled: bool = False + use_si_terminology: bool = False workers: list[WorkerRegistryEntry] = Field(default=[]) diff --git a/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index 1ef0eee761a..d338f38d862 100644 --- a/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -501,3 +501,25 @@ def test_ui_config_tells_the_dashboard_whether_stdio_mcp_servers_are_enabled(mon assert response.status_code == 200 assert response.json()["mcp_stdio_enabled"] is expected + + +@pytest.mark.parametrize(("flag", "expected"), [(None, False), ("false", False), ("true", True), ("TRUE", True)]) +def test_ui_config_tells_the_dashboard_whether_to_use_si_terminology( + monkeypatch: pytest.MonkeyPatch, flag: str | None, expected: bool +) -> None: + if flag is None: + monkeypatch.delenv("LITELLM_UI_USE_SI_TERMINOLOGY", raising=False) + else: + monkeypatch.setenv("LITELLM_UI_USE_SI_TERMINOLOGY", flag) + app: Final[FastAPI] = FastAPI() + app.include_router(router) + + with ( + patch("litellm.proxy.utils.get_server_root_path", return_value="/"), + patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), + ): + response: Final = TestClient(app).get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + assert response.json()["use_si_terminology"] is expected diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx index b392e270d33..3f599c38976 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState, useEffect } from "react"; import { Info, Plus } from "lucide-react"; import { getAgentsList, deleteAgentCall } from "@/components/networking"; @@ -31,6 +32,7 @@ interface AgentsResponse { } const AgentsPanel: React.FC = ({ accessToken, userRole, teams }) => { + const term = useTerminology(); const [agentsList, setAgentsList] = useState([]); const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [isLoading, setIsLoading] = useState(true); @@ -140,8 +142,9 @@ const AgentsPanel: React.FC = ({ accessToken, userRole, teams

Agents

- List of A2A-spec agents that are available to be used in your organization. Go to AI Hub, to make agents - public. + {term( + "List of A2A-spec agents that are available to be used in your organization. Go to AI Hub, to make agents public.", + )}

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx index 900a04e480d..78091584cde 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import { Play } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models"; @@ -47,6 +48,7 @@ export function EvaluationSettingsModal({ accessToken, onRunEvaluation, }: EvaluationSettingsModalProps) { + const term = useTerminology(); const [prompt, setPrompt] = useState(DEFAULT_PROMPT); const [schema, setSchema] = useState(DEFAULT_SCHEMA); const [model, setModel] = useState(null); @@ -95,8 +97,8 @@ export function EvaluationSettingsModal({ Evaluation Settings {guardrailName - ? `Configure AI evaluation for ${guardrailName}` - : "Configure AI evaluation for re-running on logs"} + ? `${term("Configure AI evaluation for")} ${guardrailName}` + : term("Configure AI evaluation for re-running on logs")} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx index 613fee04f68..ee214b451fa 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx @@ -1,9 +1,16 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; -import { render, screen, within } from "@testing-library/react"; +import { fireEvent, render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import { TerminologyProvider } from "@/contexts/TerminologyContext"; import GuardrailGarden from "./guardrail_garden"; import { ALL_CARDS } from "./guardrail_garden_data"; +const uiConfig = vi.hoisted(() => ({ data: undefined as { use_si_terminology?: boolean } | undefined })); + +vi.mock("@/app/(dashboard)/hooks/uiConfig/useUIConfig", () => ({ + useUIConfig: () => ({ data: uiConfig.data }), +})); + vi.mock("./guardrail_garden_detail", () => ({ __esModule: true, default: ({ card, onBack }: { card: { name: string }; onBack: () => void }) => ( @@ -21,6 +28,7 @@ const DECISION_MODEL_CARDS = ALL_CARDS.filter((c) => c.category === "decision_mo describe("GuardrailGarden", () => { beforeEach(() => { vi.clearAllMocks(); + uiConfig.data = undefined; }); const renderGarden = () => render(); @@ -114,6 +122,26 @@ describe("GuardrailGarden", () => { } }); + it.each([ + [true, "SI"], + [false, "AI"], + ])("should match the wording a card shows when SI terminology is %s", (useSi, shownTerm) => { + uiConfig.data = { use_si_terminology: useSi }; + const card = ALL_CARDS.find((c) => c.description.startsWith("AI security platform protecting")); + expect(card).toBeDefined(); + render( + + + , + ); + + fireEvent.change(screen.getByPlaceholderText("Search guardrails"), { + target: { value: `${shownTerm} security platform protecting` }, + }); + + expect(screen.getByText(card!.name)).toBeInTheDocument(); + }); + it("should show an empty result set for a query that matches nothing", async () => { const user = userEvent.setup(); renderGarden(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx index f8e1ab1ef66..d430a5a9f19 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState } from "react"; import { ArrowRight, Search } from "lucide-react"; import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; @@ -11,6 +12,7 @@ interface GuardrailGardenProps { } const GuardrailGarden: React.FC = ({ accessToken, onGuardrailCreated }) => { + const term = useTerminology(); const [searchQuery, setSearchQuery] = useState(""); const [selectedCard, setSelectedCard] = useState(null); const [showAllLitellm, setShowAllLitellm] = useState(false); @@ -22,7 +24,7 @@ const GuardrailGarden: React.FC = ({ accessToken, onGuardr const q = searchQuery.toLowerCase(); return ( card.name.toLowerCase().includes(q) || - card.description.toLowerCase().includes(q) || + term(card.description).toLowerCase().includes(q) || card.tags.some((t) => t.toLowerCase().includes(q)) ); }); @@ -106,7 +108,7 @@ const GuardrailGarden: React.FC = ({ accessToken, onGuardr Partner Guardrails

- Third-party guardrail integrations from leading AI security providers. + {term("Third-party guardrail integrations from leading AI security providers.")}

{partnerCards.map((card) => ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx index 2b7cedd6d85..8a6375d0d68 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx @@ -1,9 +1,11 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React from "react"; import { CircleCheck } from "lucide-react"; import { GuardrailCardInfo } from "./guardrail_garden_data"; import { Logo } from "@/components/molecules/logo/Logo"; const GuardrailCard: React.FC<{ card: GuardrailCardInfo; onClick: () => void }> = ({ card, onClick }) => { + const term = useTerminology(); return (
void }> {card.name}
-

{card.description}

+

{term(card.description)}

{card.eval && (
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx index 8ac0493b012..e549ebca50a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState } from "react"; import { ArrowLeft } from "lucide-react"; import { Button } from "@/components/ui/button"; @@ -21,6 +22,7 @@ interface GuardrailDetailViewProps { } const GuardrailDetailView: React.FC = ({ card, onBack, accessToken, onGuardrailCreated }) => { + const term = useTerminology(); const [isAddFormVisible, setIsAddFormVisible] = useState(false); const [activeTab, setActiveTab] = useState("overview"); @@ -66,7 +68,7 @@ const GuardrailDetailView: React.FC = ({ card, onBack,

{card.name}

-

{card.description}

+

{term(card.description)}

{/* Action buttons — outlined style like Vertex */}
@@ -101,7 +103,7 @@ const GuardrailDetailView: React.FC = ({ card, onBack, {/* Left column — overview + details table */}

Overview

-

{card.description}

+

{term(card.description)}

Guardrail Details

Details are as follows

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx index 04a095f8792..46c1d2c6848 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx @@ -1,5 +1,6 @@ "use client"; +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState } from "react"; import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { RefreshCw, Unplug } from "lucide-react"; @@ -132,6 +133,7 @@ function SessionsBody({ isLoading: boolean; onDisconnect: ((selector: MCPGatewaySessionSelector) => void) | null; }) { + const term = useTerminology(); if (isLoading) { return (

- No live MCP connections on this worker (pid {data.worker_pid}). Connect an AI client to the gateway to see it - here. + No live MCP connections on this worker (pid {data.worker_pid}).{" "} + {term("Connect an AI client to the gateway to see it here.")}

); @@ -166,11 +168,11 @@ function SessionsBody({ <>
- +
- + (null); const queryOptions = { @@ -282,9 +285,9 @@ export function MCPGatewaySessionsTab({ accessToken, canTerminate }: MCPGatewayS

Live Connections

- Stateful Streamable HTTP sessions currently open on this proxy worker, grouped by the AI client that sent - the MCP initialize request and by the authenticated LiteLLM user. Stateless requests and SSE connections are - not counted. + {term( + "Stateful Streamable HTTP sessions currently open on this proxy worker, grouped by the AI client that sent the MCP initialize request and by the authenticated LiteLLM user. Stateless requests and SSE connections are not counted.", + )}

- Use tools directly from Cursor IDE with LiteLLM MCP. Enable your AI assistant to perform real-world tasks - without leaving your coding environment. + {term( + "Use tools directly from Cursor IDE with LiteLLM MCP. Enable your AI assistant to perform real-world tasks without leaving your coding environment.", + )}
@@ -475,8 +478,9 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = []

Connect to your MCP client

- Use tools directly from any MCP client with LiteLLM MCP. Enable your AI assistant to perform real-world - tasks through a simple, secure connection. + {term( + "Use tools directly from any MCP client with LiteLLM MCP. Enable your AI assistant to perform real-world tasks through a simple, secure connection.", + )}

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx index 8acc54bee8d..01151203e65 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx @@ -1,5 +1,6 @@ "use client"; +import { useTerminology } from "@/contexts/TerminologyContext"; import { CircleX, Mic, MicOff, Send, Volume2 } from "lucide-react"; import React, { useCallback, useEffect, useRef, useState } from "react"; import { Button } from "@/components/ui/button"; @@ -27,6 +28,7 @@ const RealtimePlayground: React.FC = ({ customProxyBaseUrl, selectedGuardrails, }) => { + const term = useTerminology(); const [messages, setMessages] = useState([]); const [inputText, setInputText] = useState(""); const [isConnected, setIsConnected] = useState(false); @@ -408,8 +410,8 @@ const RealtimePlayground: React.FC = ({ Realtime Voice Playground

- Click Connect to start a realtime session. You can speak using your microphone or type messages. - The AI will respond with voice and text. + Click Connect to start a realtime session. You can speak using your microphone or type messages.{" "} + {term("The AI will respond with voice and text.")}

)} @@ -428,7 +430,7 @@ const RealtimePlayground: React.FC = ({ : "bg-muted text-foreground rounded-bl-md" }`} > -
{msg.role === "user" ? "You" : "AI"}
+
{msg.role === "user" ? "You" : term("AI")}
{msg.content}
)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx index 54ea3702cb4..db576ed03a5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx @@ -7,6 +7,7 @@ import { type CompliancePrompt, } from "@/data/compliancePrompts"; import useCan from "@/app/(dashboard)/hooks/useCan"; +import { useTerminology } from "@/contexts/TerminologyContext"; import { getGuardrailsList, testPoliciesAndGuardrails } from "@/components/networking"; import PolicySelector, { getPolicyOptionEntries } from "@/components/policies/PolicySelector"; import { Policy } from "@/components/policies/types"; @@ -131,7 +132,8 @@ export default function ComplianceUI({ proxySettings, }: ComplianceUIProps) { const canViewPolicies = useCan("viewPolicies"); - const frameworks = getFrameworks(); + const term = useTerminology(); + const frameworks = getFrameworks(term); const [policyValueToLabel, setPolicyValueToLabel] = useState>(new Map()); const [guardrailOptions, setGuardrailOptions] = useState([]); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx index c4f76b07a4e..50ffb89ae83 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useEffect, useMemo, useState } from "react"; import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; @@ -62,6 +63,7 @@ const AiSuggestionModal: React.FC = ({ accessToken, allTemplates, }) => { + const term = useTerminology(); const [attackExamples, setAttackExamples] = useState([""]); const [description, setDescription] = useState(""); const [isLoading, setIsLoading] = useState(false); @@ -815,7 +817,7 @@ const AiSuggestionModal: React.FC = ({ {/* Header */}
- AI Policy Suggestion + {term("AI Policy Suggestion")}

{showResults ? `${suggestions?.length || 0} template${(suggestions?.length || 0) !== 1 ? "s" : ""} matched your requirements` diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx index c4a45d88d1f..25a4b89ed2a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState, useEffect } from "react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -40,6 +41,7 @@ const GuardrailSelectionModal: React.FC = ({ isLoading = false, progressInfo, }) => { + const term = useTerminology(); const [selectedGuardrails, setSelectedGuardrails] = useState>(new Set()); // Prepare guardrail info with existence status @@ -199,7 +201,7 @@ const GuardrailSelectionModal: React.FC = ({

✨ - AI-Discovered Competitors ({template.discoveredCompetitors.length}) + {term("AI-Discovered Competitors")} ({template.discoveredCompetitors.length})
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_templates.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_templates.tsx index 46dda1dd11c..0db1f74f540 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_templates.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_templates.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState, useEffect, useMemo } from "react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -106,6 +107,7 @@ const PolicyTemplates: React.FC = ({ onTemplatesLoaded, accessToken, }) => { + const term = useTerminology(); const [templates, setTemplates] = useState([]); const [isLoading, setIsLoading] = useState(false); const [selectedTags, setSelectedTags] = useState>(new Set()); @@ -191,7 +193,7 @@ const PolicyTemplates: React.FC = ({ - Use AI to find templates + {term("Use AI to find templates")}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx index b90379b149d..6aa26088756 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx @@ -1,3 +1,4 @@ +import { useTerminology } from "@/contexts/TerminologyContext"; import React, { useState, useEffect } from "react"; import { Button } from "@/components/ui/button"; import { Badge } from "@/components/ui/badge"; @@ -44,6 +45,7 @@ const TemplateParameterModal: React.FC = ({ isLoading = false, accessToken, }) => { + const term = useTerminology(); const [parameterValues, setParameterValues] = useState>({}); const [competitorMode, setCompetitorMode] = useState<"ai" | "manual">("ai"); const [selectedModel, setSelectedModel] = useState(null); @@ -260,7 +262,7 @@ const TemplateParameterModal: React.FC = ({ className="grid-cols-2" >