fix(proxy): keep cursor chat framing and JSON-encode Responses hook errors

This commit is contained in:
mateo-berri 2026-10-29 05:26:01 -07:00
parent 5003ef1f31
commit ca873a1327
7 changed files with 99 additions and 41 deletions

View file

@ -1250,6 +1250,16 @@ RESPONSES_FIRST_CHUNK_FRAMING: Final = FirstChunkSseFraming(
)
@dataclass(frozen=True, slots=True)
class SseFraming:
first_chunk: FirstChunkSseFraming
late: LateSseFraming
OPENAI_SSE_FRAMING: Final = SseFraming(first_chunk=OPENAI_FIRST_CHUNK_FRAMING, late=OPENAI_LATE_SSE_FRAMING)
RESPONSES_SSE_FRAMING: Final = SseFraming(first_chunk=RESPONSES_FIRST_CHUNK_FRAMING, late=RESPONSES_LATE_SSE_FRAMING)
def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
"""`headers` plus the two that stop reverse proxies from buffering SSE (issue #28384)."""
return MappingProxyType({**headers, **_TTFT_KEEPALIVE_HEADERS})
@ -2945,7 +2955,7 @@ class ProxyBaseLLMRequestProcessing:
is_streaming_request: bool | None = False,
contents: list[object] | None = None,
skip_pre_call_logic: bool = False,
late_framing: LateSseFraming = OPENAI_LATE_SSE_FRAMING,
sse_framing: SseFraming = OPENAI_SSE_FRAMING,
) -> object:
"""Run the request, sending SSE keepalives while the upstream is still silent.
@ -2986,10 +2996,11 @@ class ProxyBaseLLMRequestProcessing:
is_streaming_request=is_streaming_request,
contents=contents,
skip_pre_call_logic=skip_pre_call_logic,
first_chunk_framing=sse_framing.first_chunk,
),
ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router),
on_late_failure=_audit_late_failure,
late_framing=late_framing,
late_framing=sse_framing.late,
)
async def _process_llm_request(
@ -3013,6 +3024,7 @@ class ProxyBaseLLMRequestProcessing:
is_streaming_request: bool | None = False,
contents: list[object] | None = None, # Add contents parameter
skip_pre_call_logic: bool = False,
first_chunk_framing: FirstChunkSseFraming = OPENAI_FIRST_CHUNK_FRAMING,
) -> Any:
"""
Common request processing logic for both chat completions and responses API endpoints
@ -3306,9 +3318,7 @@ class ProxyBaseLLMRequestProcessing:
request=request,
refresh_headers=refresh_stream_headers,
general_settings=general_settings,
framing=(
RESPONSES_FIRST_CHUNK_FRAMING if route_type == "aresponses" else OPENAI_FIRST_CHUNK_FRAMING
),
framing=first_chunk_framing,
)
### CALL HOOKS ### - modify outgoing data

View file

