fix(responses): forward the routed input and report routing rejections on the websocket

litellm/proxy/response_api_endpoints/endpoints.py

@@ -1395,6 +1395,15 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec
     return MappingProxyType({key: value for key, value in hints.items() if value is not None})
 
 
+def _responses_ws_failure_frame(failure: Exception) -> str:
+    raw_status: Final = getattr(failure, "status_code", None)
+    status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500
+    error_type: Final = (
+        "rate_limit_exceeded" if status == 429 else "invalid_request_error" if 400 <= status < 500 else "server_error"
+    )
+    return json.dumps({"type": "error", "status": status, "error": {"type": error_type, "message": str(failure)}})
+
+
 async def _enforce_responses_ws_first_frame_model_auth(
     request: Request,
     model: str,
@@ -1574,6 +1583,15 @@ async def responses_websocket_endpoint(
                 original_exception=failure,
                 request_data=data,
             )
-    except Exception:
+    except Exception as e:
         verbose_proxy_logger.exception("Responses WebSocket error")
+        try:
+            await websocket.send_text(_responses_ws_failure_frame(e))
+        except Exception:
+            pass
+        await proxy_logging_obj.post_call_failure_hook(
+            user_api_key_dict=user_api_key_dict,
+            original_exception=e,
+            request_data=data,
+        )
         await websocket.close(code=1011, reason="Internal server error")

litellm/responses/main.py

@@ -1,5 +1,6 @@
 import asyncio
 import contextvars
+import json
 from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -8,7 +9,7 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast
 
 import httpx
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel, TypeAdapter, ValidationError
 from typing_extensions import assert_never
 
 import litellm
@@ -2277,6 +2278,24 @@ def _deployment_reasoning_default(kwargs: Mapping[str, object]) -> Reasoning | d
 _RESPONSES_WS_ROUTING_HINT_KEYS: Final = frozenset({"input", "previous_response_id"})
 
 
+def _first_ws_frame_with_routed_input(first_message: str, routed_input: object) -> str:
+    try:
+        frame: Final = _JSON_OBJECT_ADAPTER.validate_json(first_message)
+    except ValidationError:
+        return first_message
+    if frame is None or routed_input is None:
+        return first_message
+    raw_nested: Final = frame.get("response")
+    nested: Final = _JSON_OBJECT_ADAPTER.validate_python(raw_nested) if isinstance(raw_nested, Mapping) else None
+    if nested is not None and nested.get("input") is not None:
+        if nested["input"] == routed_input:
+            return first_message
+        return json.dumps({**frame, "response": {**nested, "input": routed_input}})
+    if frame.get("input") == routed_input:
+        return first_message
+    return json.dumps({**frame, "input": routed_input})
+
+
 def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ResponsesWebSocketRequestDefaults:
     default_reasoning: Final = _deployment_reasoning_default(kwargs)
     candidate_params: Final[dict[str, object]] = {
@@ -2367,10 +2386,12 @@ async def _aresponses_websocket(
         "api_base",
         "api_key",
         "timeout",
+        "first_message",
         *_RESPONSES_WS_ROUTING_HINT_KEYS,
     }
     remaining_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _explicit_keys}
     deployment_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _RESPONSES_WS_ROUTING_HINT_KEYS}
