scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(proxy): gate the Responses WebSocket with the router after aliases resolve, and ignore client-set internal fallback flags

Browse source
This commit is contained in:
mateo-berri 2026-10-13 23:31:35 -07:00
parent 6728e9562d
commit 19872dd0ec
6 changed files with 65 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -36,7 +36,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
read_request_body,
safe_set_request_parsed_body,
)
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.proxy.route_llm_request import endpoint_unsupported_by_model, raise_if_required_body_param_missing
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import (
REASONING_EFFORT,
@ -1564,6 +1564,14 @@ async def _responses_websocket_session(
model=resolved_model,
route_type="_aresponses_websocket",
)
unsupported_endpoint: Final = endpoint_unsupported_by_model(
route_type="_aresponses_websocket",
data=data,
llm_router=llm_router,
team_id=user_api_key_dict.team_id,
)
if unsupported_endpoint is not None:
raise unsupported_endpoint
except Exception as e:
verbose_proxy_logger.exception("Responses WebSocket pre-call error")
try:

10
litellm/proxy/route_llm_request.py
View file

@ -15,7 +15,6 @@ from litellm.router_utils.common_utils import is_proxy_admin_request
from litellm.router_utils.fallback_event_handlers import (
check_non_standard_fallback_format,
fallback_lookup_groups,
fallbacks_disabled_for_request,
get_fallback_model_group_for_lookup_groups,
has_unattempted_fallback_target,
)
@ -514,12 +513,15 @@ def _fallback_would_answer(
team_id: str | None,
walked: frozenset[str] = frozenset(),
) -> bool:
if fallbacks_disabled_for_request(data):
if data.get("disable_fallbacks") is True:
return False
raw_fallbacks: Final = _fallbacks_for_request(data, llm_router)
if raw_fallbacks is None:
return False
try:
fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(_fallbacks_for_request(data, llm_router))
fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(raw_fallbacks)
except ValidationError:
return False
return True
walked_here: Final = walked | {model}
return any(
has_unattempted_fallback_target((target,), data)

36
tests/integration/routing/test_non_chat_model_endpoint_gate.py
View file

@ -1,4 +1,5 @@
import asyncio
import contextlib
import json
import uuid
from collections.abc import Iterator, Mapping
@ -13,6 +14,7 @@ import anthropic
import httpx
import openai
import pytest
import websockets
import yaml
from integration._support.client import (
JSON_OBJECT,
@ -30,6 +32,7 @@ from integration._support.openai_wire import answering_model_discovery, chat_rep
from integration._support.process import graceful_stop_seconds, owned_proxy_process, owned_upstream, paused
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
from websockets.exceptions import ConnectionClosed
from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse
@ -526,6 +529,39 @@ def test_key_alias_to_an_embedding_only_model_is_rejected_under_the_resolved_nam
_assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", target, "/v1/embeddings"))
async def _websocket_first_frame_and_close(
gateway: Gateway, key: str, model: str, marker: str
) -> tuple[dict[str, JsonValue], int | None, str | None]:
url: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://", 1) + "/v1/responses"
async with websockets.connect(url, additional_headers={"Authorization": f"Bearer {key}"}, open_timeout=10) as ws:
await ws.send(json.dumps({"type": "response.create", "model": model, "input": marker}))
first: Final = JSON_OBJECT.validate_json(await asyncio.wait_for(ws.recv(), 30))
with contextlib.suppress(ConnectionClosed):
await asyncio.wait_for(ws.recv(), 30)
return first, ws.close_code, ws.close_reason
@pytest.mark.parametrize("aliased", [False, True])
def test_responses_websocket_closes_on_an_embedding_only_model_before_the_upstream(rig: _Rig, aliased: bool) -> None:
target: Final = rig.names["embedding_only"]
alias: Final = f"gate-ws-alias-{uuid.uuid4().hex[:12]}"
marker: Final = _marker()
with rig.gateway.scenario() as scenario:
key: Final = scenario.key(aliases={alias: target})
frame, close_code, close_reason = asyncio.run(
_websocket_first_frame_and_close(rig.gateway, key, alias if aliased else target, marker)
)
assert frame == {
"type": "error",
"error": {
"type": "invalid_request_error",
"message": f"/v1/responses is not supported by model={target}. Supported endpoints: /v1/embeddings",
},
}
assert (close_code, close_reason) == (1008, "Pre-call error")
assert rig.upstream.calls(marker) == ()
def test_team_scoped_embedding_only_model_is_rejected_under_its_public_name(rig: _Rig) -> None:
public: Final = f"gate-team-{uuid.uuid4().hex[:12]}"
marker: Final = _marker()

22
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -1007,7 +1007,8 @@ class TestResponsesWSFirstFrameModelAuth:
class TestResponsesWSDeclaredNonChatModel:
@pytest.mark.asyncio
@pytest.mark.parametrize("query_model", [None, "embed-only"])
@pytest.mark.parametrize("requested_model", ["embed-only", "embed-alias"])
@pytest.mark.parametrize("use_query_model", [False, True])
@pytest.mark.parametrize(
"model_info, rejected",
[
@ -1016,12 +1017,11 @@ class TestResponsesWSDeclaredNonChatModel:
],
)
async def test_session_closes_on_a_model_declared_non_chat_before_routing(
self, query_model: str | None, model_info: dict[str, object], rejected: bool
self, requested_model: str, use_query_model: bool, model_info: dict[str, object], rejected: bool
):
from litellm.proxy.response_api_endpoints.endpoints import (
responses_websocket_endpoint,
)
from litellm.proxy.route_llm_request import endpoint_unsupported_by_model
router: Final = litellm.Router(
model_list=[
@ -1039,21 +1039,17 @@ class TestResponsesWSDeclaredNonChatModel:
ws.url = "ws://testserver/v1/responses"
ws.accept = AsyncMock()
ws.receive_text = AsyncMock(
return_value=json.dumps({"type": "response.create", "model": "embed-only", "input": []})
return_value=json.dumps({"type": "response.create", "model": requested_model, "input": []})
)
ws.send_text = AsyncMock()
ws.close = AsyncMock()
key_aliases: Final = {"embed-alias": "embed-only"}
async def gated_pre_call(**kwargs):
unsupported = endpoint_unsupported_by_model(
route_type=kwargs["route_type"], data={"model": kwargs["model"]}, llm_router=router, team_id=None
)
if unsupported is not None:
raise unsupported
return {"model": kwargs["model"], "litellm_metadata": {}}, MagicMock()
async def alias_resolving_pre_call(**kwargs):
return {"model": key_aliases.get(kwargs["model"], kwargs["model"]), "litellm_metadata": {}}, MagicMock()
processor = MagicMock()
processor.common_processing_pre_call_logic = AsyncMock(side_effect=gated_pre_call)
processor.common_processing_pre_call_logic = AsyncMock(side_effect=alias_resolving_pre_call)
async def fake_llm_call():
return None
@ -1076,7 +1072,7 @@ class TestResponsesWSDeclaredNonChatModel:
):
await responses_websocket_endpoint(
websocket=ws,
model=query_model,
model=requested_model if use_query_model else None,
user_api_key_dict=MagicMock(team_id=None),
)

10
tests/unit/proxy/test_common_request_processing.py
View file

@ -725,16 +725,12 @@ class TestProxyBaseLLMRequestProcessing:
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"route_type, endpoint",
[("acompletion", "/v1/chat/completions"), ("_aresponses_websocket", "/v1/responses")],
)
@pytest.mark.parametrize(
"requested_model, key_aliases",
[("embed-only", {}), ("embed-alias", {"embed-alias": "embed-only"})],
)
async def test_common_processing_pre_call_logic_rejects_chat_on_embedding_only_model_before_pre_call_hooks(
self, monkeypatch, route_type, endpoint, requested_model, key_aliases
self, monkeypatch, requested_model, key_aliases
):
processing_obj, mock_request, mock_proxy_logging_obj, mock_proxy_config, _ = self._guardrail_tag_budget_harness(
monkeypatch,
@ -749,13 +745,13 @@ class TestProxyBaseLLMRequestProcessing:
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", aliases=key_aliases),
proxy_logging_obj=mock_proxy_logging_obj,
proxy_config=mock_proxy_config,
route_type=route_type,
route_type="acompletion",
llm_router=self._router_with_embedding_only_model(),
)
assert exc_info.value.code == "400"
assert exc_info.value.message == (
f"{endpoint} is not supported by model=embed-only. Supported endpoints: /v1/embeddings"
"/v1/chat/completions is not supported by model=embed-only. Supported endpoints: /v1/embeddings"
)
mock_proxy_logging_obj.pre_call_hook.assert_not_awaited()

2
tests/unit/proxy/test_route_llm_request.py
View file

@ -1963,6 +1963,8 @@ EMBEDDING_FALLBACK_DEPLOYMENT: Final = _embedding_deployment("embed-fallback", E
({}, {"fallbacks": [{"model": "chat-fallback", "temperature": 0.2}]}),
({"fallbacks": [{"embed-only": ["chat-alias"]}], "model_group_alias": {"chat-alias": "chat-fallback"}}, {}),
({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["chat-fallback"]}]}, {}),
({"fallbacks": [{"embed-only": ["chat-fallback"]}]}, {"metadata": {"_disable_fallbacks": True}}),
({}, {"fallbacks": {"chat-fallback": ["embed-fallback"]}}),
],
)
def test_endpoint_unsupported_lets_a_configured_fallback_answer(router_kwargs, request_kwargs):