@ -61,19 +61,22 @@ class _FailureDetails(LiteLLMBaseModel):
class _ErrorPayload(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
error: Mapping[str, object]
error: Mapping[str, object] | str
def error_object(self) -> Mapping[str, object]:
return {"message": self.error} if isinstance(self.error, str) else self.error
def sse_error_object(data: str | bytes) -> Mapping[str, object] | None:
try:
return _ErrorPayload.model_validate_json(data).error
return _ErrorPayload.model_validate_json(data).error_object()
except ValidationError:
return None
def event_error_object(event: Mapping[str, object]) -> Mapping[str, object] | None:
try:
return _ErrorPayload.model_validate(event).error or None
return _ErrorPayload.model_validate(event).error_object() or None
except ValidationError:
return None

View file

@ -10202,6 +10202,12 @@ def _format_streaming_sse_chunk(chunk: object) -> str | bytes:
return f"data: {chunk}\n\n"
def _format_responses_sse_chunk(chunk: object) -> str | bytes:
if isinstance(chunk, Mapping):
return f"data: {json.dumps(chunk, default=str)}\n\n"
return _format_streaming_sse_chunk(chunk=chunk)
_SSE_FRAME_DELIMITERS: Final = ("\r\n\r\n", "\n\n", "\r\r")
_OPENAI_STREAM_DONE_FRAME: Final = "data: [DONE]\n\n"
_MAX_RAW_SSE_BUFFER_CHARS: Final = 8 * 1024 * 1024
@ -10602,7 +10608,7 @@ async def async_data_generator(
if not raw_passthrough:
try:
if error_state is not None:
yield error_state.mark_emitted(_format_streaming_sse_chunk(chunk=chunk))
yield error_state.mark_emitted(_format_responses_sse_chunk(chunk=chunk))
else:
yield _format_streaming_sse_chunk(chunk=chunk)
except Exception as e:

View file

@ -52,7 +52,7 @@ from litellm.proxy.auth.user_api_key_auth import (
user_api_key_auth_websocket,
)
from litellm.proxy.common_request_processing import (
RESPONSES_LATE_SSE_FRAMING,
RESPONSES_SSE_FRAMING,
ProxyBaseLLMRequestProcessing,
cost_headers,
create_response,
@ -438,7 +438,7 @@ async def responses_api(
if requests_background(_PROXY_SERVER_REQUEST_ADAPTER.validate_python(data))
else native_data_generator
),
late_framing=RESPONSES_LATE_SSE_FRAMING,
sse_framing=RESPONSES_SSE_FRAMING,
model=None,
user_model=user_model,
user_temperature=user_temperature,

View file

@ -21,7 +21,7 @@ from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_request_processing import RESPONSES_LATE_SSE_FRAMING, ProxyBaseLLMRequestProcessing
from litellm.proxy.common_request_processing import RESPONSES_SSE_FRAMING, ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.responses_stream_errors import (
event_error_object,
responses_error_fields,
@ -188,7 +188,7 @@ async def background_streaming_task(
user_api_base=user_api_base,
version=version,
skip_pre_call_logic=True,
late_framing=RESPONSES_LATE_SSE_FRAMING,
sse_framing=RESPONSES_SSE_FRAMING,
)
# Process streaming response following OpenAI events format

View file

@ -17,7 +17,7 @@ from __future__ import annotations
import asyncio
import json
from collections.abc import AsyncIterator, Callable
from collections.abc import AsyncIterator, Callable, Mapping
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock
@ -148,7 +148,8 @@ async def test_responses_stream_dispatches_named_events_without_a_chat_done_sent
@pytest.mark.asyncio
async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_event() -> None:
@pytest.mark.parametrize("as_mapping", [False, True])
async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_event(as_mapping: bool) -> None:
response: Final = ResponsesAPIResponse(id="resp_error_then_raise", created_at=1, model="local-model", output=[])
created: Final = ResponseCreatedEvent.model_validate(
{"type": "response.created", "sequence_number": 0, "response": response}
@ -157,9 +158,9 @@ async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_e
{"type": "error", "sequence_number": 1, "code": "server_error", "message": "upstream failed", "param": None}
)
async def upstream() -> AsyncIterator[BaseModel]:
async def upstream() -> AsyncIterator[BaseModel | Mapping[str, object]]:
yield created
yield error_event
yield error_event.model_dump() if as_mapping else error_event
raise litellm.APIError(status_code=500, message="Stream close failed", llm_provider="openai", model="local-model")
frames: Final = tuple(
@ -172,14 +173,26 @@ async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_e
)
assert tuple(frame.split("\n", 1)[0] for frame in frames) == ("event: response.created", "event: error")
assert json.loads(frames[-1].split("data: ", 1)[1])["message"] == "upstream failed"
@pytest.mark.asyncio
async def test_a_guardrail_error_string_mid_responses_stream_ends_it_as_a_named_error_event() -> None:
@pytest.mark.parametrize(
("guardrail_error", "error_object"),
[
(
{"message": "response withheld", "type": "guardrail_scan_error", "code": 500},
{"message": "response withheld", "type": "guardrail_scan_error", "code": 500},
),
("response withheld", {"message": "response withheld"}),
],
)
async def test_a_guardrail_error_string_mid_responses_stream_ends_it_as_a_named_error_event(
guardrail_error: Mapping[str, object] | str, error_object: Mapping[str, object]
) -> None:
response: Final = ResponsesAPIResponse(id="resp_hook_error", created_at=1, model="local-model", output=[])
created: Final = ResponseCreatedEvent.model_validate(
{"type": "response.created", "sequence_number": 0, "response": response}
)
guardrail_error: Final = {"message": "response withheld", "type": "guardrail_scan_error", "code": 500}
async def upstream() -> AsyncIterator[object]:
yield created
@ -203,7 +216,7 @@ async def test_a_guardrail_error_string_mid_responses_stream_ends_it_as_a_named_
"message": "response withheld",
"param": None,
"sequence_number": 1,
"error": guardrail_error,
"error": error_object,
}

View file

@ -181,14 +181,16 @@ async def test_streaming_upstream_errors_keep_the_client_protocol(
class _FirstFrameHook(CustomLogger):
def __init__(self, frame: str | bytes) -> None:
def __init__(self, first: str | bytes | Exception) -> None:
super().__init__()
self.frame: Final = frame
self.first: Final = first
async def async_post_call_streaming_iterator_hook(
self, user_api_key_dict: object, response: object, request_data: object
) -> AsyncIterator[str | bytes]:
yield self.frame
if isinstance(self.first, Exception):
raise self.first
yield self.first
_DATA_ERROR_FRAME: Final = 'data: {"error": {"message": "denied", "type": "guardrail_scan_error", "code": 403}}\n\n'
@ -197,22 +199,9 @@ _NAMED_ERROR_FRAME: Final = (
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"path,first_frame,expected_status",
[
("/v1/responses", _DATA_ERROR_FRAME, 403),
("/v1/chat/completions", _DATA_ERROR_FRAME, 403),
("/v1/chat/completions", _NAMED_ERROR_FRAME, 200),
("/v1/responses", _NAMED_ERROR_FRAME, 529),
("/v1/responses", _NAMED_ERROR_FRAME.replace(b"event: error", b"event: response.created"), 200),
("/v1/chat/completions", _DATA_ERROR_FRAME + "data: [DONE]\n\n", 200),
("/v1/responses", _DATA_ERROR_FRAME + "data: [DONE]\n\n", 200),
],
)
async def test_only_a_lone_first_error_frame_answers_with_its_status(
monkeypatch: pytest.MonkeyPatch, path: str, first_frame: str | bytes, expected_status: int
) -> None:
async def _post_through_first_frame_hook(
monkeypatch: pytest.MonkeyPatch, path: str, first: str | bytes | Exception
) -> httpx.Response:
import litellm.proxy.proxy_server as ps
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -230,14 +219,14 @@ async def test_only_a_lone_first_error_frame_answers_with_its_status(
)
monkeypatch.setattr(ps, "llm_router", router)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "callbacks", [_FirstFrameHook(first_frame)])
monkeypatch.setattr(litellm, "callbacks", [_FirstFrameHook(first)])
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override)
with respx.mock as transport:
transport.post(upstream_url + ("/chat/completions" if is_chat else "/responses")).respond(
200, content="data: [DONE]\n\n", headers={"Content-Type": "text/event-stream"}
)
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client:
result: Final = await client.post(
return await client.post(
path,
json={
"model": model,
@ -246,6 +235,26 @@ async def test_only_a_lone_first_error_frame_answers_with_its_status(
},
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"path,first_frame,expected_status",
[
("/v1/responses", _DATA_ERROR_FRAME, 403),
("/v1/chat/completions", _DATA_ERROR_FRAME, 403),
("/v1/chat/completions", _NAMED_ERROR_FRAME, 200),
("/cursor/chat/completions", _NAMED_ERROR_FRAME, 200),
("/v1/responses", _NAMED_ERROR_FRAME, 529),
("/v1/responses", _NAMED_ERROR_FRAME.replace(b"event: error", b"event: response.created"), 200),
("/v1/chat/completions", _DATA_ERROR_FRAME + "data: [DONE]\n\n", 200),
("/v1/responses", _DATA_ERROR_FRAME + "data: [DONE]\n\n", 200),
],
)
async def test_only_a_lone_first_error_frame_answers_with_its_status(
monkeypatch: pytest.MonkeyPatch, path: str, first_frame: str | bytes, expected_status: int
) -> None:
result: Final = await _post_through_first_frame_hook(monkeypatch, path, first_frame)
assert result.status_code == expected_status, result.text
if expected_status == 200:
assert result.headers["content-type"].startswith("text/event-stream"), result.text
@ -253,6 +262,23 @@ async def test_only_a_lone_first_error_frame_answers_with_its_status(
assert result.json()["error"]["message"] == "denied"
@pytest.mark.asyncio
@pytest.mark.parametrize("path", ["/cursor/chat/completions", "/v1/chat/completions"])
async def test_a_hook_failure_before_the_first_chunk_answers_in_chat_completions_shape(
monkeypatch: pytest.MonkeyPatch, path: str
) -> None:
message: Final = "rate limited before the first chunk"
result: Final = await _post_through_first_frame_hook(
monkeypatch, path, HTTPException(status_code=429, detail=message)
)
frames: Final = tuple(frame for frame in result.text.split("\n\n") if frame)
assert result.status_code == 429, result.text
assert frames[-1] == "data: [DONE]", result.text
assert message in json.loads(frames[0].removeprefix("data: "))["error"]["message"]
assert len(frames) == 2, result.text
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/chat/completions"])
def test_a_failure_after_keepalive_pings_keeps_the_client_protocol(monkeypatch: pytest.MonkeyPatch, path: str) -> None:
async def request_and_check() -> None: