From 9e0659212a02dccfdaf74711bffb75bef2a4fda0 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 4 Sep 2026 13:48:42 -0700 Subject: [PATCH 1/8] test(e2e): repair two suites broken by intentional behaviour changes Both of these are e2e assumptions that PRs #31731 and #39532 invalidated, not product regressions. They have been red in litellm-e2e builds 119-123. Wildcard readiness probe (6 errors in test_model_access_group_e2e.py) #31731 made _get_wildcard_models drop a wildcard route from /v1/models unconditionally; before it, a wildcard with a matching router deployment stayed in the list and only the no-router / no-deployment fallbacks removed it. The shared readiness helper polls /v1/models for an exact id match, so registering openai/gpt-5.4* now times out at model_servable_timeout every run and every test in the class errors in setup. return_wildcard_routes=True still re-adds the route, so the poll asks for it. The flag is a no-op for a concrete model name -- it only ever adds wildcard entries -- so it is set unconditionally rather than sniffing the name. Semantic auto-router spend assertion #39532 bills the routing embedding to the caller's key on purpose, so the key's spend logs now legitimately carry an openai/text-embedding-3-small row and _assert_served_only_by rejects it. Widening the allowlist would have weakened the assertion this test exists for -- that the request reached the target deployment. Instead the embedding row is split off and asserted separately, which turns the break into coverage for #39532. The poll gains a predicate so it waits for the embedding row rather than racing whichever row is written first. --- tests/e2e/models.py | 9 +++++++++ tests/e2e/proxy_client.py | 3 ++- .../router/test_auto_router_regressions_e2e.py | 15 +++++++++++++-- 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 79d9e011f7e..b5229744d6f 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -851,6 +851,15 @@ class ModelListEntry(BaseModel): id: str +class ModelsListParams(BaseModel): + """Query for GET /v1/models. A wildcard route such as ``openai/gpt-5.4*`` is + listed only under ``return_wildcard_routes``; without it the route is dropped + and only its expansions remain, so a readiness poll for the pattern itself + never resolves.""" + + return_wildcard_routes: bool = True + + class ModelsListResponse(BaseModel): """GET /v1/models on the data plane: the deployments the gateway can actually serve right now. Used to confirm a freshly created model has propagated from diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 2d382a610e1..cdc20e5299a 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -55,6 +55,7 @@ from models import ( ModelMode, ModelNewBody, ModelNewResponse, + ModelsListParams, ModelsListResponse, ModelUpdateBody, OcrBody, @@ -336,7 +337,7 @@ class ProxyClient: lambda poll_timeout: self.transport.get( "/v1/models", headers=headers, - params=NoBody(), + params=ModelsListParams(), response_type=ModelsListResponse, timeout=poll_timeout, ), diff --git a/tests/e2e/router/test_auto_router_regressions_e2e.py b/tests/e2e/router/test_auto_router_regressions_e2e.py index 35ba2c8d3d1..188db2a8eb5 100644 --- a/tests/e2e/router/test_auto_router_regressions_e2e.py +++ b/tests/e2e/router/test_auto_router_regressions_e2e.py @@ -597,9 +597,20 @@ class TestSemanticAutoRouterResponses: ) ) assert answer.id, "/v1/responses through the semantic auto-router returned no response id" - rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + rows: Final = proxy.poll_logs_for_key( + key, + min_rows=2, + predicate=lambda logged: any(row.model == EMBEDDING_MODEL for row in logged), + ) + embedding_rows: Final = tuple(row for row in rows if row.model == EMBEDDING_MODEL) + assert embedding_rows, ( + "the routing embedding was not billed to the caller's key; " + f"spend logs show {tuple(row.model for row in rows)}" + ) _assert_served_only_by( - rows, CHEAP_SERVED | {semantic_auto_router.target}, "semantic auto-router /v1/responses string input" + [row for row in rows if row.model != EMBEDDING_MODEL], + CHEAP_SERVED | {semantic_auto_router.target}, + "semantic auto-router /v1/responses string input", ) From 88985d00e2a1de43c616893141934cd0f444ce04 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 09:58:44 -0700 Subject: [PATCH 2/8] bump: litellm-enterprise 0.1.64 -> 0.1.65, litellm-proxy-extras 0.4.93 -> 0.4.94 --- enterprise/pyproject.toml | 4 ++-- litellm-proxy-extras/pyproject.toml | 4 ++-- pyproject.toml | 4 ++-- uv.lock | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index b6f482ccd86..3699087dbfa 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.64" +version = "0.1.65" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.64" +version = "0.1.65" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 97e9eb66bf2..82d31fec373 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.93" +version = "0.4.94" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.93" +version = "0.4.94" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/pyproject.toml b/pyproject.toml index c1fde4af3f5..b889a3a0e60 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,8 +67,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.93", - "litellm-enterprise==0.1.64", + "litellm-proxy-extras==0.4.94", + "litellm-enterprise==0.1.65", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", diff --git a/uv.lock b/uv.lock index 9e6343ff375..89205cd9527 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-01T21:00:02.682921Z" +exclude-newer = "2026-09-02T16:58:34.594994Z" exclude-newer-span = "P3D" [manifest] @@ -4771,12 +4771,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.64" +version = "0.1.65" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.93" +version = "0.4.94" source = { editable = "litellm-proxy-extras" } [[package]] From 4df284e16dcc02ceabad4576df6f2c976f20d839 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:39:24 -0700 Subject: [PATCH 3/8] fix(guardrails): record guardrail information for undecorated custom apply_guardrail overrides (#39727) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 3 + litellm/integrations/custom_guardrail.py | 9 + .../integrations/test_custom_guardrail.py | 168 ++++++++++++++++++ .../test_openai_guardrail_handler.py | 31 ++++ .../test_bedrock_guardrails.py | 6 +- 5 files changed, 216 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 7d6de612349..e4fb2a00297 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -215,6 +215,9 @@ MAX_CALLBACKS: Final = get_env_int("LITELLM_MAX_CALLBACKS", 100) # so the deployment-level hook does not re-run them for the same request PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails" +# Attribute stamped on log_guardrail_information wrappers so __init_subclass__ does not wrap them again +LOGS_GUARDRAIL_INFORMATION_MARKER: Final = "_litellm_logs_guardrail_information" + # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9bb613654e4..c462b1edb98 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -46,6 +46,7 @@ dc: Final = DualCache() from litellm.constants import ( GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + LOGS_GUARDRAIL_INFORMATION_MARKER, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) from litellm.exceptions import ( @@ -151,6 +152,13 @@ class CustomGuardrail(CustomLogger): records_own_guardrail_information: ClassVar[bool] = False + def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: forwarded to cooperative __init_subclass__ hooks + super().__init_subclass__(**kwargs) + own_apply_guardrail: Final = cls.__dict__.get("apply_guardrail") + if own_apply_guardrail is None or LOGS_GUARDRAIL_INFORMATION_MARKER in vars(own_apply_guardrail): + return + cls.apply_guardrail = log_guardrail_information(own_apply_guardrail) + def __init__( self, guardrail_name: str | None = None, @@ -1559,4 +1567,5 @@ def log_guardrail_information(func): return async_wrapper(*args, **kwargs) return sync_wrapper(*args, **kwargs) + vars(wrapper)[LOGS_GUARDRAIL_INFORMATION_MARKER] = True # rebind-ok: stamps the wrapper this call just built return wrapper diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 6edc2c9bf77..49a52157c8a 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,4 +1,5 @@ import asyncio +from typing import TYPE_CHECKING, Literal, Optional from unittest.mock import AsyncMock import pytest @@ -11,6 +12,9 @@ from litellm.integrations.custom_guardrail import ( from litellm.proxy._types import CallTypes, UserAPIKeyAuth from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class TestCustomGuardrailDeploymentHook: @@ -2239,6 +2243,170 @@ class TestRecordsOwnGuardrailInformation: assert _guardrail_entries(request_data) == [] +class _UndecoratedGuardrail(CustomGuardrail): + """apply_guardrail written like the docs example: no @log_guardrail_information.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + from litellm.exceptions import GuardrailRaisedException + + if any("forbidden" in text for text in inputs.get("texts") or []): + raise GuardrailRaisedException(guardrail_name=self.guardrail_name, message="Content blocked") + return inputs + + +class _UndecoratedSelfRecordingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"custom": True}, + request_data=request_data, + guardrail_status="success", + start_time=0.0, + end_time=0.0, + duration=0.0, + ) + return inputs + + +class _InheritedApplyGuardrail(_UndecoratedGuardrail): + pass + + +class TestUndecoratedApplyGuardrailIsLogged: + """LIT-5983 regression: a custom guardrail that overrides apply_guardrail without the + @log_guardrail_information decorator must still record guardrail information, and the + auto-wrap must not double-record decorated or self-recording implementations.""" + + @pytest.mark.asyncio + async def test_undecorated_success_is_recorded(self): + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = _UndecoratedGuardrail(guardrail_name="docs-style", event_hook=GuardrailEventHooks.pre_call) + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "docs-style" + assert entries[0]["guardrail_mode"] == "pre_call" + assert entries[0]["guardrail_status"] == "success" + + @pytest.mark.asyncio + async def test_undecorated_block_is_recorded_and_reraised(self): + from litellm.exceptions import GuardrailRaisedException + + guardrail = _UndecoratedGuardrail(guardrail_name="docs-style") + request_data: dict = {"model": "gpt-4o"} + + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["forbidden"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "docs-style" + assert entries[0]["guardrail_status"] == "guardrail_intervened" + + @pytest.mark.asyncio + async def test_undecorated_bare_exception_is_recorded_as_failed_to_respond(self): + class _BareExceptionGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + raise Exception("Content blocked: policy violation") + + guardrail = _BareExceptionGuardrail(guardrail_name="docs-style") + request_data: dict = {"model": "gpt-4o"} + + with pytest.raises(Exception, match="Content blocked"): + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["x"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_status"] == "guardrail_failed_to_respond" + + @pytest.mark.asyncio + async def test_inherited_apply_guardrail_is_recorded_once(self): + guardrail = _InheritedApplyGuardrail(guardrail_name="child") + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + assert len(_guardrail_entries(request_data)) == 1 + + @pytest.mark.asyncio + async def test_undecorated_self_recording_apply_guardrail_is_recorded_once(self): + guardrail = _UndecoratedSelfRecordingGuardrail(guardrail_name="self-recording") + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_response"] == {"custom": True} + + @pytest.mark.asyncio + async def test_base_apply_guardrail_is_not_recorded(self): + guardrail = CustomGuardrail(guardrail_name="base") + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + assert _guardrail_entries(request_data) == [] + + def test_subclass_keywords_reach_cooperative_init_subclass(self): + class _LabelMixin: + seen_label: str = "" + + def __init_subclass__(cls, label: str = "", **kwargs: object) -> None: + super().__init_subclass__(**kwargs) + cls.seen_label = label + + class _Labelled(CustomGuardrail, _LabelMixin, label="docs-style"): + pass + + assert _Labelled.seen_label == "docs-style" + + class _ApplyOnlyObserver(CustomGuardrail): """Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook.""" diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index cebab2512d0..36e715d5804 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1075,6 +1075,37 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert result == responses_so_far +class TestUndecoratedGuardrailIsRecorded: + """LIT-5983 regression: the handler calls apply_guardrail bare, so a custom guardrail + without @log_guardrail_information must still end up in the request's guardrail + information on both the request and response paths.""" + + @pytest.mark.asyncio + async def test_request_path_records_undecorated_guardrail(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="docs-style") + data = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + + await handler.process_input_messages(data, guardrail) + + entries = data["metadata"]["standard_logging_guardrail_information"] + assert [(e["guardrail_name"], e["guardrail_status"]) for e in entries] == [("docs-style", "success")] + + @pytest.mark.asyncio + async def test_response_path_records_undecorated_guardrail(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="docs-style") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="hi", role="assistant"))] + ) + request_data: dict = {"metadata": {}} + + await handler.process_output_response(response, guardrail, request_data=request_data) + + entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert [(e["guardrail_name"], e["guardrail_status"]) for e in entries] == [("docs-style", "success")] + + class TestGetStructuredMessages: """Test the get_structured_messages method.""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 9842d88e8d1..7da55f22bda 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -1137,7 +1137,11 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): mock_api.assert_called_once() kwargs = mock_api.call_args.kwargs assert kwargs["source"] == "OUTPUT" - assert kwargs["request_data"] == {"model": "gpt-4o"} + assert kwargs["request_data"]["model"] == "gpt-4o" + recorded = kwargs["request_data"]["metadata"]["standard_logging_guardrail_information"] + assert [(e["guardrail_name"], e["guardrail_status"]) for e in recorded] == [ + (guardrail.guardrail_name, "success") + ] synthetic = kwargs["response"] assert isinstance(synthetic, ModelResponse) assert len(synthetic.choices) == 2 From f66b3ebe0dda8e25c47739c1eb15637dce2d11ad Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:40:00 -0700 Subject: [PATCH 4/8] feat(responses): honor supported_endpoints /v1/responses opt-in for OpenAI-compatible deployments (#39725) * feat(responses): honor supported_endpoints /v1/responses opt-in for OpenAI-compatible deployments custom_openai and other generic OpenAI-compatible deployments have no native Responses API config, so every /v1/responses call is bridged through /v1/chat/completions. When model_info.supported_endpoints lists /v1/responses, resolve OpenAILikeResponsesConfig instead so the request is forwarded to {api_base}/responses, for streaming, non-streaming and mode: responses deployments alike. Providers with their own Responses config are unchanged. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): drop deployment supported_endpoints opt-in after cross-provider prompt swap A prompt manager that moves the request to another provider leaves kwargs['model_info'] describing the original deployment; without this the swapped provider was sent an OpenAI-like /responses request it does not serve. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(responses): carry prompt-swap deployment metadata as a return value instead of a kwargs marker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/main.py | 93 +++++-- ...sponses_supported_endpoints_passthrough.py | 254 ++++++++++++++++++ 2 files changed, 327 insertions(+), 20 deletions(-) create mode 100644 tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py diff --git a/litellm/responses/main.py b/litellm/responses/main.py index ed2d6a216fd..5e74b7324b4 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -2,6 +2,7 @@ import asyncio import contextvars from collections.abc import Coroutine, Generator, Iterable, Mapping from contextlib import contextmanager +from dataclasses import dataclass from functools import partial from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast @@ -23,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) @@ -403,8 +405,40 @@ def _bridges_to_chat_completions( return responses_api_provider_config is None or use_chat_completions_api is True +def _deployment_passes_through_responses(model_info: object) -> bool: + """Whether ``model_info.supported_endpoints`` opts the deployment into native ``{api_base}/responses``.""" + if not isinstance(model_info, dict): + return False + supported_endpoints: Final = model_info.get("supported_endpoints") + return isinstance(supported_endpoints, (list, tuple)) and "/v1/responses" in supported_endpoints + + +def _deployment_model_info_after_prompt_swap( + requested_provider: str | None, resolved_provider: str | None, model_info: object +) -> object: + """Deployment metadata only describes the upstream while the prompt manager keeps its provider.""" + return model_info if resolved_provider == requested_provider else None + + +@dataclass(frozen=True, slots=True) +class _AsyncPromptManagementOutcome: + merged_optional_params: Mapping[str, object] + deployment_model_info: object + + +def _resolve_responses_api_provider_config( + model: str, custom_llm_provider: str, model_info: object +) -> BaseResponsesAPIConfig | None: + provider_config: Final = ProviderConfigManager.get_provider_responses_api_config( + model=model, provider=custom_llm_provider + ) + if provider_config is not None or not _deployment_passes_through_responses(model_info): + return provider_config + return OpenAILikeResponsesConfig() + + def _will_bridge_to_chat_completions( - model: str, custom_llm_provider: str | None, use_chat_completions_api: bool + model: str, custom_llm_provider: str | None, use_chat_completions_api: bool, model_info: object ) -> bool: """``_bridges_to_chat_completions`` for callers running before the provider config is resolved. @@ -418,9 +452,7 @@ def _will_bridge_to_chat_completions( if custom_llm_provider is None: return True return _bridges_to_chat_completions( - ProviderConfigManager.get_provider_responses_api_config( - model=normalized_model[0], provider=custom_llm_provider - ), + _resolve_responses_api_provider_config(normalized_model[0], custom_llm_provider, model_info), use_chat_completions_api or normalized_model[1], ) @@ -527,7 +559,10 @@ async def aresponses( with _prompt_management_sees_a_provisional_message_list( kwargs, bridged=_will_bridge_to_chat_completions( - model, custom_llm_provider, bool(kwargs.get("use_chat_completions_api")) + model, + custom_llm_provider, + bool(kwargs.get("use_chat_completions_api")), + kwargs.get("model_info"), ), ): ( @@ -552,6 +587,7 @@ async def aresponses( merged_input=merged_input, ), ) + requested_provider: Final = custom_llm_provider if model != original_model: custom_llm_provider = _resolve_prompt_swapped_provider( original_model=original_model, @@ -561,7 +597,12 @@ async def aresponses( prompt_id=prompt_id, ) kwargs.pop("prompt_id", None) - kwargs["_async_prompt_merged_params"] = merged_optional_params + kwargs["_async_prompt_merged_params"] = _AsyncPromptManagementOutcome( + merged_optional_params=merged_optional_params, + deployment_model_info=_deployment_model_info_after_prompt_swap( + requested_provider, custom_llm_provider, kwargs.get("model_info") + ), + ) func: Final = partial( responses, @@ -666,12 +707,14 @@ def _apply_prompt_management_to_responses_call( kwargs: dict[str, Any], local_vars: dict[str, object], use_chat_completions_api: bool, -) -> tuple[str | ResponseInputParam, str, str | None]: - async_merged: Final[Mapping[str, object] | None] = kwargs.pop("_async_prompt_merged_params", None) - if async_merged is not None: - for key, value in async_merged.items(): +) -> tuple[str | ResponseInputParam, str, str | None, object]: + """Returns the prompt-managed input, model and provider, plus the deployment metadata that still + describes the upstream (``None`` once the prompt manager moved the request to another provider).""" + async_outcome: Final[_AsyncPromptManagementOutcome | None] = kwargs.pop("_async_prompt_merged_params", None) + if async_outcome is not None: + for key, value in async_outcome.merged_optional_params.items(): local_vars[key] = value - return input, model, custom_llm_provider + return input, model, custom_llm_provider, async_outcome.deployment_model_info prompt_id: Final = cast(str | None, kwargs.get("prompt_id", None)) prompt_variables: Final = cast(dict | None, kwargs.get("prompt_variables", None)) @@ -684,7 +727,9 @@ def _apply_prompt_management_to_responses_call( ): with _prompt_management_sees_a_provisional_message_list( kwargs, - bridged=_will_bridge_to_chat_completions(model, custom_llm_provider, use_chat_completions_api), + bridged=_will_bridge_to_chat_completions( + model, custom_llm_provider, use_chat_completions_api, kwargs.get("model_info") + ), ): ( model, @@ -710,19 +755,28 @@ def _apply_prompt_management_to_responses_call( ) local_vars["input"] = input local_vars["model"] = model - if model != original_model: - custom_llm_provider = _resolve_prompt_swapped_provider( + resolved_provider: Final = ( + custom_llm_provider + if model == original_model + else _resolve_prompt_swapped_provider( original_model=original_model, swapped_model=model, custom_llm_provider=custom_llm_provider, kwargs=kwargs, prompt_id=prompt_id, ) - local_vars["custom_llm_provider"] = custom_llm_provider + ) + local_vars["custom_llm_provider"] = resolved_provider for key, value in merged_optional_params.items(): local_vars[key] = value + return ( + input, + model, + resolved_provider, + _deployment_model_info_after_prompt_swap(custom_llm_provider, resolved_provider, kwargs.get("model_info")), + ) - return input, model, custom_llm_provider + return input, model, custom_llm_provider, kwargs.get("model_info") # Opt-in via model id (mirrors the `responses/` prefix pattern on chat completions). @@ -1052,7 +1106,7 @@ def responses( ) local_vars["custom_llm_provider"] = custom_llm_provider - input, model, custom_llm_provider = _apply_prompt_management_to_responses_call( + input, model, custom_llm_provider, deployment_model_info = _apply_prompt_management_to_responses_call( input=input, model=model, custom_llm_provider=custom_llm_provider, @@ -1123,9 +1177,8 @@ def responses( if custom_llm_provider is None: responses_api_provider_config = None else: - responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config( - model=model, - provider=custom_llm_provider, + responses_api_provider_config = _resolve_responses_api_provider_config( + model, custom_llm_provider, deployment_model_info ) local_vars.update(kwargs) diff --git a/tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py b/tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py new file mode 100644 index 00000000000..7cd04b015f9 --- /dev/null +++ b/tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py @@ -0,0 +1,254 @@ +""" +A deployment with `model_info.supported_endpoints` containing `/v1/responses` forwards +`/v1/responses` natively to `{api_base}/responses`. Without it, generic OpenAI-compatible +providers such as `custom_openai` keep bridging through `/v1/chat/completions`. +""" + +import json +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +import respx + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig +from litellm.responses.main import _resolve_responses_api_provider_config +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ModelResponse + +API_BASE = "https://backend.example/v1" +RESPONSES_URL = f"{API_BASE}/responses" +CHAT_URL = f"{API_BASE}/chat/completions" +OPT_IN = {"supported_endpoints": ["/v1/chat/completions", "/v1/responses"]} + +RESPONSES_BODY = { + "id": "resp_native", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "my-model", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "native", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} + +CHAT_BODY = { + "id": "chatcmpl_bridged", + "object": "chat.completion", + "created": 1741476542, + "model": "my-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "bridged"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + +SSE_BODY = ( + "event: response.created\n" + f"data: {json.dumps({'type': 'response.created', 'response': RESPONSES_BODY})}\n\n" + "event: response.completed\n" + f"data: {json.dumps({'type': 'response.completed', 'response': RESPONSES_BODY})}\n\n" +) + + +def _mock_backend(router: respx.MockRouter) -> tuple[respx.Route, respx.Route]: + responses_route = router.post(RESPONSES_URL).mock(return_value=httpx.Response(200, json=RESPONSES_BODY)) + chat_route = router.post(CHAT_URL).mock(return_value=httpx.Response(200, json=CHAT_BODY)) + return responses_route, chat_route + + +SWAPPED_MODEL = "deepseek/deepseek-chat" +SWAPPED_API_BASE = "https://api.deepseek.com/beta" + + +def _prompt_manager_swapping_to(model: str) -> MagicMock: + """A logging object whose prompt hook rewrites the request's model, as a prompt manager does.""" + prompt_return = (model, [{"role": "user", "content": "hi"}], {}) + logging_obj = MagicMock() + logging_obj.__class__ = LiteLLMLoggingObj + logging_obj.should_run_prompt_management_hooks.return_value = True + logging_obj.get_chat_completion_prompt.return_value = prompt_return + logging_obj.async_get_chat_completion_prompt = AsyncMock(return_value=prompt_return) + logging_obj.model_call_details = {} + return logging_obj + + +def _mock_swap_targets(router: respx.MockRouter, monkeypatch) -> tuple[respx.Route, respx.Route]: + """The swapped provider's chat endpoint, plus the `/responses` it does not serve but a stale + opt-in would send to.""" + monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-deepseek") + swapped_chat_route = router.post(f"{SWAPPED_API_BASE}/chat/completions").mock( + return_value=httpx.Response(200, json=CHAT_BODY) + ) + stale_responses_route = router.post(f"{SWAPPED_API_BASE}/responses").mock( + return_value=httpx.Response(200, json=RESPONSES_BODY) + ) + return swapped_chat_route, stale_responses_route + + +@pytest.fixture(autouse=True) +def _respx_interceptable_httpx_client(monkeypatch): + monkeypatch.setattr(litellm, "num_retries", 0) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.parametrize( + "model_info, expected_type", + [ + (OPT_IN, OpenAILikeResponsesConfig), + ({"supported_endpoints": ["/v1/chat/completions"]}, type(None)), + ({}, type(None)), + (None, type(None)), + ("/v1/responses", type(None)), + ], +) +def test_resolver_opt_in_gates_openai_like_config(model_info, expected_type): + config = _resolve_responses_api_provider_config("my-model", "custom_openai", model_info) + assert type(config) is expected_type + + +def test_resolver_keeps_native_provider_config(): + """`openai/` already routes /v1/responses natively; the opt-in must not swap its config.""" + config = _resolve_responses_api_provider_config("gpt-4.1", "openai", OPT_IN) + assert type(config) is OpenAIResponsesAPIConfig + + +@respx.mock +async def test_opt_in_forwards_responses_natively(): + responses_route, chat_route = _mock_backend(respx.mock) + + result = await litellm.aresponses( + model="custom_openai/my-model", + input="hi", + api_base=API_BASE, + api_key="sk-backend", + model_info=OPT_IN, + ) + + assert responses_route.call_count == 1 + assert chat_route.call_count == 0 + request = responses_route.calls.last.request + assert request.headers["authorization"] == "Bearer sk-backend" + assert json.loads(request.content)["input"] == "hi" + assert isinstance(result, ResponsesAPIResponse) + assert result.output[0].content[0].text == "native" + + +@respx.mock +async def test_opt_in_forwards_streaming_responses_natively(monkeypatch): + """The router registers each deployment in `litellm.model_cost`; an unregistered model is + treated as non-streaming and would be faked, so mirror that registration here.""" + monkeypatch.setitem(litellm.model_cost, "custom_openai/my-model", {"litellm_provider": "custom_openai"}) + responses_route = respx.post(RESPONSES_URL).mock( + return_value=httpx.Response(200, text=SSE_BODY, headers={"content-type": "text/event-stream"}) + ) + chat_route = respx.post(CHAT_URL).mock(return_value=httpx.Response(200, json=CHAT_BODY)) + + stream = await litellm.aresponses( + model="custom_openai/my-model", + input="hi", + stream=True, + api_base=API_BASE, + api_key="sk-backend", + model_info=OPT_IN, + ) + events = [event async for event in stream] + + assert responses_route.call_count == 1 + assert chat_route.call_count == 0 + assert json.loads(responses_route.calls.last.request.content)["stream"] is True + assert [event.type for event in events] == ["response.created", "response.completed"] + + +@respx.mock +async def test_without_opt_in_still_bridges_through_chat_completions(): + responses_route, chat_route = _mock_backend(respx.mock) + + result = await litellm.aresponses( + model="custom_openai/my-model", + input="hi", + api_base=API_BASE, + api_key="sk-backend", + model_info={"supported_endpoints": ["/v1/chat/completions"]}, + ) + + assert chat_route.call_count == 1 + assert responses_route.call_count == 0 + assert isinstance(result, ResponsesAPIResponse) + assert result.output[0].content[0].text == "bridged" + + +@respx.mock +async def test_prompt_swap_to_other_provider_drops_deployment_opt_in(monkeypatch): + """When a prompt manager moves the request to another provider, the original deployment's + `supported_endpoints` no longer describes the upstream, so the swapped provider bridges.""" + swapped_chat_route, stale_responses_route = _mock_swap_targets(respx.mock, monkeypatch) + + result = await litellm.aresponses( + model="custom_openai/my-model", + input="hi", + prompt_id="p1", + litellm_logging_obj=_prompt_manager_swapping_to(SWAPPED_MODEL), + model_info=OPT_IN, + ) + + assert swapped_chat_route.call_count == 1 + assert stale_responses_route.call_count == 0 + assert isinstance(result, ResponsesAPIResponse) + assert result.output[0].content[0].text == "bridged" + + +@respx.mock +def test_sync_prompt_swap_to_other_provider_drops_deployment_opt_in(monkeypatch): + swapped_chat_route, stale_responses_route = _mock_swap_targets(respx.mock, monkeypatch) + + result = litellm.responses( + model="custom_openai/my-model", + input="hi", + prompt_id="p1", + litellm_logging_obj=_prompt_manager_swapping_to(SWAPPED_MODEL), + model_info=OPT_IN, + ) + + assert swapped_chat_route.call_count == 1 + assert stale_responses_route.call_count == 0 + assert isinstance(result, ResponsesAPIResponse) + assert result.output[0].content[0].text == "bridged" + + +@respx.mock +async def test_mode_responses_chat_completion_reaches_native_responses(monkeypatch): + """A `mode: responses` deployment bridges chat completions into the Responses API; with + the opt-in that inner call must reach `{api_base}/responses` instead of bouncing back + to `/chat/completions`.""" + responses_route, chat_route = _mock_backend(respx.mock) + monkeypatch.setitem( + litellm.model_cost, + "custom_openai/my-model", + {"mode": "responses", "litellm_provider": "custom_openai"}, + ) + + result = await litellm.acompletion( + model="custom_openai/my-model", + messages=[{"role": "user", "content": "hi"}], + api_base=API_BASE, + api_key="sk-backend", + model_info={"mode": "responses", **OPT_IN}, + ) + + assert responses_route.call_count == 1 + assert chat_route.call_count == 0 + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "native" From e6705510f82c0c70b274c922210bbdd8edab5379 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:41:35 -0700 Subject: [PATCH 5/8] fix(router): hold max_parallel_requests slot until streaming response is exhausted or closed (#39859) * fix(router): hold max_parallel_requests slot until streaming response is exhausted or closed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(router): normalize deployment_slot once to keep stream_with_fallbacks under the C901 ceiling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): close upstream stream before releasing max_parallel_requests slot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 108 ++++++++++++----------- tests/test_litellm/test_router.py | 141 ++++++++++++++++++++++++++++++ 2 files changed, 198 insertions(+), 51 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 6943eece90f..490836f5f0a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2597,14 +2597,20 @@ class Router: model_response: CustomStreamWrapper, messages: list[dict[str, str]], initial_kwargs: dict, + deployment_slot: contextlib.AsyncExitStack | None = None, ) -> CustomStreamWrapper: """ Helper to iterate over a streaming response. Catches errors for fallbacks using the router's fallback system + + `deployment_slot` holds the deployment's max_parallel_requests semaphore; it is + released when the stream is exhausted, closed, or falls back to another deployment """ from litellm.exceptions import MidStreamFallbackError + held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack() + class FallbackStreamWrapper(CustomStreamWrapper): def __init__(self, async_generator: AsyncGenerator): # Copy attributes from the original model_response @@ -2628,12 +2634,26 @@ class Router: async def __anext__(self): return await self._async_generator.__anext__() + async def close_model_response() -> None: + if not hasattr(model_response, "aclose"): + return + try: + await model_response.aclose() + except BaseException as e: + verbose_router_logger.debug( + "stream_with_fallbacks: error closing model_response: %s", + e, + ) + async def stream_with_fallbacks(): fallback_response = None # Track for cleanup in finally try: async for item in model_response: yield item except MidStreamFallbackError as e: + with anyio.CancelScope(shield=True): + await close_model_response() + await held_slot.aclose() if not e.is_pre_first_chunk and ( e.generated_content or _stream_chunks_have_generated_content(model_response.chunks) ): @@ -2707,14 +2727,8 @@ class Router: # (e.g. on client disconnect). # Shield from anyio cancellation so the awaits can complete. with anyio.CancelScope(shield=True): - if hasattr(model_response, "aclose"): - try: - await model_response.aclose() - except BaseException as e: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing model_response: %s", - e, - ) + await close_model_response() + await held_slot.aclose() if fallback_response is not None and hasattr(fallback_response, "aclose"): try: await fallback_response.aclose() @@ -3379,61 +3393,53 @@ class Router: kwargs=kwargs, client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): - async with rpm_semaphore: - """ - - Check rpm limits before making the call - - If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe) - """ - await self.async_routing_strategy_pre_call_checks( - deployment=deployment, - logging_obj=logging_obj, - parent_otel_span=parent_otel_span, - ) - response = await _response - else: + async with contextlib.AsyncExitStack() as deployment_slot: + if isinstance(rpm_semaphore, asyncio.Semaphore): + await deployment_slot.enter_async_context(rpm_semaphore) await self.async_routing_strategy_pre_call_checks( deployment=deployment, logging_obj=logging_obj, parent_otel_span=parent_otel_span, ) - response = await _response - ## CHECK CONTENT FILTER ERROR ## - if isinstance(response, ModelResponse): - _should_raise = self._should_raise_content_policy_error(model=model, response=response, kwargs=kwargs) - if _should_raise: - raise litellm.ContentPolicyViolationError( - message="Response output was blocked.", - model=model, - llm_provider="", + ## CHECK CONTENT FILTER ERROR ## + if isinstance(response, ModelResponse): + _should_raise = self._should_raise_content_policy_error( + model=model, response=response, kwargs=kwargs ) + if _should_raise: + raise litellm.ContentPolicyViolationError( + message="Response output was blocked.", + model=model, + llm_provider="", + ) - if ( - isinstance(response, CustomStreamWrapper) - and response.completion_stream is None - and response.make_call is not None - ): - await response.fetch_stream() + if ( + isinstance(response, CustomStreamWrapper) + and response.completion_stream is None + and response.make_call is not None + ): + await response.fetch_stream() - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - # debug how often this deployment picked - self._track_deployment_metrics( - deployment=deployment, - response=response, - parent_otel_span=parent_otel_span, - ) - - if isinstance(response, CustomStreamWrapper): - return await self._acompletion_streaming_iterator( - model_response=response, - messages=messages, - initial_kwargs=input_kwargs_for_streaming_fallback, + self.success_calls[model_name] += 1 + verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) + # debug how often this deployment picked + self._track_deployment_metrics( + deployment=deployment, + response=response, + parent_otel_span=parent_otel_span, ) - return response + if isinstance(response, CustomStreamWrapper): + return await self._acompletion_streaming_iterator( + model_response=response, + messages=messages, + initial_kwargs=input_kwargs_for_streaming_fallback, + deployment_slot=deployment_slot.pop_all(), + ) + + return response except litellm.Timeout as e: deployment_request_timeout_param: Final = _timeout_debug_deployment_dict.get("litellm_params", {}).get( "request_timeout", None diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 31eb46f1458..ffd6c5f97ce 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12938,3 +12938,144 @@ async def test_router_retry_policy_controls_upstream_attempt_count( await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) assert upstream.call_count == expected_upstream_calls + + +class _InFlightTracker: + def __init__(self) -> None: + self.current = 0 + self.peak = 0 + + def enter(self) -> None: + self.current += 1 + self.peak = max(self.peak, self.current) + + def exit(self) -> None: + self.current -= 1 + + +_SSE_CHUNKS: Final[tuple[bytes, ...]] = tuple( + b'data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"gpt-5.6",' + b'"choices":[{"index":0,"delta":{"content":"x"},"finish_reason":null}]}\n\n' + for _ in range(5) +) + + +class _CountingSSEStream(httpx.AsyncByteStream): + def __init__(self, tracker: _InFlightTracker) -> None: + self._tracker = tracker + self._in_flight = False + + def _finish(self) -> None: + if self._in_flight: + self._in_flight = False + self._tracker.exit() + + async def __aiter__(self): + self._in_flight = True + self._tracker.enter() + try: + for chunk in _SSE_CHUNKS: + await asyncio.sleep(0.02) + yield chunk + finally: + await self.aclose() + yield b"data: [DONE]\n\n" + + async def aclose(self) -> None: + await asyncio.sleep(0.02) + self._finish() + + +def _max_parallel_router(max_parallel_requests: int) -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": "https://max-parallel.local/v1", + "max_parallel_requests": max_parallel_requests, + }, + } + ], + num_retries=0, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_router_max_parallel_requests_bounds_in_flight_upstream_calls( + monkeypatch: pytest.MonkeyPatch, stream: bool +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + tracker: Final = _InFlightTracker() + router: Final = _max_parallel_router(max_parallel_requests=2) + + async def upstream(request: httpx.Request) -> httpx.Response: + if stream: + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CountingSSEStream(tracker) + ) + tracker.enter() + await asyncio.sleep(0.05) + tracker.exit() + return httpx.Response( + 200, + json={ + "id": "c", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}], + }, + ) + + async def one_call() -> None: + response = await router.acompletion( + model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=stream + ) + if stream: + async for _ in response: + pass + + with respx.mock(assert_all_called=True) as respx_mock: + respx_mock.post("https://max-parallel.local/v1/chat/completions").mock(side_effect=upstream) + await asyncio.wait_for(asyncio.gather(*(one_call() for _ in range(10))), timeout=10) + + assert tracker.peak <= 2 + assert tracker.current == 0 + + +@pytest.mark.asyncio +async def test_router_max_parallel_requests_slot_released_when_stream_closed_early(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + tracker: Final = _InFlightTracker() + router: Final = _max_parallel_router(max_parallel_requests=1) + + with respx.mock() as respx_mock: + respx_mock.post("https://max-parallel.local/v1/chat/completions").mock( + side_effect=lambda request: httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CountingSSEStream(tracker) + ) + ) + first: Final = await router.acompletion( + model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=True + ) + await first.__anext__() + + async def second_call() -> None: + second = await router.acompletion( + model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=True + ) + async for _ in second: + pass + + second_task: Final = asyncio.create_task(second_call()) + await asyncio.sleep(0.05) + assert tracker.current == 1 + await first.aclose() + await asyncio.wait_for(second_task, timeout=2) + + assert tracker.peak == 1 + assert tracker.current == 0 From cba3dd58287114588dad0624cd2c8d0d040b902d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:41:56 -0700 Subject: [PATCH 6/8] fix(proxy): retry deadlocks and requeue spend logs on any DB write error (#39883) * fix(proxy): retry deadlocks and requeue spend logs on any DB write error update_spend_logs dequeued the batch and only retried/requeued on transport errors. A 40P01 deadlock surfaced as a plain prisma DataError and went through poison-row isolation, which dropped every row it hit; every other DB error was re-raised with the batch already gone from the queue. Treat deadlocks as transient (retry, then requeue), keep them out of poison-row isolation, and requeue the batch at the head of the queue on any other prisma error so it lands once the DB is healthy. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): drop redundant docstrings and tighten test typing for spend-log requeue Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert deadlock retries from mock call history instead of mutable lists Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/exception_handler.py | 6 + litellm/proxy/utils.py | 17 ++- .../test_proxy_update_spend.py | 109 +++++++++++++++++- 3 files changed, 128 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 19bddee618b..f469587ab8e 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -199,6 +199,12 @@ class PrismaDBExceptionHandler: return True return False + @staticmethod + def is_prisma_error(e: Exception) -> bool: + import prisma + + return isinstance(e, _exception_types(prisma.errors.PrismaError)) + @staticmethod def is_deadlock_error(e: Exception) -> bool: """True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma.""" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1f453d3b1ba..accf7b720fb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6386,10 +6386,17 @@ class ProxyUpdateSpend: ) break except Exception as e: - if not PrismaDBExceptionHandler.is_database_transport_error(e): + if not _is_transient_spend_log_write_error(e): + if PrismaDBExceptionHandler.is_prisma_error(e): + await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) + verbose_proxy_logger.warning( + "Spend tracking - DB error writing spend logs, requeued %d rows for the next flush. error=%s", + len(logs_to_process), + str(e), + ) raise verbose_proxy_logger.warning( - "Spend tracking - DB connection error writing spend logs, retry %d/%d. logs_count=%d, error=%s", + "Spend tracking - transient DB error writing spend logs, retry %d/%d. logs_count=%d, error=%s", i + 1, n_retry_times, len(logs_to_process), @@ -6732,6 +6739,10 @@ async def _monitor_spend_logs_queue( MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256 +def _is_transient_spend_log_write_error(e: Exception) -> bool: + return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) + + async def _create_spend_logs_with_poison_isolation( repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]], @@ -6767,6 +6778,8 @@ async def _create_spend_logs_with_poison_isolation( raise if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): raise + if PrismaDBExceptionHandler.is_deadlock_error(e): + raise budget_left: Final = max(failure_budget - 1, 0) if len(rows) == 1: request_id: Final = rows[0].get("request_id") diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index 048fddb10d6..d671a4ffc1f 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -473,6 +473,110 @@ async def test_update_spend_logs_retries_and_requeues_batch_on_db_outage( assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"] +def _deadlock_error() -> Exception: + return _data_error( + 'Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: None, ' + 'kind: QueryError(PostgresError { code: "40P01", message: "deadlock detected", severity: "ERROR" }) })' + ) + + +@pytest.mark.asyncio +async def test_update_spend_logs_retries_deadlock_and_keeps_every_row( + mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """A 40P01 deadlock aborts the whole insert, so the same rows succeed on replay. + Before the fix the deadlock surfaced as a plain ``DataError`` and went through + poison-row isolation, which bisected the batch and dropped every row the + deadlock happened to hit as if Postgres had rejected it. + """ + + async def _fake_sleep(_: float) -> None: + return None + + monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) + create_many = AsyncMock(side_effect=[_deadlock_error(), _deadlock_error(), None]) + mock_prisma_client.db.litellm_spendlogs.create_many = create_many + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [] + + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=2, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")], + ) + + attempts = tuple( + tuple(row["request_id"] for row in call.kwargs["data"]) for call in create_many.await_args_list + ) + assert attempts == (("a", "b"), ("a", "b"), ("a", "b")) + assert mock_prisma_client.spend_log_transactions == [] + + +@pytest.mark.asyncio +async def test_update_spend_logs_requeues_batch_once_deadlock_retries_exhaust( + mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """If every retry deadlocks, the batch goes back to the head of the queue for + the next flush instead of being dropped. + """ + + async def _fake_sleep(_: float) -> None: + return None + + monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_deadlock_error()) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")] + + with pytest.raises(type(_deadlock_error())): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=1, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")], + ) + + assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 2 + assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_update_spend_logs_requeues_batch_on_non_transport_db_error( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """A DB error that is neither transport nor deadlock (here P2021, the table is + gone mid-migration) is not retried in place, but the dequeued batch must not + be lost either: it goes back to the head of the queue so it lands once the + DB is healthy again. + """ + from prisma.errors import TableNotFoundError + + err = TableNotFoundError( + {"user_facing_error": {"error_code": "P2021", "message": "The table does not exist", "meta": {}}} + ) + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=err) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")] + + with pytest.raises(TableNotFoundError): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=2, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")], + ) + + assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1 + assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"] + + @pytest.mark.asyncio async def test_requeue_after_outage_drops_oldest_logs_past_the_byte_budget( mock_prisma_client: Any, make_spend_log_row: Any @@ -549,8 +653,9 @@ async def test_flush_returns_the_bytes_it_took_off_the_queue(mock_prisma_client: async def test_update_spend_logs_does_not_requeue_non_transport_failures( mock_prisma_client: Any, make_spend_log_row: Any ) -> None: - """Only transport failures are worth replaying. A rejection the DB will keep - rejecting must not be requeued, or it would wedge the queue forever. + """Only DB failures are worth replaying. A row the proxy itself cannot + serialize would fail the same way on every flush, so requeueing it would + wedge the head of the queue forever. """ mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=ValueError("bad payload")) proxy_logging = MagicMock() From e3b4a82ff9991369f0a79e34f44a5da732506b0f Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 5 Sep 2026 11:44:34 -0700 Subject: [PATCH 7/8] Merge pull request #39926 from BerriAI/litellm_lit6981_none_url_auth fix(mcp): reject URL credentials for none auth --- .../_experimental/mcp_server/exceptions.py | 14 +++++++ .../outbound_credentials/adapter.py | 3 ++ .../outbound_credentials/resolver.py | 11 ++++- .../mcp_server/outbound_credentials/types.py | 12 ++++++ .../mcp_server/rest_endpoints.py | 3 ++ .../outbound_credentials/test_adapter.py | 12 ++++++ .../outbound_credentials/test_resolver.py | 29 +++++++++++++ .../outbound_credentials/test_types.py | 9 ++++ .../mcp_server/test_mcp_server_manager.py | 42 +++++++++++++++++++ .../mcp_server/test_rest_endpoints.py | 30 +++++++++++++ 10 files changed, 164 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/exceptions.py b/litellm/proxy/_experimental/mcp_server/exceptions.py index a1b3b167a4a..c818f6b05bd 100644 --- a/litellm/proxy/_experimental/mcp_server/exceptions.py +++ b/litellm/proxy/_experimental/mcp_server/exceptions.py @@ -5,6 +5,20 @@ from typing import Final from fastapi import HTTPException +class MCPServerURLCredentialsError(HTTPException): + """A fixed, sanitized URL-credential migration error safe for operator previews.""" + + def __init__(self) -> None: + super().__init__( + status_code=500, + detail=( + "misconfigured: auth_type none cannot be used with credentials embedded in the upstream URL; " + "remove them from the URL and configure Basic Auth with auth_type: basic and " + "auth_value: username:password" + ), + ) + + class MCPUpstreamAuthError(Exception): """Raised when an upstream MCP server returns an authentication failure (typically HTTP 401) and the gateway should surface it transparently to diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 6a95a93a2a8..ea2318bd6f1 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -19,6 +19,7 @@ from pydantic import SecretStr from typing_extensions import assert_never from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials +from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( DEFAULT_CREDENTIAL_HEADER, @@ -293,6 +294,8 @@ def raise_public(error: CredError) -> NoReturn: ) case "misconfigured": raise HTTPException(status_code=500, detail=error.summary) + case "url_credentials_not_allowed": + raise MCPServerURLCredentialsError() case "upstream_unavailable": raise HTTPException(status_code=503, detail=error.summary) case "unsupported_mode": diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 3af7b51f432..404baa14350 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -134,7 +134,7 @@ class UpstreamCredentialProvider: async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: case NoneConfig(): - return Ok(NoOpAuth()) + return self._none(server) case ApiKeyConfig() as config: return self._api_key(config) case PassthroughConfig(): @@ -151,6 +151,15 @@ class UpstreamCredentialProvider: return _not_implemented(AuthSpecKind.aws_sigv4) assert_never(server.config) + def _none(self, server: ServerSpec) -> Result[httpx.Auth, CredError]: + try: + resource: Final = httpx.URL(server.resource) + except httpx.InvalidURL: + return Ok(NoOpAuth()) + if resource.userinfo: + return Error(CredError.of_url_credentials_not_allowed()) + return Ok(NoOpAuth()) + async def has_user_token(self, subject: Subject, server: ServerSpec) -> bool: """Whether a usable per-user token exists for this server (the preemptive 401's check). diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 67aad3e443e..632dc57dcf6 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -95,6 +95,7 @@ class CredError: tag: Literal[ "unauthorized", "misconfigured", + "url_credentials_not_allowed", "upstream_unavailable", "unsupported_mode", "precondition_required", @@ -103,6 +104,7 @@ class CredError: unauthorized: Unauthorized = case() # no usable credential for this (subject, server) -> 401 challenge misconfigured: str = case() # the declared mode is missing required config -> 5xx (operator) + url_credentials_not_allowed: None = case() upstream_unavailable: str = case() # the IdP / token endpoint could not be reached -> 503 unsupported_mode: str = case() # a raw mode string did not parse into AuthSpecKind (boundary) precondition_required: str = case() # a required per-user value (e.g. an env var) has not been provided -> 412 @@ -129,6 +131,10 @@ class CredError: def of_misconfigured(detail: str) -> CredError: return CredError(misconfigured=detail) + @staticmethod + def of_url_credentials_not_allowed() -> CredError: + return CredError(url_credentials_not_allowed=None) + @staticmethod def of_upstream_unavailable(detail: str) -> CredError: return CredError(upstream_unavailable=detail) @@ -154,6 +160,12 @@ class CredError: return f"unauthorized: {self.unauthorized.detail}" case "misconfigured": return f"misconfigured: {self.misconfigured}" + case "url_credentials_not_allowed": + return ( + "misconfigured: auth_type none cannot be used with credentials embedded in the upstream URL; " + "remove them from the URL and configure Basic Auth with auth_type: basic and " + "auth_value: username:password" + ) case "upstream_unavailable": return f"upstream unavailable: {self.upstream_unavailable}" case "unsupported_mode": diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index b3469da9071..5fbfad54a39 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -18,6 +18,7 @@ from litellm.exceptions import ( ) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, + MCPServerURLCredentialsError, MCPUpstreamAuthError, ) from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( @@ -75,6 +76,8 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + if isinstance(exc, MCPServerURLCredentialsError): + return str(exc.detail) if isinstance(exc, TimeoutError): return ( f"Failed to connect to MCP server: no response from {url or 'the server'} " diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index c6f3b9cb1f4..d67d0df4d0e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -13,6 +13,7 @@ from fastapi import HTTPException from pydantic import ValidationError from litellm.experimental_mcp_client.client import MCPClient +from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( oauth_protected_resource_path, raise_public, @@ -442,6 +443,17 @@ def test_raise_public_maps_each_error_to_its_status(error, status): assert exc_info.value.status_code == status +def test_raise_public_marks_only_url_credentials_error_as_safe_for_preview(): + with pytest.raises(HTTPException) as generic_exc_info: + raise_public(CredError.of_misconfigured("private operator detail")) + assert not isinstance(generic_exc_info.value, MCPServerURLCredentialsError) + + error = CredError.of_url_credentials_not_allowed() + with pytest.raises(MCPServerURLCredentialsError) as url_exc_info: + raise_public(error) + assert url_exc_info.value.detail == error.summary + + def test_raise_public_emits_unauthorized_challenge(): body = {"error": "byok_auth_required", "server_id": "s1"} error = CredError.of_unauthorized("needs key", www_authenticate='Bearer resource_metadata="/x"', body=body) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 9d63e8c2c1c..5e2f2cf97d7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -116,6 +116,35 @@ async def test_none_mode_yields_a_no_op_auth(): assert isinstance(result.ok, NoOpAuth) +@pytest.mark.asyncio +async def test_none_mode_rejects_url_userinfo(): + spec = ServerSpec( + server_id="s", + resource="https://lit-user:s3cr3t@upstream.example.com/mcp", + config=NoneConfig(), + ) + + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, spec) + + assert isinstance(result, Error) + assert result.error.tag == "url_credentials_not_allowed" + assert "Basic Auth" in result.error.summary + assert "auth_type: basic" in result.error.summary + assert "auth_value: username:password" in result.error.summary + assert "lit-user" not in result.error.summary + assert "s3cr3t" not in result.error.summary + + +@pytest.mark.asyncio +async def test_none_mode_does_not_validate_non_credential_resource(): + spec = ServerSpec(server_id="s", resource="https://[::1", config=NoneConfig()) + + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, spec) + + assert isinstance(result, Ok) + assert isinstance(result.ok, NoOpAuth) + + @pytest.mark.asyncio async def test_api_key_shared_emits_the_configured_header(): config = ApiKeyConfig( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py index d4b51b08e06..bacbb5c1236 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py @@ -75,6 +75,15 @@ def test_crederror_factory_sets_the_matching_tag(factory, expected_tag): assert "detail text" in err.summary +def test_url_credentials_error_has_a_fixed_actionable_summary(): + err = CredError.of_url_credentials_not_allowed() + + assert err.tag == "url_credentials_not_allowed" + assert "Basic Auth" in err.summary + assert "auth_type: basic" in err.summary + assert "auth_value: username:password" in err.summary + + def test_apikeyconfig_requires_a_key_source(): with pytest.raises(ValidationError): ApiKeyConfig() # type: ignore[call-arg] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 9745e508703..34dc067e7a3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -8966,6 +8966,24 @@ class TestCreateMcpClientV2Graft: assert isinstance(client._resolved_auth, NoOpAuth) assert client._mcp_auth_value is None + @pytest.mark.parametrize("auth_type", [None, MCPAuth.none]) + async def test_none_mode_rejects_url_userinfo(self, auth_type): + with pytest.raises(HTTPException) as exc_info: + await MCPServerManager()._create_mcp_client( + self._http_server( + auth_type=auth_type, + url="https://lit-user:s3cr3t@upstream.example.com/mcp", + ) + ) + + detail = str(exc_info.value.detail) + assert exc_info.value.status_code == 500 + assert "Basic Auth" in detail + assert "auth_type: basic" in detail + assert "auth_value: username:password" in detail + assert "lit-user" not in detail + assert "s3cr3t" not in detail + @pytest.mark.parametrize( "auth_type, token, expected_name, expected_value", [ @@ -11369,6 +11387,30 @@ class TestResolveOpenapiToolAuth: assert "Authorization" not in (forwarded or {}) + @pytest.mark.asyncio + async def test_none_mode_without_url_keeps_spec_path_server_unauthenticated(self): + server = MCPServer( + server_id="openapi-only", + name="report_api", + server_name="report_api", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.none, + spec_path="https://api.example.com/openapi.json", + ) + + resolved, forwarded = await MCPServerManager().resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers=None, + raw_headers=None, + mcp_auth_header=None, + user_api_key_auth=None, + forwarded_headers={"X-Trace": "trace-id"}, + ) + + assert resolved is None + assert forwarded == {"X-Trace": "trace-id"} + class TestOpenApiHandlerRelaysUpstreamAuth: """`_call_openapi_tool_handler` must not flatten a re-auth signal into a generic message. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index f2c8f8c80c5..c007d22117f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -177,6 +177,36 @@ class TestExecuteWithMcpClient: assert "https://api.example.com/mcp/" in message assert "30s" in message + def test_connection_error_message_hides_arbitrary_http_exception_detail(self): + message = rest_endpoints._connection_error_message( + HTTPException(status_code=500, detail="secret upstream detail"), + "https://api.example.com/mcp/", + 30.0, + ) + + assert "secret upstream detail" not in message + + @pytest.mark.asyncio + async def test_none_mode_url_credentials_returns_actionable_redacted_error(self): + async def unreached_operation(client): + raise AssertionError("operation must not run for an invalid server configuration") + + payload = NewMCPServerRequest( + server_name="example", + url="https://lit-user:s3cr3t@upstream.example.com/mcp", + auth_type=MCPAuth.none, + ) + + result = await rest_endpoints._execute_with_mcp_client(payload, unreached_operation) + + message = str(result["message"]) + assert result["error"] is True + assert "Basic Auth" in message + assert "auth_type: basic" in message + assert "auth_value: username:password" in message + assert "lit-user" not in message + assert "s3cr3t" not in message + @pytest.mark.asyncio async def test_forwards_static_headers(self, monkeypatch): """Ensure static_headers are forwarded to the MCP client during test calls. From a0058ed15759febc3acb96ab27c823671a232715 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 5 Sep 2026 11:47:36 -0700 Subject: [PATCH 8/8] fix(hide-secrets): stop redacting benign identifiers (#39879) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(hide-secrets): stop redacting benign identifiers and make redaction deterministic The OpenAI key detector matched `sk-` anywhere inside a word, so `` became ``, and the Base64 entropy limit of 3.0 flagged ordinary quoted identifiers such as `"application/json"` and model ids. Redaction also iterated a hash-seeded set, so the same request produced different bytes on different workers and broke prompt caching. - require a standalone `sk-`/`sk_` token with a digit (still catches sk-proj-/sk-ant-) - raise Base64HighEntropyString limit from 3.0 to the detect-secrets default 4.5 - redact overlapping matches longest-first in a stable order Resolves LIT-7049 * fix(hide-secrets): treat separators as key boundaries and defer sk_live_ to the stripe detector The standalone-token boundary also rejected keys glued to a preceding `_`, `-` or percent-encoded delimiter (`openai_sk-…`, `key-sk-…`, `Bearer%20sk-…`), which the old pattern redacted, and `sk_live_…` was counted by both the OpenAI and the Stripe detector. * fix(hide-secrets): keep the openai key scan linear on repeated sk separators The digit requirement was a lookahead, so every `sk` inside a long `[a-zA-Z0-9_-]` run re-scanned the rest of that run looking for a digit. 100 KB of `-sk-` took over 5s in the worker's event loop and the proxy closed the connection without a response. The check now runs once per match in `analyze_string` instead. * chore(hide-secrets): remove redundant performance test comment * fix(hide-secrets): consume complete openai key tokens * chore(hide-secrets): remove redundant fixture comment * chore(hide-secrets): remove redundant test docstrings * fix(hide-secrets): redact whole stripe live keys * style(hide-secrets): wrap secret sorting key --- .../enterprise_callbacks/secret_detection.py | 27 ++--- .../secrets_plugins/openai_api_key.py | 15 ++- .../test_secret_detection.py | 101 ++++++++++++++++-- 3 files changed, 124 insertions(+), 19 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index 1fddc527ec8..bfbfd7bfb15 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -433,9 +433,9 @@ _default_detect_secrets_config = { "name": "ZendeskSecretKeyDetector", "path": _custom_plugins_path + "/zendesk_secret_key.py", }, - {"name": "Base64HighEntropyString", "limit": 3.0}, + {"name": "Base64HighEntropyString", "limit": 4.5}, {"name": "HexHighEntropyString", "limit": 3.0}, - ] + ], } @@ -466,16 +466,19 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): os.remove(temp_file.name) - detected_secrets = [] - for file in secrets.files: - for found_secret in secrets[file]: - if found_secret.secret_value is None: - continue - detected_secrets.append( - {"type": found_secret.type, "value": found_secret.secret_value} - ) - - return detected_secrets + return [ + {"type": found_secret.type, "value": found_secret.secret_value} + for file in sorted(secrets.files) + for found_secret in sorted( + secrets[file], + key=lambda secret: ( + -len(secret.secret_value or ""), + secret.type, + secret.secret_value or "", + ), + ) + if found_secret.secret_value is not None + ] def redact_text(self, text: str, source: str = "message") -> str: """Replace every detected secret in ``text`` with ``[REDACTED]`` and diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py b/enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py index c5d20f75909..32652703326 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py @@ -3,6 +3,7 @@ This plugin searches for OpenAI API Keys. """ import re +from collections.abc import Generator from detect_secrets.plugins.base import RegexBasedDetector @@ -16,4 +17,16 @@ class OpenAIApiKeyDetector(RegexBasedDetector): @property def denylist(self) -> list[re.Pattern]: - return [re.compile(r"""(sk-[a-zA-Z0-9]{5,})""")] + return [ + re.compile( + r"((?:(? Generator[str, None, None]: + # the digit check lives outside the regex: a lookahead re-scans the token + # from every `sk` inside it, which is quadratic on `-sk-sk-sk-...` input + yield from (match for match in super().analyze_string(string) if re.search(r"[0-9]", match)) diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py index dc1cbb9983e..f46df5baadf 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py @@ -10,6 +10,8 @@ Covers the three defects from the ticket: handling live only on the native path). """ +import time + import pytest from litellm_enterprise.enterprise_callbacks.secret_detection import ( @@ -19,12 +21,16 @@ from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth AWS_KEY = "AKIAIOSFODNN7EXAMPLE" +OPENAI_KEY = "sk-test-abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH" +SHORT_OPENAI_KEY = "sk-12345" +UNICODE_DIGIT_SUFFIX = "sk-notification٣" +STRIPE_LIVE_KEY = f"sk_live_{'1234567890' * 3}" +URL_ENCODED_KEY = "Bearer%20sk-Ab3dEf6Gh7Ij8Kl9Mn0Pq2Rs3Tu4Vw5X" +AWS_KEYS = [f"AKIAIOSFODNN7EXAMPL{suffix}" for suffix in "FEDCBA"] def _guardrail() -> _ENTERPRISE_SecretDetection: - return _ENTERPRISE_SecretDetection( - guardrail_name="hide-secrets", event_hook="pre_call", default_on=True - ) + return _ENTERPRISE_SecretDetection(guardrail_name="hide-secrets", event_hook="pre_call", default_on=True) def _recorded(request_data: dict) -> dict: @@ -33,6 +39,91 @@ def _recorded(request_data: dict) -> dict: return entries[0] +def test_scan_message_preserves_benign_identifiers_and_xml_tags(): + guardrail = _guardrail() + content = " model: claude-sonnet-4-5-20250929 " + + assert guardrail.scan_message_for_secrets(content) == [] + assert guardrail.redact_text(content) == content + assert guardrail.redact_text("result = compute(x) ") == ( + "result = compute(x) " + ) + + +def test_scan_message_preserves_quoted_benign_identifiers(): + guardrail = _guardrail() + content = '{"content-type": "application/json", "model": "claude-sonnet-4-5-20250929"}' + + assert guardrail.scan_message_for_secrets(content) == [] + assert guardrail.redact_text(content) == content + + +def test_scan_message_redacts_every_openai_key_occurrence(): + guardrail = _guardrail() + content = f"first {OPENAI_KEY}, second {OPENAI_KEY}" + + assert guardrail.redact_text(content) == "first [REDACTED], second [REDACTED]" + + +def test_scan_message_redacts_short_numeric_openai_like_values(): + guardrail = _guardrail() + + assert guardrail.redact_text(f"value {SHORT_OPENAI_KEY}") == "value [REDACTED]" + + +def test_scan_message_requires_ascii_digits_for_openai_like_values(): + guardrail = _guardrail() + + assert guardrail.scan_message_for_secrets(UNICODE_DIGIT_SUFFIX) == [] + assert guardrail.redact_text(UNICODE_DIGIT_SUFFIX) == UNICODE_DIGIT_SUFFIX + + +def test_scan_message_redacts_openai_key_after_separator(): + guardrail = _guardrail() + + assert guardrail.redact_text(f"openai_{OPENAI_KEY} key-{OPENAI_KEY}") == ( + "openai_[REDACTED] key-[REDACTED]" + ) + assert guardrail.redact_text(URL_ENCODED_KEY) == "Bearer%20[REDACTED]" + + +def test_scan_message_does_not_stop_openai_key_at_token_characters(): + guardrail = _guardrail() + + assert guardrail.redact_text("key sk-proj-abcde12345/extra") == "key [REDACTED]/extra" + + +def test_scan_message_stays_linear_on_repeated_sk_separators(): + guardrail = _guardrail() + content = "-sk-" * 25_000 + + started = time.perf_counter() + assert guardrail.scan_message_for_secrets(content) == [] + assert time.perf_counter() - started < 2.0 + + +def test_scan_message_redacts_whole_stripe_live_key(): + guardrail = _guardrail() + + assert guardrail.redact_text(f"stripe {STRIPE_LIVE_KEY} end") == "stripe [REDACTED] end" + + +def test_scan_message_returns_matches_in_stable_order(): + guardrail = _guardrail() + detected = guardrail.scan_message_for_secrets(" ".join(AWS_KEYS)) + + assert [secret["value"] for secret in detected] == sorted(AWS_KEYS) + + +def test_scan_message_replaces_longest_overlapping_match_first(): + guardrail = _guardrail() + content = f'token = "{OPENAI_KEY}/extra"' + + detected = guardrail.scan_message_for_secrets(content) + assert [secret["value"] for secret in detected] == [f"{OPENAI_KEY}/extra", OPENAI_KEY] + assert guardrail.redact_text(content) == 'token = "[REDACTED]"' + + @pytest.mark.asyncio async def test_apply_guardrail_redacts_secrets(): """Playground path: the returned texts must carry [REDACTED], not the secret.""" @@ -199,9 +290,7 @@ async def test_apply_guardrail_without_texts_records_nothing(): "messages": [ { "role": "user", - "content": [ - {"type": "image_url", "image_url": {"url": "https://x/y.png"}} - ], + "content": [{"type": "image_url", "image_url": {"url": "https://x/y.png"}}], } ], "metadata": {},