+    first_message: Final = kwargs.get("first_message")
 
     return await base_llm_http_handler.async_responses_websocket(
         model=resolved_model,
@@ -2380,6 +2401,11 @@ async def _aresponses_websocket(
         api_base=resolved_api_base,
         api_key=resolved_api_key,
         timeout=timeout,
+        first_message=(
+            _first_ws_frame_with_routed_input(first_message, kwargs.get("input"))
+            if isinstance(first_message, str)
+            else None
+        ),
         user_api_key_dict=kwargs.get("user_api_key_dict"),
         litellm_metadata=_build_litellm_metadata_for_ws(kwargs),
         custom_llm_provider=_custom_llm_provider,

tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py

@@ -638,6 +638,71 @@ class TestResponsesWSFirstFrameModelAuth:
         assert booked["user_api_key_dict"] is user_api_key_dict
         assert booked["request_data"]["model"] == "gpt-4o-mini"
 
+    @pytest.mark.asyncio
+    async def test_endpoint_sends_an_error_frame_when_routing_rejects_the_connection(self):
+        from litellm.proxy.response_api_endpoints.endpoints import (
+            responses_websocket_endpoint,
+        )
+
+        ws = MagicMock()
+        ws.headers = {}
+        ws.query_params = {}
+        ws.scope = {"headers": []}
+        ws.url = "ws://testserver/v1/responses"
+        ws.accept = AsyncMock()
+        ws.receive_text = AsyncMock(
+            return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []})
+        )
+        ws.send_text = AsyncMock()
+        ws.close = AsyncMock()
+
+        processor = MagicMock()
+        processor.common_processing_pre_call_logic = AsyncMock(
+            return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock())
+        )
+        rejection = litellm.RateLimitError(
+            message="origin deployment is cooling down", model="gpt-4o-mini", llm_provider="openai"
+        )
+        proxy_logging_obj = MagicMock()
+        proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
+        user_api_key_dict = MagicMock()
+
+        with (
+            patch(  # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above
+                "litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
+                new_callable=AsyncMock,
+            ),
+            patch(  # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint tells the client is under test
+                "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
+                return_value=processor,
+            ),
+            patch(  # test-quality-ok: routing is the seam that raises the affinity rejection
+                "litellm.proxy.route_llm_request.route_request",
+                new_callable=AsyncMock,
+                side_effect=rejection,
+            ),
+            patch(  # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row
+                "litellm.proxy.proxy_server.proxy_logging_obj",
+                proxy_logging_obj,
+            ),
+        ):
+            await responses_websocket_endpoint(
+                websocket=ws,
+                model=None,
+                user_api_key_dict=user_api_key_dict,
+            )
+
+        frame = json.loads(ws.send_text.await_args.args[0])
+        assert frame["type"] == "error"
+        assert frame["status"] == 429
+        assert frame["error"]["type"] == "rate_limit_exceeded"
+        assert "cooling down" in frame["error"]["message"]
+        ws.close.assert_awaited_once_with(code=1011, reason="Internal server error")
+        booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs
+        assert booked["original_exception"] is rejection
+        assert booked["user_api_key_dict"] is user_api_key_dict
+        assert booked["request_data"]["model"] == "gpt-4o-mini"
+
     @pytest.mark.asyncio
     async def test_reruns_model_auth_for_first_frame_model(self):
         from starlette.requests import Request

tests/test_litellm/responses/test_responses_api_request_body.py

@@ -448,6 +448,70 @@ async def test_aresponses_websocket_keeps_routing_hints_out_of_the_relay_kwargs(
         assert "previous_response_id" not in mock_ws.call_args.kwargs
 
 
+_STRIPPED_WS_INPUT = [{"role": "user", "content": "hi"}]
+_ORIGINAL_WS_INPUT = [
+    {"type": "reasoning", "id": "rs_1", "encrypted_content": "blob-from-a-removed-deployment", "summary": []},
+    *_STRIPPED_WS_INPUT,
+]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("nested", [False, True])
+async def test_aresponses_websocket_forwards_the_routed_input_in_the_first_frame(nested):  # test-quality-ok: the first frame handed to the relay is the only place the routed input is observable before the provider socket
+    from unittest.mock import MagicMock
+
+    from litellm.responses.main import _aresponses_websocket
+
+    body = {"model": "gpt-5.6", "input": _ORIGINAL_WS_INPUT, "store": False}
+    first_message = json.dumps(
+        {"type": "response.create", "response": body} if nested else {"type": "response.create", **body}
+    )
+
+    with patch.object(
+        import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket",
+        new_callable=AsyncMock,
+    ) as mock_ws:
+        await _aresponses_websocket(
+            model="openai/gpt-5.6",
+            websocket=MagicMock(),
+            api_key="sk-test",
+            litellm_logging_obj=MagicMock(),
+            input=list(_STRIPPED_WS_INPUT),
+            first_message=first_message,
+        )
+
+    forwarded = json.loads(mock_ws.call_args.kwargs["first_message"])
+    container = forwarded["response"] if nested else forwarded
+    assert container["input"] == _STRIPPED_WS_INPUT
+    assert container["store"] is False
+    assert container["model"] == "gpt-5.6"
+    assert forwarded["type"] == "response.create"
+
+
+@pytest.mark.asyncio
+async def test_aresponses_websocket_forwards_the_first_frame_verbatim_when_routing_left_the_input_alone():  # test-quality-ok: the relay kwargs are the boundary; byte-identical passthrough is only observable there
+    from unittest.mock import MagicMock
+
+    from litellm.responses.main import _aresponses_websocket
+
+    first_message = '{"type": "response.create", "model": "gpt-5.6",  "input": [{"role": "user", "content": "hi"}]}'
+
+    with patch.object(
+        import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket",
+        new_callable=AsyncMock,
+    ) as mock_ws:
+        await _aresponses_websocket(
+            model="openai/gpt-5.6",
+            websocket=MagicMock(),
+            api_key="sk-test",
+            litellm_logging_obj=MagicMock(),
+            input=list(_STRIPPED_WS_INPUT),
+            first_message=first_message,
+        )
+
+    assert mock_ws.call_args.kwargs["first_message"] == first_message
+
+
 _INJECTION_POINT_INPUT = [{"role": "system", "content": "You are terse."}, {"role": "user", "content": "hi"}]
 _SYSTEM_POINT = {"location": "message", "role": "system"}
 _USER_POINT = {"location": "message", "role": "user"}