From 764a7a6536ee804af2399c53ec6715399fb34979 Mon Sep 17 00:00:00 2001 From: Alon Sokolovski Date: Thu, 29 Oct 2026 06:35:32 +0200 Subject: [PATCH 01/42] fix(context_management)!: count the system prompt toward the clear_tool_uses input_tokens trigger (#47857) * fix(context_management): count the system prompt toward the clear_tool_uses input_tokens trigger Anthropic's API compares an input_tokens trigger with the whole prompt, system included; the polyfill counted messages and tools only, so it cleared late by the system prompt's size. compact_20260112 already counts it: share its system_to_text. * fix(context_management): take read-only types in the system prompt count The type-discipline gate (LIT001) rejects new list and dict annotations. * chore: rerun CI (proxy-build-from-pip timed out) * fix(context_management): count the system prompt through compact's shared counter Finishes the main merge's conflict in clear_tool_uses.py. The input_tokens trigger now adds the system prompt through count_system_tokens, the same counter compact uses, while the count reused for cleared_input_tokens stays messages and tools only, so the reported savings are just the cleared tool results * perf(context_management): skip the token count when nothing can be cleared The keep check runs before the trigger count, the system prompt reader accepts any mapping block, and the system prompt test gets short parameter ids * test(context_management): keep the off-loop count test counting after the keep check moved first * test(context_management): cover a client compaction summary counting toward the trigger The trigger also skips the system prompt count when messages and tools alone already pass it * fix(context_management): move a client compaction summary into the system prompt when clearing runs before compact * fix(context_management): clear only the history after a client compaction block so compact still replaces its summary * fix(context_management): keep tool calls beside a client compaction block clearable --------- Co-authored-by: AlonSoko <25619863+AlonSoko@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../editors/clear_tool_uses.py | 40 ++++--- .../context_management/editors/compact.py | 27 ++++- .../test_clear_tool_uses.py | 111 +++++++++++++++++- .../context_management/test_dispatcher.py | 1 + 4 files changed, 158 insertions(+), 21 deletions(-) diff --git a/litellm/llms/anthropic/messages/context_management/editors/clear_tool_uses.py b/litellm/llms/anthropic/messages/context_management/editors/clear_tool_uses.py index d4430ca5ab9..c4f7966f429 100644 --- a/litellm/llms/anthropic/messages/context_management/editors/clear_tool_uses.py +++ b/litellm/llms/anthropic/messages/context_management/editors/clear_tool_uses.py @@ -31,6 +31,7 @@ from ..constants import ( ) from ..errors import AnthropicContextManagementError from ..placeholders import build_cleared_tool_result_content +from .compact import count_system_tokens, history_after_client_compaction if TYPE_CHECKING: from litellm.types.llms.openai import ChatCompletionToolParam @@ -203,9 +204,10 @@ def _trigger_met( model: str, messages: list[dict[str, object]], tools: list[dict[str, object]] | None, + system: str | Sequence[Mapping[str, object]] | None, counted_tool_uses: int, ) -> tuple[bool, int | None]: - """Return (trigger_met, input_tokens if counted for reuse).""" + """Return (trigger_met, message and tool tokens for reuse); like the API, input_tokens adds the system prompt.""" trigger_type: Final = trigger.get("type", "input_tokens") threshold = trigger.get("value") @@ -216,14 +218,17 @@ def _trigger_met( if not isinstance(threshold, int): threshold = DEFAULT_INPUT_TOKENS_TRIGGER - current_tokens: Final = litellm.token_counter( + message_tokens: Final = litellm.token_counter( model=model, messages=messages, tools=cast("list[ChatCompletionToolParam] | None", tools), ) + if message_tokens > threshold: + return True, message_tokens + current_tokens: Final = message_tokens + count_system_tokens(model, system) verbose_logger.debug("context_management polyfill: current_tokens: %s", current_tokens) verbose_logger.debug("context_management polyfill: threshold: %s", threshold) - return current_tokens > threshold, current_tokens + return current_tokens > threshold, message_tokens def _resolve_keep_count(keep: dict[str, object]) -> int: @@ -237,7 +242,7 @@ def _resolve_keep_count(keep: dict[str, object]) -> int: def _last_completed_tool_use_id( - messages: list[dict[str, object]], + messages: Sequence[Mapping[str, object]], ) -> str | None: """Latest completed tool_result id; never cleared.""" last_id: str | None = None @@ -253,7 +258,7 @@ def _last_completed_tool_use_id( def _clear_tool_results( - messages: list[dict[str, object]], ids_to_clear: set[str] + messages: list[dict[str, object]], ids_to_clear: Set[str] ) -> tuple[list[dict[str, object]], int]: """Clear matching tool_result content; return (messages, cleared_count).""" cleared = 0 @@ -311,7 +316,11 @@ def apply_clear_tool_uses_20250919( system: str | list[dict[str, object]] | None, edit_spec: ClearToolUsesEditSpec, ) -> tuple[list[dict[str, object]], AppliedEdit | None]: - """Apply clear_tool_uses to an edit ``check_clear_tool_uses_spec`` returned; return (messages, AppliedEdit or None).""" + """Apply clear_tool_uses to an edit ``check_clear_tool_uses_spec`` returned; return (messages, AppliedEdit or None). + + Like the API, it reads only the history after a client compaction block and counts the block's summary as system. + """ + kept, counted_system = history_after_client_compaction(messages, system) fields: Final = _ClearToolUsesFields.model_validate(edit_spec) clear_at_least: Final = 0 if fields.clear_at_least is None else fields.clear_at_least.value exclude_tools: Final = frozenset(fields.exclude_tools or ()) @@ -326,7 +335,7 @@ def apply_clear_tool_uses_20250919( "value": DEFAULT_KEEP_TOOL_USES, } - blocks: Final = _all_blocks(messages) + blocks: Final = _all_blocks(kept) results: Final = _blocks_by_id(blocks, "tool_result", "tool_use_id") uses: Final = _blocks_by_id(blocks, "tool_use", "id") tiny: Final = _tiny_result_ids(results) @@ -335,15 +344,15 @@ def apply_clear_tool_uses_20250919( for tool_use_id, name in _tool_uses_in_order(blocks) if name not in exclude_tools and (tool_use_id not in tiny or _clears_input(name, input_tools)) ) - met, counted_tokens = _trigger_met(trigger, model, messages, tools, len(clearable)) - if not met: - return messages, None - keep_count: Final = _resolve_keep_count(keep) if len(clearable) <= keep_count: return messages, None - last_completed_id: Final = _last_completed_tool_use_id(messages) + met, counted_tokens = _trigger_met(trigger, model, kept, tools, counted_system, len(clearable)) + if not met: + return messages, None + + last_completed_id: Final = _last_completed_tool_use_id(kept) candidates: Final = tuple( (tool_use_id, name) for tool_use_id, name in clearable[: len(clearable) - keep_count] @@ -358,7 +367,7 @@ def apply_clear_tool_uses_20250919( _select_in_batches(candidate_ids, tokens, clear_at_least) if clear_at_least > 0 else candidate_ids ) - edited, cleared_count = _clear_tool_results(messages, set(ids_to_clear)) + edited, cleared_count = _clear_tool_results(messages, ids_to_clear) verbose_logger.debug("context_management polyfill: edited: %s", edited) if cleared_count == 0: return messages, None @@ -370,11 +379,12 @@ def apply_clear_tool_uses_20250919( counted_tokens if counted_tokens is not None else litellm.token_counter( - model=model, messages=messages, tools=cast("list[ChatCompletionToolParam] | None", tools) + model=model, messages=kept, tools=cast("list[ChatCompletionToolParam] | None", tools) ) ) + kept_after, _ = history_after_client_compaction(final_messages, system) tokens_after: Final = litellm.token_counter( - model=model, messages=final_messages, tools=cast("list[ChatCompletionToolParam] | None", tools) + model=model, messages=kept_after, tools=cast("list[ChatCompletionToolParam] | None", tools) ) cleared_input_tokens: Final = max(tokens_before - tokens_after, 0) diff --git a/litellm/llms/anthropic/messages/context_management/editors/compact.py b/litellm/llms/anthropic/messages/context_management/editors/compact.py index 254d1c4c058..0bf8cac7344 100644 --- a/litellm/llms/anthropic/messages/context_management/editors/compact.py +++ b/litellm/llms/anthropic/messages/context_management/editors/compact.py @@ -189,6 +189,21 @@ def _strip_compaction_blocks( return cleaned +def history_after_client_compaction( + messages: Sequence[dict[str, object]], + system: str | Sequence[Mapping[str, object]] | None, +) -> tuple[list[dict[str, object]], str | Sequence[Mapping[str, object]] | None]: + """Return the history the API keeps after the latest compaction block, without it, and the system with its summary.""" + sliced, block = _slice_around_compaction_block(messages) + if block is None: + return list(messages), system + tail: Final = _strip_compaction_blocks(sliced) + summary: Final = block.get("content") + if not isinstance(summary, str) or not summary: + return tail, system + return tail, _augment_system_with_summary(_system_to_text(system) or None, summary) + + def _augment_system_with_summary( system: str | list[dict[str, object]] | None, summary_text: str, @@ -320,14 +335,16 @@ def _count_effective_tokens( content: Final = compaction_block.get("content") or "" if content: total += litellm.token_counter(model=model, text=content) + return total + count_system_tokens(model, system) + + +def count_system_tokens(model: str, system: str | Sequence[Mapping[str, object]] | None) -> int: system_text: Final = _system_to_text(system) - if system_text: - total += litellm.token_counter(model=model, text=system_text) - return total + return litellm.token_counter(model=model, text=system_text) if system_text else 0 def _system_to_text( - system: str | list[dict[str, object]] | None, + system: str | Sequence[Mapping[str, object]] | None, ) -> str: """Flatten an Anthropic-style ``system`` value into a single string for token counting. Returns ``""`` when ``system`` carries no text.""" @@ -338,7 +355,7 @@ def _system_to_text( return "\n".join( text for block in system - if isinstance(block, dict) + if isinstance(block, Mapping) and block.get("type") == "text" and isinstance(text := block.get("text"), str) and text diff --git a/tests/unit/llms/anthropic/messages/context_management/test_clear_tool_uses.py b/tests/unit/llms/anthropic/messages/context_management/test_clear_tool_uses.py index 2567d1b76db..f03ffb25b2b 100644 --- a/tests/unit/llms/anthropic/messages/context_management/test_clear_tool_uses.py +++ b/tests/unit/llms/anthropic/messages/context_management/test_clear_tool_uses.py @@ -10,7 +10,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest import litellm -from litellm.llms.anthropic.messages.adapters.handler import _run_polyfill_if_enabled +from litellm.llms.anthropic.messages.adapters.handler import _run_polyfill_if_enabled, prepare_context_managed_request from litellm.llms.anthropic.messages.context_management import ( AnthropicContextManagementError, PolyfillResult, @@ -690,6 +690,115 @@ def test_tool_uses_trigger_does_not_count_excluded_tools(trigger_value: int, tri assert (applied is not None) == triggers +@pytest.mark.parametrize( + "system", + ["Reference notes. " * 400, [{"type": "text", "text": "Reference notes. " * 400}]], + ids=["string", "blocks"], +) +def test_input_tokens_trigger_counts_the_system_prompt(system: str | list[dict[str, object]]) -> None: + """Like the API, the trigger compares the whole prompt: a system prompt can push it over.""" + messages = _make_history(n_pairs=3) + history_tokens: Final = litellm.token_counter(model=MODEL, messages=messages) + edit_spec: Final = { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": history_tokens + 100}, + "keep": {"type": "tool_uses", "value": 1}, + } + _, without_system = apply_clear_tool_uses_20250919( + model=MODEL, messages=deepcopy(messages), tools=None, system=None, edit_spec=edit_spec + ) + _, with_system = apply_clear_tool_uses_20250919( + model=MODEL, messages=deepcopy(messages), tools=None, system=system, edit_spec=edit_spec + ) + _, met_without_system = apply_clear_tool_uses_20250919( + model=MODEL, + messages=deepcopy(messages), + tools=None, + system=None, + edit_spec={**edit_spec, "trigger": {"type": "input_tokens", "value": 0}}, + ) + assert without_system is None + assert with_system is not None + assert met_without_system is not None + assert with_system["cleared_tool_uses"] == 2 + assert with_system["cleared_input_tokens"] == met_without_system["cleared_input_tokens"] > 0 + + +@pytest.mark.parametrize("compact_position", [None, "after", "before"], ids=["alone", "before_compact", "after_compact"]) +async def test_input_tokens_trigger_counts_a_client_compaction_summary(compact_position: str | None) -> None: + """A client compaction summary counts toward the trigger and calls beside it stay clearable, in any edit order.""" + tail: Final = _make_history(n_pairs=3) + edit_spec: Final = { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": litellm.token_counter(model=MODEL, messages=tail) + 100}, + "keep": {"type": "tool_uses", "value": 1}, + } + compaction: Final = {"type": "compaction", "content": "Earlier work. " * 400} + compacted_history: Final = [tail[0], {**tail[1], "content": [compaction, *tail[1]["content"]]}, *tail[2:]] + compact_spec: Final = {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50_000}} + edits: Final = {None: [edit_spec], "after": [edit_spec, compact_spec], "before": [compact_spec, edit_spec]}[ + compact_position + ] + + async def prepare(messages: list[dict[str, object]]) -> PolyfillResult | None: + return await prepare_context_managed_request( + model=MODEL, + messages=messages, + tools=None, + system=None, + context_management_spec={"edits": edits}, + litellm_metadata=None, + additional_drop_params=None, + llm_router=None, + ) + + def cleared(result: PolyfillResult | None) -> list[object]: + applied: Final = result.applied_edits if result is not None else [] + return [edit.get("cleared_tool_uses") for edit in applied if edit.get("type") == "clear_tool_uses_20250919"] + + without_summary: Final = await prepare(deepcopy(tail)) + with_summary: Final = await prepare(deepcopy(compacted_history)) + + assert cleared(without_summary) == [] + assert cleared(with_summary) == [2] + + +async def test_compact_after_clearing_replaces_the_client_compaction_summary(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.proxy_server import general_settings + + monkeypatch.setitem(general_settings, COMPACT_SUMMARY_MODEL_SETTING_KEY, "summary-model") + summary: Final = MagicMock() + summary.choices = [MagicMock(message=MagicMock(content="Fresh summary"))] + summary.usage = MagicMock(prompt_tokens=50, completion_tokens=100) + router: Final = MagicMock() + router.acompletion = AsyncMock(return_value=summary) + tail: Final = _make_history(n_pairs=3) + result: Final = await prepare_context_managed_request( + model=MODEL, + messages=[{"role": "assistant", "content": [{"type": "compaction", "content": "Old summary. " * 20_000}]}, *tail], + tools=None, + system="You are a helpful assistant.", + context_management_spec={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 1_000}, + "keep": {"type": "tool_uses", "value": 1}, + }, + {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50_000}}, + ] + }, + litellm_metadata=None, + additional_drop_params=None, + llm_router=router, + ) + assert result is not None + assert isinstance(result.system, str) + assert "Fresh summary" in result.system + assert "Old summary" not in result.system + assert result.system.endswith("You are a helpful assistant.") + + @pytest.mark.parametrize(("clear_tool_inputs", "triggers"), [(False, False), (True, True)]) def test_tool_uses_trigger_counts_tiny_results_only_when_their_inputs_are_cleared( clear_tool_inputs: bool, triggers: bool diff --git a/tests/unit/llms/anthropic/messages/context_management/test_dispatcher.py b/tests/unit/llms/anthropic/messages/context_management/test_dispatcher.py index b79e33b0c3b..fabb05c2e61 100644 --- a/tests/unit/llms/anthropic/messages/context_management/test_dispatcher.py +++ b/tests/unit/llms/anthropic/messages/context_management/test_dispatcher.py @@ -151,6 +151,7 @@ async def test_sync_editor_counts_tokens_off_the_event_loop(): { "type": "clear_tool_uses_20250919", "trigger": {"type": "input_tokens", "value": 10_000_000}, + "keep": {"type": "tool_uses", "value": 0}, } ] }, From a11e76fe8f861b91664935ec8a997fdc2cfd8a01 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:46:06 -0700 Subject: [PATCH 02/42] fix(google_genai): read one-object and null spellings of repeated fields in countTokens and generateContent (#48771) * fix(count_tokens): normalize Google system parts and null tools * fix(count_tokens): preserve typed validator return * fix(count_tokens): validate single system part without typing.cast * fix(count_tokens): accept every Google spelling of a repeated field Google countTokens now reads a repeated field sent as one object, null, or a list with null items the same way as its list form, for system parts, tools, and function declarations, matching what Vertex accepts. * fix(count_tokens): read every Google spelling of contents and skip null system text contents and each content's parts now read one object, null, and null items like their list form, the same as system parts and tools. A one-object parts used to translate to no messages, so non-Google models counted it as empty. A system part whose text is null is skipped, as Vertex AI does. * fix(google_genai): read null-text media system parts as media and one-object content parts on generateContent * fix(google_proxy): count null-text system parts that carry only part metadata, like Vertex AI A system part with null or missing text and only metadata such as thought or thoughtSignature, or an empty part, now counts as empty instead of answering 400. Parts that carry an image or file with null text still answer 400. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../google_genai/adapters/transformation.py | 7 +- litellm/proxy/google_endpoints/endpoints.py | 73 +++++++++++++-- .../google_genai/test_google_genai_adapter.py | 12 +++ .../proxy/google_endpoints/test_endpoints.py | 90 ++++++++++++++++++- 4 files changed, 172 insertions(+), 10 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index e1085dd7857..c1895376675 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -198,6 +198,11 @@ def _genai_message_part(part: object) -> _GenAIMessagePart | None: return ChatCompletionTextObject(type="text", text=text) if isinstance(text, str) else _genai_media_part(fields) +def _genai_content_parts(value: object) -> Sequence[object]: + single_part: Final = _validated(_CONFIG_FIELDS, value) + return (single_part,) if single_part is not None else _validated(_GENAI_CONTENT_PARTS, value) or () + + def _genai_message_parts(parts: object) -> tuple[_GenAIMessagePart, ...]: items: Final = _validated(_GENAI_CONTENT_PARTS, parts) or () return tuple(converted for part in items if (converted := _genai_message_part(part)) is not None) @@ -248,7 +253,7 @@ def _genai_content_messages(content: object) -> tuple[AllMessageValues, ...]: fields: Final = _validated(_CONFIG_FIELDS, content) if fields is None: return () - parts: Final = _validated(_GENAI_CONTENT_PARTS, fields.get("parts", ())) or () + parts: Final = _genai_content_parts(fields.get("parts", ())) converted: Final = _genai_message_parts(parts) role: Final = fields.get("role", "user") if role == "user": diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 54976857584..c018f071b3c 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,11 +1,11 @@ from collections.abc import Mapping from itertools import chain from types import MappingProxyType -from typing import Final +from typing import Annotated, Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, ORJSONResponse -from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * @@ -27,14 +27,67 @@ from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.types.llms.vertex_ai import TokenCountDetailsResponse +def _google_repeated_field(value: JsonValue) -> JsonValue: + match value: + case None: + return [] + case dict(): + return [value] + case list(): + return [item for item in value if item is not None] + case _: + return value + + +def _google_content(value: JsonValue) -> JsonValue: + match value: + case {"parts": parts}: + return {**value, "parts": _google_repeated_field(parts)} + case _: + return value + + +def _google_contents(value: JsonValue) -> JsonValue: + match _google_repeated_field(value): + case list() as contents: + return [_google_content(content) for content in contents] + case other: + return other + + +_GOOGLE_PART_TEXT_AND_METADATA_FIELDS: Final = frozenset( + {"text", "thought", "thoughtSignature", "thought_signature", "partMetadata", "part_metadata"} +) + + +def _google_part_has_non_text_content(part: Mapping[str, JsonValue]) -> bool: + return any(field is not None for key, field in part.items() if key not in _GOOGLE_PART_TEXT_AND_METADATA_FIELDS) + + +def _google_system_part(value: JsonValue) -> JsonValue: + match value: + case dict() if value.get("text") is None and _google_part_has_non_text_content(value): + return {key: field for key, field in value.items() if key != "text"} + case dict(): + return {"text": None, **value} + case _: + return value + + +_GOOGLE_REPEATED_FIELD: Final = BeforeValidator(_google_repeated_field) + + class _GoogleCountTokensTextPart(BaseModel): model_config = ConfigDict(frozen=True) - text: str + text: str | None class _GoogleCountTokensSystemInstruction(BaseModel): model_config = ConfigDict(frozen=True) - parts: tuple[_GoogleCountTokensTextPart, ...] + parts: Annotated[ + tuple[Annotated[_GoogleCountTokensTextPart, BeforeValidator(_google_system_part)], ...], + _GOOGLE_REPEATED_FIELD, + ] = () class _GoogleCountTokensFunctionDeclaration(BaseModel): @@ -47,16 +100,16 @@ class _GoogleCountTokensFunctionDeclaration(BaseModel): class _GoogleCountTokensTool(BaseModel): model_config = ConfigDict(frozen=True, populate_by_name=True) - function_declarations: tuple[_GoogleCountTokensFunctionDeclaration, ...] = Field( + function_declarations: Annotated[tuple[_GoogleCountTokensFunctionDeclaration, ...], _GOOGLE_REPEATED_FIELD] = Field( default=(), alias="functionDeclarations" ) class _GoogleCountTokensRequest(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True) - contents: tuple[Mapping[str, JsonValue], ...] = () + contents: Annotated[tuple[Mapping[str, JsonValue], ...], BeforeValidator(_google_contents)] = () system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction") - tools: tuple[_GoogleCountTokensTool, ...] = () + tools: Annotated[tuple[_GoogleCountTokensTool, ...], _GOOGLE_REPEATED_FIELD] = () router: Final = APIRouter( @@ -112,7 +165,11 @@ def _google_count_tokens_request(model_name: str, data: object) -> TokenCountReq GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", []) ) system: Final = ( - [{"type": "text", "text": part.text} for part in validated_data.system_instruction.parts] + [ + {"type": "text", "text": part.text} + for part in validated_data.system_instruction.parts + if part.text is not None + ] if validated_data.system_instruction is not None else None ) diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 0292e58c029..4ab642dc99d 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -158,6 +158,18 @@ def test_native_function_response_preserves_json_when_media_is_malformed(invalid assert messages == [{"role": "tool", "tool_call_id": "call_lookup", "content": '{"ok": true}'}] +@pytest.mark.parametrize("role", ["user", "model"]) +def test_one_object_content_parts_translate_like_their_list_form(role: str) -> None: + part: Final = {"text": "Reply with OK"} + + def messages(parts: object) -> object: + return GoogleGenAIAdapter().translate_generate_content_to_completion( + model="anthropic/claude-haiku-5-5", contents=[{"role": role, "parts": parts}] + )["messages"] + + assert messages(part) == messages([part]) != [] + + @pytest.mark.parametrize( ("args", "expected_arguments"), [(None, "null"), ('{"query":"x"}', '"{\\"query\\":\\"x\\"}"')], diff --git a/tests/unit/proxy/google_endpoints/test_endpoints.py b/tests/unit/proxy/google_endpoints/test_endpoints.py index 80e45edf897..9dd2e267e64 100644 --- a/tests/unit/proxy/google_endpoints/test_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_endpoints.py @@ -6,7 +6,7 @@ from typing import Final import pytest from dotenv import load_dotenv -from pydantic import JsonValue +from pydantic import JsonValue, ValidationError from starlette.requests import Request from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens @@ -105,3 +105,91 @@ async def test_proxy_gemini_to_openai_like_model_token_counting(): ) assert response.get("totalTokens") > 0 + + +_DECLARATION: Final = {"name": "lookup", "description": "Find a record"} +_SYSTEM: Final = {"systemInstruction": {"parts": [{"text": "Follow these rules"}]}} +_NO_SYSTEM_PARTS: Final = {"systemInstruction": {"parts": []}} +_TOOLS: Final = {"tools": [{"functionDeclarations": [_DECLARATION]}]} +_TURN: Final = {"role": "user", "parts": [{"text": "Count this request"}]} +_CONTENTS: Final = {"contents": [_TURN]} + + +@pytest.mark.parametrize( + ("spelled", "canonical"), + ( + ({"systemInstruction": {"parts": {"text": "Follow these rules"}}}, _SYSTEM), + ({"system_instruction": {"parts": {"text": "Follow these rules"}}}, _SYSTEM), + ({"systemInstruction": {"parts": [{"text": "Follow these rules"}, None]}}, _SYSTEM), + ({"systemInstruction": {"parts": None}}, _NO_SYSTEM_PARTS), + ({"systemInstruction": {"parts": [{"text": None}, {"text": "Follow these rules"}]}}, _SYSTEM), + ({"systemInstruction": {"parts": {"text": None}}}, _NO_SYSTEM_PARTS), + ({"systemInstruction": {"parts": {"text": None, "inlineData": None}}}, _NO_SYSTEM_PARTS), + ({"systemInstruction": {"parts": {"text": None, "thought": False}}}, _NO_SYSTEM_PARTS), + ({"systemInstruction": {"parts": [{"thoughtSignature": "c2ln"}, {"text": "Follow these rules"}]}}, _SYSTEM), + ({"systemInstruction": {"parts": [{}]}}, _NO_SYSTEM_PARTS), + ({"systemInstruction": {}}, _NO_SYSTEM_PARTS), + ({"tools": None}, {}), + ({"tools": [None]}, {}), + ({"tools": [{"functionDeclarations": None}]}, {}), + ({"tools": {"functionDeclarations": [_DECLARATION]}}, _TOOLS), + ({"tools": [{"functionDeclarations": _DECLARATION}]}, _TOOLS), + ({"tools": {"functionDeclarations": _DECLARATION}}, _TOOLS), + ({"contents": _TURN}, _CONTENTS), + ({"contents": [None, _TURN]}, _CONTENTS), + ({"contents": None}, {"contents": []}), + ({"contents": [{"role": "user", "parts": {"text": "Count this request"}}]}, _CONTENTS), + ({"contents": [{"role": "user", "parts": [None, {"text": "Count this request"}]}]}, _CONTENTS), + ({"contents": [{"role": "user", "parts": None}]}, {"contents": [{"role": "user", "parts": []}]}), + ), +) +def test_google_count_tokens_counts_every_google_spelling_of_a_repeated_field_like_its_list_form( + spelled: dict[str, JsonValue], canonical: dict[str, JsonValue] +) -> None: + contents: Final = [{"role": "user", "parts": [{"text": "Count this request"}]}] + + assert _google_count_tokens_request("counted", {"contents": contents, **spelled}) == _google_count_tokens_request( + "counted", {"contents": contents, **canonical} + ) + + +@pytest.mark.parametrize( + "malformed", + ( + {"tools": "lookup"}, + {"tools": [{"functionDeclarations": "lookup"}]}, + {"systemInstruction": {"parts": 7}}, + {"contents": 5}, + {"contents": ["Count this request"]}, + ), +) +def test_google_count_tokens_rejects_a_scalar_where_google_expects_an_object_or_list( + malformed: dict[str, JsonValue], +) -> None: + with pytest.raises(ValidationError): + _google_count_tokens_request("counted", {"contents": [_TURN], **malformed}) + + +@pytest.mark.parametrize( + "parts", + ( + {"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}, + [{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}], + {"text": None, "inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}, + [{"text": "Follow these rules"}, {"text": None, "fileData": {"fileUri": "gs://bucket/rules.pdf"}}], + ), +) +def test_google_count_tokens_rejects_media_only_system_parts(parts: JsonValue) -> None: + with pytest.raises(ValidationError): + _google_count_tokens_request( + "counted", + {"contents": [{"role": "user", "parts": [{"text": "Hello"}]}], "systemInstruction": {"parts": parts}}, + ) + + +def test_google_count_tokens_keeps_counting_system_text_that_sits_beside_media() -> None: + text_with_media: Final = {"text": "Follow these rules", "inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}} + + assert _google_count_tokens_request( + "counted", {**_CONTENTS, "systemInstruction": {"parts": [text_with_media]}} + ) == _google_count_tokens_request("counted", {**_CONTENTS, **_SYSTEM}) From cbf4e40d381042e454209f8c54cfb3e76062ceb4 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:49:03 -0700 Subject: [PATCH 03/42] fix(proxy)!: stop labelling management and health rejections as auth failures (#48739) * test(proxy): cover management and health routes labelling rejections as auth failures * fix(proxy): name management and health route errors by their status instead of auth_error * fix(proxy): keep auth_error for 401 and 403 and answer database outages with 503 in route catch-alls /key/list and /key/aliases now share the catch-all helper too, so they stop adding the Authentication Error prefix and typing a 400 as internal_server_error * fix(proxy): answer 503 when the database drops mid-query and prisma fails to name the error A real outage surfaces from prisma's query engine as a builtin AttributeError, which the prisma-module check skipped, so the route catch-alls still answered 500 or 400 * fix(proxy): report wrapped and stalled database outages as 503 and label upstream failures by their own status * fix(proxy): answer the upstream status on the catch-all, keep auth_error on tag denials, report model insert outages as 503 * fix(proxy): answer 503 for database failovers and label provider auth errors like chat A database that refuses connections while starting up, shutting down, in recovery, or read-only after a failover now answers 503 no_db_connection instead of 400. A provider's 401 or 403 keeps the type /chat/completions gives it, while a 401 or 403 the proxy raises keeps auth_error. An exception with no message is named by its type. * fix(proxy): scope failover 503s to management route errors and keep a re-raised error's cause * test(proxy): expect the OpenAI error shape from management route refusals Team-not-found 404s now answer invalid_request_error, and every handler-raised refusal answers param null instead of the string None. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/_types.py | 18 +- .../common_utils/openai_error_payload.py | 96 +++++- litellm/proxy/db/exception_handler.py | 4 +- .../health_endpoints/_health_endpoints.py | 20 +- .../internal_user_endpoints.py | 19 +- .../key_management_endpoints.py | 55 +--- .../model_management_endpoints.py | 69 ++-- .../organization_endpoints.py | 19 +- .../management_endpoints/team_endpoints.py | 19 +- litellm/proxy/proxy_server.py | 59 +--- .../authorization/test_team_admin_gate.py | 2 +- .../test_user_self_update_guard.py | 4 +- .../test_team_delete_large_membership.py | 4 +- .../common_utils/test_openai_error_payload.py | 294 ++++++++++++++++++ .../health_endpoints/test_health_endpoints.py | 9 + .../test_internal_user_endpoints.py | 15 + .../test_key_management_endpoints.py | 72 ++++- .../test_model_management_endpoints.py | 75 ++++- .../test_organization_endpoints.py | 31 ++ .../test_team_endpoints.py | 45 +++ tests/unit/proxy/test_proxy_server.py | 99 ++++++ 21 files changed, 808 insertions(+), 220 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5ce6cd93a22..df139a9b84d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3816,6 +3816,16 @@ class ManagementEndpointLoggingPayload(LiteLLMPydanticObjectBase): end_time: datetime | None = None +def proxy_exception_status_code(message: str, code: str) -> str: + # Litellm router.py returns "No healthy deployment available" when there are no deployments available + # Should map to 429 errors https://github.com/BerriAI/litellm/issues/2487 + if "No healthy deployment available" in message or "No deployments available" in message: + return "429" + if RouterErrors.no_deployments_with_tag_routing.value in message: + return "401" + return code + + class ProxyException(Exception): # NOTE: DO NOT MODIFY THIS # This is used to map exactly to OPENAI Exceptions @@ -3844,13 +3854,7 @@ class ProxyException(Exception): headers[k] = str(v) self.headers = headers or {} self.provider_specific_fields = provider_specific_fields - # rules for proxyExceptions - # Litellm router.py returns "No healthy deployment available" when there are no deployments available - # Should map to 429 errors https://github.com/BerriAI/litellm/issues/2487 - if "No healthy deployment available" in self.message or "No deployments available" in self.message: - self.code = "429" - elif RouterErrors.no_deployments_with_tag_routing.value in self.message: - self.code = "401" + self.code = proxy_exception_status_code(self.message, self.code) def to_dict(self) -> dict: """Converts the ProxyException instance to a dictionary.""" diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index 2f01dbe11f7..3c91e87ee51 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -2,14 +2,17 @@ ``type`` is a required string and ``param`` is nullable, neither of which the literal string ``"None"`` satisfies.""" +import sys from collections.abc import Mapping from types import MappingProxyType from typing import Final -from fastapi import status +from fastapi import HTTPException, status from litellm.constants import STRINGIFIED_NONE -from litellm.proxy._types import ProxyException +from litellm.proxy._types import ProxyErrorTypes, ProxyException, proxy_exception_status_code +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain LITELLM_CALL_ID_HEADER: Final = "x-litellm-call-id" @@ -21,6 +24,19 @@ _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( } ) +_AUTH_FAILURE_STATUS_CODES: Final = frozenset({status.HTTP_401_UNAUTHORIZED, status.HTTP_403_FORBIDDEN}) + +_FAILOVER_SQLSTATES: Final = frozenset({"25006", "57P01", "57P02", "57P03"}) + +_FAILOVER_REFUSALS: Final = ( + "is not currently accepting connections", + "the database system is starting up", + "the database system is shutting down", + "the database system is in recovery mode", + "the database system is not yet accepting connections", + "terminating connection due to administrator command", +) + def attribute_of(value: object, name: str, default: object = None) -> object: return getattr(value, name, default) @@ -61,6 +77,82 @@ def openai_error_param(exc: object) -> str | None: return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None +def proxy_exception_for(exc: Exception, default_status_code: int) -> ProxyException: + match exc: + case ProxyException(): + return exc + case _ if _is_database_outage(exc): + return ProxyException( + message=PrismaDBExceptionHandler.database_unavailable_message(exc), + type=ProxyErrorTypes.no_db_connection.value, + param=None, + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) + case HTTPException(status_code=status_code, detail=detail): + return _openai_shaped_proxy_exception(exc, str(detail), status_code) + case _: + return _openai_shaped_proxy_exception( + exc, str(exc) or type(exc).__name__, error_status_code(exc, default_status_code) + ) + + +def _is_database_outage(exc: Exception) -> bool: + return any(map(_is_database_outage_link, exception_chain(exc))) + + +def _is_database_outage_link(exc: BaseException) -> bool: + match exc: + case DBLookupDeadlineExceeded(): + return True + case Exception() if _prisma_is_loaded(): + return _is_prisma_outage(exc) + case _: + return False + + +def _prisma_is_loaded() -> bool: + return sys.modules.get("prisma") is not None + + +def _is_prisma_outage(exc: Exception) -> bool: + if PrismaDBExceptionHandler.is_prisma_engine_internal_error(exc): + return True + if type(exc).__module__.partition(".")[0] != "prisma": + return False + return PrismaDBExceptionHandler.is_database_service_unavailable_error(exc) or _is_failover_refusal(exc) + + +def _is_failover_refusal(exc: Exception) -> bool: + if PrismaDBExceptionHandler.is_read_only_transaction_error(exc): + return True + if PrismaDBExceptionHandler.postgres_sqlstate(exc) in _FAILOVER_SQLSTATES: + return True + message: Final = str(exc).lower() + return any(refusal in message for refusal in _FAILOVER_REFUSALS) + + +def _openai_shaped_proxy_exception(exc: Exception, message: str, status_code: int) -> ProxyException: + answered_status: Final = int(proxy_exception_status_code(message, str(status_code))) + return ProxyException( + message=message, + type=_catch_all_error_type(exc, answered_status), + param=openai_error_param(exc), + code=answered_status, + ) + + +def _catch_all_error_type(exc: Exception, status_code: int) -> str: + if _rejects_the_caller(exc, status_code): + return ProxyErrorTypes.auth_error.value + return openai_error_type(exc, status_code) + + +def _rejects_the_caller(exc: Exception, status_code: int) -> bool: + raised_as_http_answer: Final = isinstance(exc, HTTPException) + status_from_upstream: Final = error_status_code(exc, 0) == status_code + return status_code in _AUTH_FAILURE_STATUS_CODES and (raised_as_http_answer or not status_from_upstream) + + def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None: # mutable-ok: ProxyException.headers if litellm_call_id is None: return None diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 5dfcaa4b2f8..3f7ea334a99 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -36,7 +36,7 @@ _LOCK_TIMEOUT_SQLSTATE: Final = "55P03" _LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout" -def _exception_chain(e: BaseException) -> Iterator[BaseException]: +def exception_chain(e: BaseException) -> Iterator[BaseException]: current = e # rebind-ok: advances one link per iteration of the bounded walk for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH): yield current @@ -49,7 +49,7 @@ def _exception_chain(e: BaseException) -> Iterator[BaseException]: def _database_service_unavailable_errors(e: BaseException) -> tuple[Exception, ...]: return tuple( link - for link in _exception_chain(e) + for link in exception_chain(e) if isinstance(link, Exception) and PrismaDBExceptionHandler.is_database_service_unavailable_error(link) ) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 869e4953df0..396a00ab1a1 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -32,7 +32,6 @@ from litellm.proxy._types import ( EnterpriseLicenseData, Litellm_EntityType, LitellmUserRoles, - ProxyErrorTypes, ProxyException, SpecialModelNames, UserAPIKeyAuth, @@ -47,6 +46,7 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.model_checks import get_key_models from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils import http_parsing_utils +from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.db.health_check_latest import ( @@ -683,21 +683,9 @@ async def health_services_endpoint( except Exception as e: verbose_proxy_logger.error("litellm.proxy.proxy_server.health_services_endpoint(): Exception occured - %s", e) verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) from e def convert_health_check_to_dict(check) -> dict: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index f2f0bc5be0b..bb901f22649 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -46,6 +46,7 @@ from litellm.proxy.auth.password_policy import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast +from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, object_permission_cache_key, @@ -1729,21 +1730,9 @@ async def user_update( except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.user_update(): Exception occured - %s", e) verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e async def bulk_update_processed_users( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a19ecb8a9e9..55e59ca5776 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -75,6 +75,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, publish_config_change, ) +from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache @@ -4087,21 +4088,9 @@ async def update_key_fn( # update based on remaining passed in values except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.update_key_fn(): Exception occured - %s", e) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e @router.post( @@ -7244,21 +7233,9 @@ async def list_keys( except Exception as e: verbose_proxy_logger.exception("Error in list_keys: %s", e) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"error({e})"), - type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), - code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR) from e async def _apply_non_admin_alias_scope( @@ -7402,21 +7379,9 @@ async def key_aliases( except Exception as e: verbose_proxy_logger.exception("Error in key_aliases: %s", e) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"error({e})"), - type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) from e def _validate_sort_params( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 8cd3559ba86..e7ad7d04070 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -77,6 +77,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.proxy.management.teams.authz import TEAM_ADMIN_ONLY, is_team_admin @@ -2845,21 +2846,9 @@ async def delete_model( except Exception as e: verbose_proxy_logger.exception("Failed to delete model. Due to error - %s", e) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e async def delete_team_model_alias( @@ -2899,6 +2888,13 @@ async def delete_team_model_alias( return removed_model_aliases +def _model_not_added_to_db() -> HTTPException: + return HTTPException( + status_code=500, + detail={"error": "Failed to add model to db. Check your server logs for more details."}, + ) + + @router.post( "/model/new", description="Allows adding new models to the model list in the config.yaml", @@ -3080,6 +3076,8 @@ async def add_new_model( if isinstance(e, HTTPException): raise verbose_proxy_logger.exception("Exception in add_new_model: %s", e) + if model_response is None: + raise _model_not_added_to_db() from e else: raise HTTPException( @@ -3088,10 +3086,7 @@ async def add_new_model( ) if model_response is None: - raise HTTPException( - status_code=500, - detail={"error": "Failed to add model to db. Check your server logs for more details."}, - ) + raise _model_not_added_to_db() ## CREATE AUDIT LOG ## asyncio.create_task( @@ -3121,21 +3116,9 @@ async def add_new_model( except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.add_new_model(): Exception occured - %s", e) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e #### MODEL MANAGEMENT #### @@ -3364,21 +3347,9 @@ async def update_model( return None if model_response is None else _model_write_response(model_response, member_write) except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.update_model(): Exception occured - %s", e) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e @router.post( diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index e1d19d8efae..8a1b9c112dc 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -41,6 +41,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast +from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.budget_management_endpoints import ( @@ -1396,21 +1397,9 @@ async def organization_member_add( ) except Exception as e: verbose_proxy_logger.exception("Error adding member to organization: %s", e) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) from e async def find_member_if_email(user_email: str, prisma_client: PrismaClient) -> LiteLLM_UserTable: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index a256f417dfa..50b700e712d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -132,6 +132,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( ) from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch +from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, team_membership_auth_cache_key, @@ -5526,21 +5527,9 @@ async def team_info( e, traceback.format_exc(), ) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e @router.get( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d6ec62b25c2..bb8fffd85bc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -524,6 +524,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( LITELLM_CALL_ID_HEADER, headers_with_litellm_call_id, litellm_call_id_headers, + proxy_exception_for, with_litellm_call_id, ) from litellm.proxy.common_utils.periodic_reload_schedule import ( @@ -18536,25 +18537,13 @@ async def async_queue_request( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - elif isinstance(e, litellm.BadRequestError): + if isinstance(e, litellm.BadRequestError): raise ProxyException( message=e.message, type="invalid_request_error", param=getattr(e, "param", None), code=e.status_code - ) - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + ) from e + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e @app.get("/fallback/login", tags=["experimental"], include_in_schema=False) @@ -19833,21 +19822,9 @@ async def update_config( except Exception as e: verbose_proxy_logger.error("litellm.proxy.proxy_server.update_config(): Exception occured - %s", e) verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e ### CONFIG GENERAL SETTINGS @@ -21030,21 +21007,9 @@ async def get_config( } except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.get_config(): Exception occured - %s", e) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) + if isinstance(e, ProxyException): + raise + raise proxy_exception_for(e, default_status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) from e @router.get( diff --git a/tests/integration/authorization/test_team_admin_gate.py b/tests/integration/authorization/test_team_admin_gate.py index a7513277609..eff4060631e 100644 --- a/tests/integration/authorization/test_team_admin_gate.py +++ b/tests/integration/authorization/test_team_admin_gate.py @@ -746,7 +746,7 @@ def test_internal_user_model_creation_prohibition_refuses_team_admin_but_not_pro "error": { "message": "Model creation is disabled for internal users by disable_model_add_for_internal_users.", "type": "auth_error", - "param": "None", + "param": None, "code": "403", } }, diff --git a/tests/integration/authorization/test_user_self_update_guard.py b/tests/integration/authorization/test_user_self_update_guard.py index d83960e1e85..d058b3c6998 100644 --- a/tests/integration/authorization/test_user_self_update_guard.py +++ b/tests/integration/authorization/test_user_self_update_guard.py @@ -38,7 +38,7 @@ def _escalation_error(field: str) -> dict[str, JsonValue]: "error": { "message": f"{{'error': \"Non-admin users cannot modify '{field}' on their own record. Contact your proxy admin.\"}}", "type": "auth_error", - "param": "None", + "param": None, "code": "403", } } @@ -49,7 +49,7 @@ def _role_change_error() -> dict[str, JsonValue]: "error": { "message": "Only proxy admins can modify user roles.", "type": "auth_error", - "param": "None", + "param": None, "code": "403", } } diff --git a/tests/integration/management/test_team_delete_large_membership.py b/tests/integration/management/test_team_delete_large_membership.py index 1a8140373b0..8b3c6b06bba 100644 --- a/tests/integration/management/test_team_delete_large_membership.py +++ b/tests/integration/management/test_team_delete_large_membership.py @@ -255,8 +255,8 @@ def _team_not_found_body(team_id: str) -> dict[str, JsonValue]: return { "error": { "message": f"{{'message': 'Team not found, passed team id: {team_id}.'}}", - "type": "auth_error", - "param": "None", + "type": "invalid_request_error", + "param": None, "code": "404", } } diff --git a/tests/unit/proxy/common_utils/test_openai_error_payload.py b/tests/unit/proxy/common_utils/test_openai_error_payload.py index 40bb84ff538..acefc4c1e2c 100644 --- a/tests/unit/proxy/common_utils/test_openai_error_payload.py +++ b/tests/unit/proxy/common_utils/test_openai_error_payload.py @@ -1,16 +1,26 @@ import json +import sys +from typing import Final +import httpx +import prisma.errors import pytest from fastapi import HTTPException +from prisma._async_http import Response +from prisma.engine.utils import handle_response_errors +import litellm from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, litellm_call_id_headers, openai_error_param, openai_error_type, + proxy_exception_for, with_litellm_call_id, ) +from litellm.types.router import RouterErrors @pytest.mark.parametrize( @@ -203,3 +213,287 @@ def test_a_proxy_error_keeps_the_call_id_it_was_raised_with(): assert with_litellm_call_id(raised_with_id, "second").headers == {"x-litellm-call-id": "first"} assert with_litellm_call_id(ProxyException(message="nope", type="None", param=None, code=400), None).headers == {} + + +@pytest.mark.parametrize( + "raised_status, expected_type", + [ + (400, "invalid_request_error"), + (404, "invalid_request_error"), + (401, "auth_error"), + (403, "auth_error"), + (422, "invalid_request_error"), + (500, "internal_server_error"), + ], +) +def test_a_catch_all_handler_keeps_the_status_its_route_answered_with(raised_status: int, expected_type: str) -> None: + wrapped: Final = proxy_exception_for( + HTTPException(status_code=raised_status, detail={"message": "boom"}), + default_status_code=400, + ) + + assert (wrapped.code, wrapped.type) == (str(raised_status), expected_type) + + +@pytest.mark.parametrize( + "default_status_code, expected_type", [(400, "invalid_request_error"), (500, "internal_server_error")] +) +def test_an_exception_carrying_no_status_falls_back_to_the_routes_default(default_status_code: int, expected_type: str) -> None: + wrapped: Final = proxy_exception_for(ValueError("boom"), default_status_code=default_status_code) + + assert (wrapped.code, wrapped.type) == (str(default_status_code), expected_type) + + +def test_a_proxy_exception_passes_through_with_the_type_it_already_named() -> None: + raised: Final = ProxyException( + message="Budget has been exceeded", + type=ProxyErrorTypes.budget_exceeded.value, + param="max_budget", + code=400, + ) + + assert proxy_exception_for(raised, default_status_code=500) is raised + + +def test_the_message_is_what_the_route_raised_without_an_invented_prefix() -> None: + assert proxy_exception_for(ValueError("boom"), default_status_code=400).message == "boom" + + +def test_an_http_exceptions_detail_survives_the_wrapping() -> None: + wrapped: Final = proxy_exception_for( + HTTPException(status_code=400, detail={"error": "Invalid budget_duration 'not-a-duration'."}), + default_status_code=400, + ) + + assert "Invalid budget_duration" in wrapped.message + + +def test_a_carried_type_survives_the_wrapping() -> None: + class _Carrier(Exception): + type: str = "context_window_exceeded" + status_code: int = 400 + + assert proxy_exception_for(_Carrier("boom"), default_status_code=500).type == "context_window_exceeded" + + +@pytest.mark.parametrize("exc", [HTTPException(status_code=404, detail="nope"), ValueError("boom")]) +def test_the_wrapped_error_serializes_as_an_openai_error_object(exc: Exception) -> None: + body: Final = json.loads(json.dumps({"error": proxy_exception_for(exc, default_status_code=400).to_dict()})) + + assert body["error"]["param"] is None + assert body["error"]["type"] == "invalid_request_error" + + +@pytest.mark.parametrize( + "provider_failure, expected", + [ + ( + litellm.AuthenticationError(message="invalid x-api-key", llm_provider="anthropic", model="claude-haiku-5-5"), + ("401", "authentication_error"), + ), + ( + litellm.PermissionDeniedError( + message="model not allowed", + llm_provider="anthropic", + model="claude-haiku-5-5", + response=httpx.Response(403, request=httpx.Request("POST", "https://api.anthropic.com/v1/messages")), + ), + ("403", "permission_error"), + ), + ], +) +def test_a_provider_credential_failure_is_labelled_the_way_chat_completions_labels_it( + provider_failure: Exception, expected: tuple[str, str] +) -> None: + wrapped: Final = proxy_exception_for(provider_failure, default_status_code=400) + + assert (wrapped.code, wrapped.type) == expected + + +@pytest.mark.parametrize("status_code", [401, 403]) +def test_a_rejection_the_route_raised_keeps_the_auth_label(status_code: int) -> None: + wrapped: Final = proxy_exception_for(HTTPException(status_code=status_code, detail="no"), default_status_code=400) + + assert (wrapped.code, wrapped.type) == (str(status_code), "auth_error") + + +def test_an_exception_without_a_message_is_named_by_its_type() -> None: + wrapped: Final = proxy_exception_for(TimeoutError(), default_status_code=500) + + assert wrapped.message == "TimeoutError" + + +def test_a_database_outage_answers_503_so_callers_retry_instead_of_fixing_their_request() -> None: + wrapped: Final = proxy_exception_for( + prisma.errors.PrismaError("Can't reach database server at `db:5432`"), default_status_code=400 + ) + + assert (wrapped.code, wrapped.type) == ("503", "no_db_connection") + + +@pytest.mark.parametrize( + "postgres_refusal", + [ + 'FATAL: database "litellm" is not currently accepting connections', + "FATAL: the database system is starting up", + "FATAL: the database system is shutting down", + "FATAL: the database system is in recovery mode", + "FATAL: the database system is not yet accepting connections", + "FATAL: terminating connection due to administrator command", + "ERROR: cannot execute UPDATE in a read-only transaction", + ], +) +def test_a_database_refusing_work_during_a_failover_answers_503(postgres_refusal: str) -> None: + refused: Final = prisma.errors.DataError( + {"user_facing_error": {"message": f"Error in connector: Error querying the database: {postgres_refusal}"}} + ) + + wrapped: Final = proxy_exception_for(refused, default_status_code=400) + + assert (wrapped.code, wrapped.type) == ("503", "no_db_connection") + + +@pytest.mark.parametrize( + "localized_refusal", + [ + { + "error_code": "P2010", + "message": "Raw query failed. Message: `FEHLER: kann UPDATE nicht in einer Nur-Lese-Transaktion ausführen`", + "meta": {"code": "25006"}, + }, + { + "message": 'PostgresError { code: "57P01", message: "Verbindung wird abgebrochen", severity: "FATAL" }', + }, + ], +) +def test_a_failover_refusal_in_a_localized_postgres_answers_503(localized_refusal: dict[str, object]) -> None: + refused: Final = prisma.errors.DataError({"user_facing_error": localized_refusal}) + + wrapped: Final = proxy_exception_for(refused, default_status_code=400) + + assert (wrapped.code, wrapped.type) == ("503", "no_db_connection") + + +def test_a_failover_phrase_outside_the_database_is_not_reported_as_a_database_outage() -> None: + callback_down: Final = ValueError("FATAL: the database system is starting up") + + wrapped: Final = proxy_exception_for(callback_down, default_status_code=400) + + assert (wrapped.code, wrapped.type) == ("400", "invalid_request_error") + + +def test_a_database_dropped_mid_query_answers_503_though_prisma_fails_to_name_the_error() -> None: + engine_reply_with_null_meta: Final = [{"user_facing_error": {"error_code": "P1001", "message": "", "meta": None}}] + with pytest.raises(AttributeError) as raised: + handle_response_errors(Response(httpx.Response(500)), engine_reply_with_null_meta) + + wrapped: Final = proxy_exception_for(raised.value, default_status_code=500) + + assert (wrapped.code, wrapped.type) == ("503", "no_db_connection") + + +def test_a_network_failure_outside_the_database_is_not_reported_as_a_database_outage() -> None: + callback_unreachable: Final = httpx.ConnectError("[Errno 111] Connection refused") + + wrapped: Final = proxy_exception_for(callback_unreachable, default_status_code=500) + + assert (wrapped.code, wrapped.type) == ("500", "internal_server_error") + + +def _add_member_failing_with(db_error: Exception) -> None: + try: + raise db_error + except Exception as e: + raise ValueError(f"Error adding member to organization: {e}") + + +@pytest.mark.parametrize( + "underlying_error, expected", + [ + (prisma.errors.PrismaError("Can't reach database server at `db:5432`"), ("503", "no_db_connection")), + (httpx.ConnectError("[Errno 111] Connection refused"), ("500", "internal_server_error")), + ], +) +def test_a_route_that_rewraps_its_error_keeps_the_database_outage_it_wrapped( + underlying_error: Exception, expected: tuple[str, str] +) -> None: + with pytest.raises(ValueError, match="Error adding member") as raised: + _add_member_failing_with(underlying_error) + + wrapped: Final = proxy_exception_for(raised.value, default_status_code=500) + + assert (wrapped.code, wrapped.type) == expected + + +def test_a_proxy_without_prisma_installed_still_answers_with_an_error_object(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(sys.modules, "prisma", None) + + wrapped: Final = proxy_exception_for(httpx.ConnectError("[Errno 111] Connection refused"), default_status_code=500) + + assert (wrapped.code, wrapped.type) == ("500", "internal_server_error") + + +def _look_up_team_failing_with(lookup_error: Exception) -> None: + try: + raise lookup_error + except Exception: + raise HTTPException(status_code=404, detail="Team doesn't exist in db") + + +@pytest.mark.parametrize( + "lookup_error, expected", + [ + (DBLookupDeadlineExceeded("team", 10.0), ("503", "no_db_connection")), + (prisma.errors.PrismaError("Can't reach database server at `db:5432`"), ("503", "no_db_connection")), + (KeyError("team_id"), ("404", "invalid_request_error")), + ], +) +def test_a_not_found_answer_raised_while_the_database_was_down_reports_the_outage( + lookup_error: Exception, expected: tuple[str, str] +) -> None: + with pytest.raises(HTTPException) as raised: + _look_up_team_failing_with(lookup_error) + + wrapped: Final = proxy_exception_for(raised.value, default_status_code=400) + + assert (wrapped.code, wrapped.type) == expected + + +def test_a_lookup_that_missed_its_deadline_answers_503() -> None: + wrapped: Final = proxy_exception_for(DBLookupDeadlineExceeded("team", 10.0), default_status_code=400) + + assert (wrapped.code, wrapped.type) == ("503", "no_db_connection") + + +def test_a_database_outage_answer_does_not_reveal_where_the_database_lives() -> None: + wrapped: Final = proxy_exception_for( + prisma.errors.PrismaError("Can't reach database server at `prod-db.internal:5432`"), default_status_code=400 + ) + + assert wrapped.code == "503" + assert "prod-db.internal" not in wrapped.message + + +def test_an_upstream_failure_answers_with_its_own_status_not_the_routes_default() -> None: + upstream_overloaded: Final = litellm.ServiceUnavailableError( + message="upstream overloaded", llm_provider="openai", model="gpt-5.4-mini" + ) + + wrapped: Final = proxy_exception_for(upstream_overloaded, default_status_code=400) + + assert (wrapped.code, wrapped.type) == ("503", "internal_server_error") + + +@pytest.mark.parametrize( + "router_rejection, expected", + [ + (f"{RouterErrors.no_deployments_with_tag_routing.value}. Passed model=gpt-5.4-mini", ("401", "auth_error")), + (f"{RouterErrors.no_deployments_available.value}. Passed model=gpt-5.4-mini", ("429", "rate_limit_error")), + ], +) +def test_a_router_rejection_is_labelled_by_the_status_it_answers_with( + router_rejection: str, expected: tuple[str, str] +) -> None: + wrapped: Final = proxy_exception_for(ValueError(router_rejection), default_status_code=400) + + assert (wrapped.code, wrapped.type) == expected diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 49183a143c0..53c163b89ae 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -5480,6 +5480,15 @@ async def test_health_services_endpoint_langfuse_missing_keys_errors(monkeypatch assert "connection refused by lf.internal.example" in str(raised.value.message) +@pytest.mark.asyncio +async def test_health_services_reports_an_unknown_service_as_a_request_error_not_an_auth_error() -> None: + with pytest.raises(ProxyException) as exc_info: + await health_services_endpoint(service="totally_unknown_service_xyz") + + assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error") + assert "Authentication Error" not in exc_info.value.message + + @pytest.mark.asyncio @pytest.mark.parametrize("admin", [False, True]) async def test_unsaved_chatgpt_account_probe_requires_proxy_admin( diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index cb8a73d581b..d1bbf444d2c 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -5211,3 +5211,18 @@ async def test_user_update_without_password_revokes_nothing(_admin_prisma, mocke await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) revoke_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error(_admin_prisma: MagicMock) -> None: + from litellm.proxy.management_endpoints.internal_user_endpoints import user_update + + with pytest.raises(ProxyException) as exc_info: + await user_update( + data=UpdateUserRequest(user_id="target-user", budget_duration="not-a-duration"), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error") + assert "Invalid budget_duration" in exc_info.value.message + assert "Authentication Error" not in exc_info.value.message diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index b13a7bfc33e..859824e4a72 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -6538,7 +6538,7 @@ async def test_list_keys_with_invalid_status(): ) # Verify ProxyException properties - assert exc_info.value.code == "400" + assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error") assert "Invalid status value" in str(exc_info.value.message) assert "deleted" in str(exc_info.value.message) @@ -23377,3 +23377,73 @@ def test_key_info_database_not_connected_sanitized_500( assert "Database not connected" in caplog.text finally: proxy_server.app.dependency_overrides.pop(user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_update_key_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None: + from fastapi import Request + + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + monkeypatch.setattr( # test-quality-ok: the endpoint reads proxy_server.prisma_client itself, no parameter to inject + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ) + + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=MagicMock(spec=Request), + data=UpdateKeyRequest(key="sk-whatever", budget_duration="not-a-duration"), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error") + assert "Authentication Error" not in exc_info.value.message + + +@pytest.mark.asyncio +async def test_update_key_passes_its_own_rejection_through_without_making_it_its_own_cause( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import Request + + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr( # test-quality-ok: the endpoint reads proxy_server.prisma_client itself, no parameter to inject + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) + + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=MagicMock(spec=Request), + data=UpdateKeyRequest(key="sk-missing", max_budget=5), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("404", "not_found_error") + assert exc_info.value.__cause__ is None + + +@pytest.mark.asyncio +async def test_key_aliases_answers_503_when_the_database_is_unreachable(monkeypatch: pytest.MonkeyPatch) -> None: + import prisma.errors + + from litellm.proxy.management_endpoints.key_management_endpoints import key_aliases + + unreachable_db: Final = MagicMock() + unreachable_db.db.query_raw = AsyncMock( + side_effect=prisma.errors.PrismaError("Can't reach database server at `db:5432`") + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", unreachable_db) + + with pytest.raises(ProxyException) as exc_info: + await key_aliases( + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + page=1, + size=50, + search=None, + team_id=None, + ) + + assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection") diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 437d948f8f5..b892dc97e5c 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -311,7 +311,7 @@ async def test_model_creation_fails_closed_when_authoritative_policy_is_unavaila Deployment(model_name="cannot-create", litellm_params={"model": "openai/test-model"}), UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.INTERNAL_USER), ) - assert error.value.code == "503" + assert (error.value.code, error.value.type) == ("503", "internal_server_error") prisma.db.litellm_proxymodeltable.create.assert_not_called() prisma.db.litellm_teamtable.find_unique.assert_not_called() @@ -10913,3 +10913,76 @@ class TestNonAdminCannotSetPerUserOauthOnModel: user_api_key_dict=admin, ) assert result is created_row + + +@pytest.mark.asyncio +async def test_delete_model_reports_an_unknown_model_id_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import ModelInfoDelete + from litellm.proxy.management_endpoints.model_management_endpoints import delete_model + + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + with pytest.raises(ProxyException) as exc_info: + await delete_model( + model_info=ModelInfoDelete(id="no-such-model"), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error") + assert "Authentication Error" not in exc_info.value.message + + +@pytest.mark.asyncio +async def test_update_model_reports_an_unknown_model_id_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.model_management_endpoints import update_model + + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", None) + + with pytest.raises(ProxyException) as exc_info: + await update_model( + model_params=updateDeployment(model_info=ModelInfo(id="no-such-model")), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error") + assert exc_info.value.message == "model not found" + + +@pytest.mark.asyncio +async def test_add_new_model_reports_a_database_outage_during_the_insert_as_503( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import prisma.errors + + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + unreachable_db: Final = MagicMock() + unreachable_db.db.query_raw = AsyncMock(return_value=[]) + unreachable_db.db.litellm_proxymodeltable.create = AsyncMock( + side_effect=prisma.errors.PrismaError("Can't reach database server at `db:5432`") + ) + monkeypatch.setattr(proxy_server, "prisma_client", unreachable_db) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-outage-salt") + + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="outage-model", + litellm_params=LiteLLM_Params(model="openai/gpt-5.4-mini", api_key="fake-key"), + model_info={"id": "outage-model-1"}, + ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection") diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 16c59aa20b9..5e06756f3a9 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1617,6 +1617,37 @@ async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu prisma_client.db.litellm_organizationtable.delete.assert_not_called() +@pytest.mark.asyncio +async def test_organization_member_add_reports_an_unknown_org_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None: + from fastapi import Request + + from litellm.proxy._types import ( + LitellmUserRoles, + OrganizationMemberAddRequest, + OrgMember, + ProxyException, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add + + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.litellm_organizationtable.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + with pytest.raises(ProxyException) as exc_info: + await organization_member_add( + data=OrganizationMemberAddRequest( + organization_id="no-such-org", + member=OrgMember(role="internal_user", user_id="someone"), + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("404", "invalid_request_error") + assert "Authentication Error" not in exc_info.value.message + + @pytest.mark.asyncio @pytest.mark.parametrize("budget_duration", ["30d", None]) async def test_new_organization_schedules_budget_reset_from_duration(monkeypatch, budget_duration): diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index e51f26bf1c3..c6d12ade2a0 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -19372,3 +19372,48 @@ async def test_default_budget_cache_eviction_does_not_fail_after_redis_error() - ) user_api_key_cache.async_delete_cache.assert_awaited_once_with(key="team_member_default_budget:budget_456") + + +@pytest.mark.asyncio +async def test_team_info_reports_a_missing_team_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None: + from fastapi import Request + + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints import team_endpoints + + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + with pytest.raises(ProxyException) as exc_info: + await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="no-such-team", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("404", "invalid_request_error") + assert "Authentication Error" not in exc_info.value.message + + +@pytest.mark.asyncio +async def test_team_info_keeps_the_auth_label_for_a_caller_outside_the_team(monkeypatch: pytest.MonkeyPatch) -> None: + from fastapi import Request + + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints import team_endpoints + + other_team: Final = MagicMock() + other_team.model_dump.return_value = {"team_id": "team-a", "members_with_roles": []} + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=other_team) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + with pytest.raises(ProxyException) as exc_info: + await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-a", + user_api_key_dict=UserAPIKeyAuth(user_id="outsider", user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("403", "auth_error") diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 976ed4f28a5..22d6316c448 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -3,6 +3,7 @@ import sys import traceback from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping from contextlib import asynccontextmanager +from pathlib import Path from typing import Final, Literal, NoReturn, cast from unittest import mock @@ -3817,6 +3818,48 @@ def test_live_session_preserves_pretty_printed_start_frame( assert routed_data["live_session_start"] == expected_start +@pytest.mark.asyncio +async def test_update_config_reports_an_unsupported_router_setting_as_a_request_error_not_an_auth_error() -> None: + from fastapi import Request + + from litellm.proxy._types import ConfigYAML, ProxyException + from litellm.proxy.proxy_server import update_config + + body: Final = {"router_settings": {"not_a_router_setting": True}} + request: Final = MagicMock(spec=Request) + request.json = AsyncMock(return_value=body) + + with pytest.raises(ProxyException) as exc_info: + await update_config( + config_info=ConfigYAML(**body), + request=request, + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error") + assert "Authentication Error" not in exc_info.value.message + + +@pytest.mark.asyncio +async def test_get_config_callbacks_reports_an_unreadable_config_as_a_server_error_not_an_auth_error( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import ProxyException + from litellm.proxy.proxy_server import get_config + + empty_config: Final = tmp_path / "config.yaml" + empty_config.write_text("") + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr(proxy_server, "user_config_file_path", str(empty_config)) + + with pytest.raises(ProxyException) as exc_info: + await get_config(user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)) + + assert (exc_info.value.code, exc_info.value.type) == ("500", "internal_server_error") + assert exc_info.value.message == "Config cannot be None or Empty." + + @pytest.mark.parametrize("scope", [None, "expand"]) @pytest.mark.parametrize( "general_settings, wildcard_query, wildcard_listed", @@ -3861,6 +3904,62 @@ async def test_model_list_return_wildcard_routes_defaults_to_general_setting( assert ("openai/*" in listed) is wildcard_listed, listed +@pytest.mark.asyncio +async def test_queue_request_labels_an_upstream_failure_a_server_error_not_an_auth_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from starlette.requests import Request + from starlette.responses import Response + + from litellm.proxy import proxy_server + from litellm.proxy._types import ProxyException + from litellm.proxy.proxy_server import async_queue_request + + body: Final = json.dumps( + {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "hi"}], "priority": 0} + ).encode() + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/queue/chat/completions", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + "scheme": "http", + "server": ("testserver", 80), + "root_path": "", + }, + receive, + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-5.4-mini", + "litellm_params": { + "model": "openai/gpt-5.4-mini", + "api_key": "sk-test", + "mock_response": "litellm.InternalServerError", + }, + } + ], + num_retries=0, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + + with pytest.raises(ProxyException) as exc_info: + await async_queue_request( + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert (exc_info.value.code, exc_info.value.type) == ("500", "internal_server_error") + + _DEBUG_LOG_SENTINEL: Final = "sentinel-request-body-7f3a" _LONG_PROMPT: Final = f"{_DEBUG_LOG_SENTINEL} " + "Keep edits minimal and explain them briefly. " * 20 _TOOL_PARAMETERS: Final = {"type": "object", "properties": {"cmd": {"type": "string"}}} From a66dc50efbdb6984797304d6fdd91dc5358647e8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:04:47 -0700 Subject: [PATCH 04/42] test(json_fragment_accumulator): assert peak allocation instead of a drain timing ratio (#43370) * test(json_fragment_accumulator): assert peak allocation instead of a drain timing ratio The drain regression guard compared the wall-clock time of draining 80,000 values against 40,000 and tripped its 3x bound whenever the CI runner was busy. A cursor drain allocates only per-value decode scratch while every re-slice of the remainder copies it, so the guard now bounds peak allocation during the drain, which does not move with runner load and still fails on every quadratic variant * test(json_fragment_accumulator): widen the drain's peak bound to a quarter of the payload * test(json_fragment_accumulator): guard copying drains by peak memory and rescanning drains by CPU-time scaling * test(json_fragment_accumulator): move the drain's timing guard from the unit suite into a CodSpeed benchmark * test(json_fragment_accumulator): measure only the pops after the buffer is joined, so a drain that copies once passes * test(json_fragment_accumulator): separate every other streamed value with a space, so a copy after whitespace fails the drain test --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- ...st_json_fragment_accumulator_benchmarks.py | 26 +++++++ .../test_json_fragment_accumulator.py | 71 +++++++++++-------- 2 files changed, 67 insertions(+), 30 deletions(-) create mode 100644 tests/benchmarks/test_json_fragment_accumulator_benchmarks.py diff --git a/tests/benchmarks/test_json_fragment_accumulator_benchmarks.py b/tests/benchmarks/test_json_fragment_accumulator_benchmarks.py new file mode 100644 index 00000000000..807dd8ff236 --- /dev/null +++ b/tests/benchmarks/test_json_fragment_accumulator_benchmarks.py @@ -0,0 +1,26 @@ +from collections.abc import Iterator +from typing import Final + +import pytest + +from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator + +CONCATENATED_VALUES: Final = 20_000 + + +def _popped_values(accumulator: JSONFragmentAccumulator) -> Iterator[object]: + while True: + found, value = accumulator.pop_next_value() + if not found: + return + yield value + + +@pytest.mark.benchmark +def test_draining_many_concatenated_values() -> None: + accumulator: Final = JSONFragmentAccumulator() + accumulator.append('{"a": 1}' * CONCATENATED_VALUES) + + drained: Final = sum(1 for _ in _popped_values(accumulator)) + + assert drained == CONCATENATED_VALUES diff --git a/tests/unit/litellm_core_utils/test_json_fragment_accumulator.py b/tests/unit/litellm_core_utils/test_json_fragment_accumulator.py index d16230afe3a..f27a5fd2307 100644 --- a/tests/unit/litellm_core_utils/test_json_fragment_accumulator.py +++ b/tests/unit/litellm_core_utils/test_json_fragment_accumulator.py @@ -1,5 +1,8 @@ import json import time +import tracemalloc +from collections.abc import Iterator +from typing import Final from unittest.mock import patch from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator @@ -175,40 +178,48 @@ def test_accumulation_of_many_fragments_is_not_quadratic(): assert elapsed_ms < 50, f"1000-fragment append took {elapsed_ms:.1f} ms (expected < 50 ms); O(n^2) regression?" -def test_draining_many_concatenated_values_is_not_quadratic(): - """ - Regression guard: peeling N JSON values already sitting in one buffer, - one pop_next_value() call per value with no new fragments in between, - must be O(n) total. Re-copying the shrinking remainder on every pop - (slicing a new string instead of advancing a cursor) makes total drain - time scale with the square of the buffer size. +def _popped_values(accumulator: JSONFragmentAccumulator) -> Iterator[object]: + while True: + found, value = accumulator.pop_next_value() + if not found: + return + yield value - Uses a doubling ratio rather than an absolute ms budget so it isn't - flaky on a slower or busier CI runner: doubling the input should - roughly double an O(n) drain's time but roughly quadruple an O(n^2) - drain's time, and that ratio holds regardless of machine speed. - """ - def drain_time_ms(n: int) -> float: - accumulator = JSONFragmentAccumulator() - accumulator.append('{"a": 1}' * n) - start = time.perf_counter() - drained = 0 - while True: - found, _ = accumulator.pop_next_value() - if not found: - break - drained += 1 - assert drained == n - return (time.perf_counter() - start) * 1000 +def _streamed(fragments: tuple[str, ...]) -> JSONFragmentAccumulator: + accumulator: Final = JSONFragmentAccumulator() + for fragment in fragments: + accumulator.append(fragment) + return accumulator - small_ms = drain_time_ms(40_000) - large_ms = drain_time_ms(80_000) - ratio = large_ms / max(small_ms, 0.001) - assert ratio < 3.0, ( - f"doubling drained values scaled time by {ratio:.2f}x ({small_ms:.1f} ms -> {large_ms:.1f} ms); " - "expected roughly 2x for O(n); O(n^2) regression?" +def _drain_peak_bytes(fragments: tuple[str, ...]) -> int: + accumulator: Final = _streamed(fragments) + first_found, _ = accumulator.pop_next_value() + assert first_found is True + already_tracing: Final = tracemalloc.is_tracing() + if not already_tracing: + tracemalloc.start() + tracemalloc.reset_peak() + baseline, _ = tracemalloc.get_traced_memory() + try: + drained: Final = sum(1 for _ in _popped_values(accumulator)) + _, peak = tracemalloc.get_traced_memory() + finally: + if not already_tracing: + tracemalloc.stop() + assert drained == len(fragments) - 1 + return peak - baseline + + +def test_draining_many_concatenated_values_never_copies_the_remainder(): + value: Final = '{"text": "' + "x" * 20 + '"}' + fragments: Final = tuple(value if index % 2 else f" {value}" for index in range(40_000)) + payload_bytes: Final = sum(len(fragment) for fragment in fragments) + peak_bytes: Final = _drain_peak_bytes(fragments) + assert peak_bytes < payload_bytes // 4, ( + f"draining {len(fragments)} values allocated a peak of {peak_bytes} bytes over a {payload_bytes} byte payload, " + "copying the remainder on every pop?" ) From b7da9a275581ccec82538428dfd0b84bbc4fb6fc Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:51:39 -0700 Subject: [PATCH 05/42] fix(credentials): keep the stored secret when PATCH echoes its masked read-back (#48950) * test(credentials): pin that PATCH keeps the stored secret when it echoes the masked read-back * fix(credentials): keep the stored secret when PATCH echoes its masked read-back * fix(credentials): only treat a mask on a key GET masks as a placeholder PATCH now skips mask-shaped values only on the keys get_masked_values masks, using the masker's own keyword check, so a value like api_base ab****cd is saved as sent again. Also drops the new test docstrings and tightens the new test helper types --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 34 +-- .../proxy/credential_endpoints/endpoints.py | 84 +++++- .../credential_endpoints/test_endpoints.py | 265 +++++++++++++++++- 3 files changed, 353 insertions(+), 30 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 620c8413f29..7f4b069ded6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5052,6 +5052,23 @@ class Logging(LiteLLMLoggingBaseClass): return result_copy +_SENSITIVE_KEYWORDS: Final = ( + "authorization", + "token", + "key", + "secret", + "vertex_credentials", + "credentials", + "password", + "passwd", +) + + +def has_sensitive_keyword(key: str) -> bool: + lowered: Final = key.lower() + return any(keyword in lowered for keyword in _SENSITIVE_KEYWORDS) + + def get_masked_values( sensitive_object: Mapping[str, object], ignore_sensitive_values: bool = False, @@ -5070,16 +5087,6 @@ def get_masked_values( masked_length: Optional length for the masked portion (number of *). If set, will use exactly this many * regardless of original string length. The total length will be unmasked_length + masked_length. """ - sensitive_keywords: Final = ( - "authorization", - "token", - "key", - "secret", - "vertex_credentials", - "credentials", - "password", - "passwd", - ) def _mask_value(v: object) -> object: if isinstance(v, dict): @@ -5106,12 +5113,7 @@ def get_masked_values( return v[: unmasked_length // 2] + "*" * (len(v) - unmasked_length) + v[-unmasked_length // 2 :] return { - k: ( - v - if ignore_sensitive_values - or not any(sensitive_keyword in k.lower() for sensitive_keyword in sensitive_keywords) - else _mask_value(v) - ) + k: (v if ignore_sensitive_values or not has_sensitive_keyword(k) else _mask_value(v)) for k, v in sensitive_object.items() } diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index d20b5602186..57c7a753d71 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -5,6 +5,7 @@ CRUD endpoints for storing reusable credentials. import time import unicodedata from collections.abc import Mapping +from types import MappingProxyType from typing import ( Annotated, Final, @@ -23,7 +24,7 @@ from litellm.constants import ( GITHUB_COPILOT_PER_USER_AUTH_TYPE, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor -from litellm.litellm_core_utils.litellm_logging import get_masked_values +from litellm.litellm_core_utils.litellm_logging import get_masked_values, has_sensitive_keyword from litellm.llms.anthropic.wif import ( ExportedJwks, NotAnInternalIssuerCredential, @@ -83,6 +84,11 @@ _CREDENTIAL_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _GITHUB_COPILOT_PROVIDER: Final = "github_copilot" _DISPLAY_NAME_MAX_LENGTH: Final = 255 _DISPLAY_NAME_FORBIDDEN_CATEGORIES: Final = frozenset({"Cc", "Cf", "Cs", "Zl", "Zp"}) +_MASK_UNMASKED_LENGTH: Final = 4 +_MASK_ASTERISKS: Final = 4 +_SHORT_SECRET_MASK: Final = "*****" +_MASKED_READ_BACK_DEPTH: Final = 20 +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) def _normalized_github_copilot_values( @@ -395,7 +401,14 @@ async def get_credentials( """ try: masked_credentials: Final = [ - _credential_view(credential, get_masked_values(credential.credential_values)) + _credential_view( + credential, + get_masked_values( + credential.credential_values, + unmasked_length=_MASK_UNMASKED_LENGTH, + number_of_asterisks=_MASK_ASTERISKS, + ), + ) for credential in litellm.credential_list ] return {"success": True, "credentials": masked_credentials} @@ -423,7 +436,11 @@ async def get_credential_by_name( if credential.credential_name == credential_name: return _credential_view( credential, - get_masked_values(credential.credential_values, unmasked_length=4, number_of_asterisks=4), + get_masked_values( + credential.credential_values, + unmasked_length=_MASK_UNMASKED_LENGTH, + number_of_asterisks=_MASK_ASTERISKS, + ), ) raise HTTPException( status_code=404, @@ -521,8 +538,8 @@ async def get_credential_by_model( raise HTTPException(status_code=404, detail="Model not found") masked_credential_values: Final = get_masked_values( credential_values, - unmasked_length=4, - number_of_asterisks=4, + unmasked_length=_MASK_UNMASKED_LENGTH, + number_of_asterisks=_MASK_ASTERISKS, ) credential: Final = CredentialItem( credential_name=f"{model.model_name}-credential-{model_id}", @@ -908,6 +925,48 @@ def update_db_credential( return merged_credential +def _is_masked_read_back(value: object) -> bool: + if not isinstance(value, str): + return False + if value == _SHORT_SECRET_MASK: + return True + shown: Final = _MASK_UNMASKED_LENGTH // 2 + return len(value) == _MASK_UNMASKED_LENGTH + _MASK_ASTERISKS and value[shown:-shown] == "*" * _MASK_ASTERISKS + + +def _as_fields(value: object) -> Mapping[str, object] | None: + try: + return _CREDENTIAL_DICT_ADAPTER.validate_python(value, strict=True) + except ValidationError: + return None + + +def _kept_value(key: str, value: object, stored: Mapping[str, object], depth: int) -> tuple[str, object] | None: + if not has_sensitive_keyword(key): + return key, value + fields: Final = _as_fields(value) + if fields is not None and depth < _MASKED_READ_BACK_DEPTH: + return key, _without_masked_read_backs(fields, _as_fields(stored.get(key)) or _NO_FIELDS, depth + 1) + if not _is_masked_read_back(value): + return key, value + if key in stored: + return key, stored[key] + return None + + +def _without_masked_read_backs( + requested: Mapping[str, object], stored: Mapping[str, object], depth: int = 0 +) -> Mapping[str, object]: + kept: Final = (_kept_value(key, value, stored, depth) for key, value in requested.items()) + return _CREDENTIAL_DICT_ADAPTER.validate_python( + MappingProxyType(dict(entry for entry in kept if entry is not None)) + ) + + +def _stored_objects(stored_values: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in stored_values.items() if _as_fields(value) is not None}) + + @router.patch( "/credentials/{credential_name:path}", dependencies=[Depends(user_api_key_auth)], @@ -936,12 +995,12 @@ async def update_credential( ) requested_display_name: Final = _normalized_display_name(credential.display_name) _reject_overlapping_credential_values(credential) - incoming_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + requested_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( _resolve_deployment_credentials(llm_router, credential.model_id) if credential.model_id else credential.credential_values or {} ) - _reject_non_admin_wif_fields(_incoming_wif_fields(incoming_values, credential), user_api_key_dict) + _reject_non_admin_wif_fields(_incoming_wif_fields(requested_values, credential), user_api_key_dict) if prisma_client is None: raise HTTPException( status_code=500, @@ -952,19 +1011,20 @@ async def update_credential( if db_credential is None: raise _not_found_unless_config_defined(credential_name, "Credential not found in DB.") _reject_non_admin_wif_fields(_stored_wif_fields(db_credential), user_api_key_dict) + decrypted_stored_credential: Final = decrypted_credential(db_credential) + decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + decrypted_stored_credential.credential_values + ) + incoming_values: Final = _without_masked_read_backs(requested_values, _stored_objects(decrypted_values)) patch: Final = CredentialItem( credential_name=credential_name, display_name=( requested_display_name if "display_name" in credential.model_fields_set else db_credential.display_name ), credential_info=_CREDENTIAL_DICT_ADAPTER.validate_python(credential.credential_info), - credential_values=incoming_values, + credential_values=_CREDENTIAL_DICT_ADAPTER.validate_python(incoming_values), credential_values_to_delete=credential.credential_values_to_delete, ) - decrypted_stored_credential: Final = decrypted_credential(db_credential) - decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( - decrypted_stored_credential.credential_values - ) stored_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( db_credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field ) diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index c905d38dd41..9dd36a535e4 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -2,8 +2,9 @@ import json from contextlib import contextmanager +from functools import reduce from types import SimpleNamespace -from typing import Final +from typing import Final, Protocol from unittest.mock import AsyncMock, MagicMock, call, patch import httpx @@ -12,14 +13,16 @@ from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import ec from fastapi import HTTPException, Request from fastapi.testclient import TestClient +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.litellm_logging import get_masked_values from litellm.llms.github_copilot.per_user_auth import _SESSION_CACHE, _session_cache_key from litellm.models.credentials import CredentialSource, UserConnectionFlowHandle from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper from litellm.proxy.credential_endpoints.endpoints import _decode_flow_handle, get_llm_router from litellm.proxy.proxy_server import app from litellm.types.utils import CredentialItem @@ -460,6 +463,264 @@ def test_update_credential_never_syncs_a_null_into_the_in_memory_credential(rest assert synced.credential_values["api_key"] == "sk-rotated" +_WRITTEN_JSON: Final = TypeAdapter(dict[str, object]) + + +class _InstallCredentialStore(Protocol): + def __call__( + self, + *, + connected: bool = True, + in_memory: tuple[object, ...] = (), + llm_router: object | None = None, + **repository_calls: AsyncMock, + ) -> None: ... + + +def _encrypted_row(served: CredentialItem) -> CredentialItem: + return CredentialItem( + credential_name=served.credential_name, + credential_values={key: encrypt_value_helper(value) for key, value in served.credential_values.items()}, + credential_info=dict(served.credential_info), + ) + + +def _written_values(update_by_name: AsyncMock) -> dict[str, object]: + written: Final = _WRITTEN_JSON.validate_json(update_by_name.await_args.kwargs["data"]["credential_values"]) + return { + key: decrypt_value_helper(value, key) if isinstance(value, str) else value for key, value in written.items() + } + + +def _written_info(update_by_name: AsyncMock) -> dict[str, object]: + return _WRITTEN_JSON.validate_json(update_by_name.await_args.kwargs["data"]["credential_info"]) + + +def _serve_credential( + credential_store: _InstallCredentialStore, + served: CredentialItem, + in_memory: tuple[CredentialItem, ...] | None = None, +) -> AsyncMock: + update_by_name = AsyncMock(return_value=None) + credential_store( + in_memory=(served,) if in_memory is None else in_memory, + find_by_name=AsyncMock(side_effect=lambda _name: _encrypted_row(served)), + update_by_name=update_by_name, + ) + return update_by_name + + +def _service_account( + private_key: str = "pk-1234567890", client_email: str = "svc@example.com" +) -> dict[str, dict[str, str]]: + return {"vertex_credentials": {"private_key": private_key, "client_email": client_email}} + + +def test_update_credential_keeps_the_stored_secret_when_the_patch_echoes_its_masked_read_back( + credential_store: _InstallCredentialStore, +): + served = CredentialItem( + credential_name="existing", + credential_values={"api_key": "sk-real-secret-1234"}, + credential_info={"custom_llm_provider": "openai", "description": "original"}, + ) + update_by_name = _serve_credential(credential_store, served) + + response = _patch_credential( + "existing", + { + "credential_values": {"api_key": "sk****34"}, + "credential_info": {"custom_llm_provider": "openai", "description": "edited"}, + }, + ) + + assert response.status_code == 200, response.text + assert _written_values(update_by_name) == {"api_key": "sk-real-secret-1234"} + assert _written_info(update_by_name)["description"] == "edited" + assert CredentialAccessor.get_credential_values("existing") == {"api_key": "sk-real-secret-1234"} + assert litellm.credential_list[0].credential_info["description"] == "edited" + + +def test_update_credential_keeps_the_secret_when_the_patch_echoes_what_get_by_name_served( + credential_store: _InstallCredentialStore, +): + served = CredentialItem( + credential_name="existing", + credential_values={"api_key": "sk-real-secret-1234", "api_base": "https://api.example.com"}, + credential_info={"custom_llm_provider": "openai", "description": "original"}, + ) + update_by_name = _serve_credential(credential_store, served) + read_back = _call_as("GET", "/credentials/by_name/existing").json() + assert read_back["credential_values"]["api_key"] != "sk-real-secret-1234" + + response = _patch_credential( + "existing", {**read_back, "credential_info": {**read_back["credential_info"], "description": "edited"}} + ) + + assert response.status_code == 200, response.text + assert _written_values(update_by_name) == {"api_key": "sk-real-secret-1234", "api_base": "https://api.example.com"} + assert _written_info(update_by_name)["description"] == "edited" + + +def test_update_credential_keeps_a_short_secret_whose_read_back_is_all_asterisks( + credential_store: _InstallCredentialStore, +): + served = CredentialItem(credential_name="existing", credential_values={"api_key": "abc"}, credential_info={}) + update_by_name = _serve_credential(credential_store, served) + read_back = _call_as("GET", "/credentials/by_name/existing").json() + assert read_back["credential_values"]["api_key"] != "abc" + + response = _patch_credential( + "existing", {"credential_values": read_back["credential_values"], "credential_info": {}} + ) + + assert response.status_code == 200, response.text + assert _written_values(update_by_name) == {"api_key": "abc"} + assert CredentialAccessor.get_credential_values("existing") == {"api_key": "abc"} + + +@pytest.mark.parametrize( + "rotated_values", + ( + {"api_key": "sk-rotated-key-5678"}, + {"api_key": "new**secret"}, + _service_account(private_key="pk**rotated"), + ), +) +def test_update_credential_still_rotates_a_secret_sent_in_full( + credential_store: _InstallCredentialStore, rotated_values: dict[str, object] +): + served = CredentialItem( + credential_name="existing", + credential_values={"api_key": "sk-old-key-000011", **_service_account()}, + credential_info={}, + ) + update_by_name = _serve_credential(credential_store, served) + + response = _patch_credential("existing", {"credential_values": rotated_values, "credential_info": {}}) + + assert response.status_code == 200, response.text + assert _written_values(update_by_name) == {**served.credential_values, **rotated_values} + assert CredentialAccessor.get_credential_values("existing") == {**served.credential_values, **rotated_values} + + +def test_update_credential_saves_a_mask_shaped_value_on_a_key_get_never_masks( + credential_store: _InstallCredentialStore, +): + served = CredentialItem( + credential_name="existing", + credential_values={ + "api_key": "sk-real-secret-1234", + "api_base": "https://api.example.com", + "extra_headers": {"x-team": "platform"}, + }, + credential_info={}, + ) + update_by_name = _serve_credential(credential_store, served) + sent: Final = {"api_base": "ab****cd", "organization": "*****", "extra_headers": {"x-team": "*****"}} + assert get_masked_values(sent, unmasked_length=4, number_of_asterisks=4) == sent + + response = _patch_credential("existing", {"credential_values": sent, "credential_info": {}}) + + assert response.status_code == 200, response.text + assert _written_values(update_by_name) == {"api_key": "sk-real-secret-1234", **sent} + + +@pytest.mark.parametrize("client_email", ("svc@example.com", "new@example.com")) +def test_update_credential_keeps_a_masked_secret_inside_a_stored_object( + credential_store: _InstallCredentialStore, client_email: str +): + served = CredentialItem(credential_name="existing", credential_values=_service_account(), credential_info={}) + update_by_name = _serve_credential(credential_store, served) + + response = _patch_credential( + "existing", + { + "credential_values": _service_account(private_key="pk****90", client_email=client_email), + "credential_info": {}, + }, + ) + + assert response.status_code == 200, response.text + assert _written_values(update_by_name) == _service_account(client_email=client_email) + assert CredentialAccessor.get_credential_values("existing") == _service_account(client_email=client_email) + + +def test_update_credential_drops_a_masked_leaf_the_stored_object_no_longer_has( + credential_store: _InstallCredentialStore, +): + served = CredentialItem(credential_name="existing", credential_values=_service_account(), credential_info={}) + update_by_name = _serve_credential(credential_store, served) + + response = _patch_credential( + "existing", + { + "credential_values": { + "vertex_credentials": { + "private_key": "pk****90", + "refresh_token": "rt****89", + "client_email": "svc@example.com", + } + }, + "credential_info": {}, + }, + ) + + assert response.status_code == 200, response.text + assert _written_values(update_by_name) == _service_account() + + +@pytest.mark.parametrize( + "in_memory", + ( + pytest.param( + ( + CredentialItem( + credential_name="existing", + credential_values={"api_key": "sk-old-key-000011"}, + credential_info={"description": "original"}, + ), + ), + id="replica_not_resynced", + ), + pytest.param((), id="replica_never_loaded_it"), + ), +) +def test_update_credential_keeps_the_table_secret_whatever_this_replica_last_rendered( + credential_store: _InstallCredentialStore, in_memory: tuple[CredentialItem, ...] +): + served = CredentialItem( + credential_name="existing", + credential_values={"api_key": "sk-rotated-key-000099"}, + credential_info={"description": "original"}, + ) + update_by_name = _serve_credential(credential_store, served, in_memory=in_memory) + + response = _patch_credential( + "existing", + {"credential_values": {"api_key": "sk****11"}, "credential_info": {"description": "edited"}}, + ) + + assert response.status_code == 200, response.text + assert _written_values(update_by_name) == {"api_key": "sk-rotated-key-000099"} + assert _written_info(update_by_name)["description"] == "edited" + + +@pytest.mark.parametrize("levels", (20, 21)) +def test_update_credential_keeps_a_secret_exactly_as_deep_as_the_masker_masks( + credential_store: _InstallCredentialStore, levels: int +): + nested = reduce(lambda inner, _level: {"credentials": inner}, range(levels), {"api_key": "leaf-secret-1234"}) + served = CredentialItem(credential_name="existing", credential_values=nested, credential_info={}) + update_by_name = _serve_credential(credential_store, served) + read_back = get_masked_values(served.credential_values, unmasked_length=4, number_of_asterisks=4) + + response = _patch_credential("existing", {"credential_values": read_back, "credential_info": {}}) + + assert response.status_code == 200, response.text + assert _written_values(update_by_name) == nested + + def _generate_es256_pem() -> str: key = ec.generate_private_key(ec.SECP256R1()) return key.private_bytes( From bad98c4b02f09076cd11fcf788888f36dce23e73 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 28 Oct 2026 23:08:24 -0700 Subject: [PATCH 06/42] fix(responses): call the renamed drops_params helper on the compact path (#49030) Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/openai/responses/transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 64b2e61130d..45fc1fc3ff7 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1218,7 +1218,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): def compact_dropped_params(self, litellm_params: GenericLiteLLMParams) -> frozenset[str]: if self.custom_llm_provider not in _PROVIDERS_WITH_STRICT_COMPACT_PARAMS: return frozenset() - if _drops_params(litellm_params): + if drops_params(litellm_params): return _COMPACT_TRACKING_PARAMS | _COMPACT_REJECTED_PARAMS return _COMPACT_TRACKING_PARAMS From c666df3dace0b5fa17140ea357c8d56d6e576c36 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:24:52 -0700 Subject: [PATCH 07/42] docs(qa): allow focused proof for mechanical PRs (#49031) --- .github/PULL_REQUEST_TEMPLATE/general.md | 7 ++++--- AGENTS.md | 4 +++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/.github/PULL_REQUEST_TEMPLATE/general.md b/.github/PULL_REQUEST_TEMPLATE/general.md index 243d121d626..86a23966671 100644 --- a/.github/PULL_REQUEST_TEMPLATE/general.md +++ b/.github/PULL_REQUEST_TEMPLATE/general.md @@ -84,9 +84,10 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac ## Screenshots / Proof of Fix