diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 58f55ac73fd..33adb663021 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1899,6 +1899,12 @@ async def _admit_frame_turn(turn_admission: Callable[[], Awaitable[None]] | None await turn_admission() +def _turn_refusal_frame(refusal: RateLimitError) -> str: + return json.dumps( + {"type": "error", "status": 429, "error": {"type": "rate_limit_exceeded", "message": str(refusal)}} + ) + + RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ "response.created", "response.completed", @@ -2547,17 +2553,7 @@ class ResponsesWebSocketStreaming: await _admit_frame_turn(turn_admission, message) except RateLimitError as e: try: - await self.websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "rate_limit_exceeded", - "message": str(e), - }, - } - ) - ) + await self.websocket.send_text(_turn_refusal_frame(e)) except Exception: # noqa: BLE001, S110 # client may already be gone pass return False @@ -2702,10 +2698,11 @@ class ManagedResponsesWebSocketHandler: return None async def _send_error(self, message: str, error_type: str = "server_error") -> None: + await self._send_frame(json.dumps({"type": "error", "error": {"type": error_type, "message": message}})) + + async def _send_frame(self, frame: str) -> None: try: - await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) - ) + await self.websocket.send_text(frame) except Exception: pass @@ -3085,7 +3082,7 @@ class ManagedResponsesWebSocketHandler: ) await _admit_frame_turn(turn_admission, raw_message) except RateLimitError as e: - await self._send_error(str(e), error_type="rate_limit_exceeded") + await self._send_frame(_turn_refusal_frame(e)) return call_kwargs: Final = self._build_base_call_kwargs(msg_obj) diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 92c815aa4ab..1d80d87a75d 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1388,6 +1388,7 @@ class TestWebSocketTurnAdmission: ] assert turn_admission.await_count == 2 refusal: Final = json.loads(websocket.send_text.await_args.args[0]) + assert refusal["status"] == 429 assert refusal["error"]["type"] == "rate_limit_exceeded" assert "crossed budget" in refusal["error"]["message"] @@ -1442,7 +1443,7 @@ class TestWebSocketTurnAdmission: for call in websocket.send_text.await_args_list if json.loads(call.args[0]).get("type") == "error" ] - assert [refusal["error"]["type"] for refusal in refusals] == ["rate_limit_exceeded"] + assert [(refusal["status"], refusal["error"]["type"]) for refusal in refusals] == [(429, "rate_limit_exceeded")] def _deployment_defaults():