From 86bfa33516a13574e511f82ddf81cddb37876c5c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:07:15 -0700 Subject: [PATCH] fix(responses): drop malformed resume params under drop_params and keep JSON answers to a streamed retrieve A bad stream flag drops to a plain retrieve and a bad cursor to a full replay when drop_params is on, and a cursor without stream is ignored even when malformed. A streamed retrieve that the upstream answers with JSON returns the response object instead of an empty stream. get_responses forwards only the include keys from extra_query. --- litellm/llms/custom_httpx/llm_http_handler.py | 12 +++- .../proxy/response_api_endpoints/endpoints.py | 3 +- litellm/responses/main.py | 10 +++- litellm/responses/retrieve_stream_query.py | 15 +++-- .../responses/test_responses_api_lifecycle.py | 55 ++++++++++++++++++- 5 files changed, 83 insertions(+), 12 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 65502198071..aa87df5b68f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -371,6 +371,10 @@ def _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) - return {"json": data} +def _answered_json(response: httpx.Response) -> bool: + return "application/json" in response.headers.get("content-type", "").lower() + + def _stream_query_params(stream: bool, starting_after: int | None) -> Mapping[str, object]: if not stream: return {} @@ -3571,7 +3575,7 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) - if stream: + if stream and not _answered_json(response): return SyncResponsesAPIStreamingIterator( response=response, model="", @@ -3583,6 +3587,8 @@ class BaseLLMHTTPHandler: replays_stored_stream=True, ) + if stream: + response.read() return responses_api_provider_config.transform_get_response_api_response( raw_response=response, logging_obj=logging_obj, @@ -3667,7 +3673,7 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) - if stream: + if stream and not _answered_json(response): return ResponsesAPIStreamingIterator( response=response, model="", @@ -3679,6 +3685,8 @@ class BaseLLMHTTPHandler: replays_stored_stream=True, ) + if stream: + await response.aread() return responses_api_provider_config.transform_get_response_api_response( raw_response=response, logging_obj=logging_obj, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index cbae1adf858..9f1924d7f23 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -83,6 +83,7 @@ from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.proxy.utils import count_request_input_tokens, handle_exception_on_proxy from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD +from litellm.responses.retrieve_stream_query import RETRIEVE_INCLUDE_QUERY_KEYS, RETRIEVE_STREAM_QUERY_KEYS from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel @@ -102,7 +103,7 @@ router: Final = APIRouter() _RESPONSES_WS_CONFIG_VALUE_ADAPTER: Final[TypeAdapter[object | None]] = TypeAdapter(object | None) _DECODED_JSON: Final = TypeAdapter(object) _PROXY_SERVER_REQUEST_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) -_RETRIEVE_QUERY_PARAMS: Final = ("include", "include[]", "include_obfuscation", "starting_after", "stream") +_RETRIEVE_QUERY_PARAMS: Final = tuple(sorted(RETRIEVE_INCLUDE_QUERY_KEYS | RETRIEVE_STREAM_QUERY_KEYS)) _ResponseDocSchemas: TypeAlias = dict[int | str, dict[str, object]] # fastapi's responses kwarg diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 86f832b6fba..816eaade21c 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -107,7 +107,7 @@ else: MCPTool = Any from .retrieve_stream_query import ( - RETRIEVE_STREAM_QUERY_KEYS, + RETRIEVE_INCLUDE_QUERY_KEYS, InvalidRetrieveStreamQuery, RetrieveStreamQuery, parse_retrieve_stream_query, @@ -2029,12 +2029,15 @@ def _requested_retrieve_stream_query( starting_after: int | None, extra_query: Mapping[str, object] | None, custom_llm_provider: str | None, + drop_malformed: bool, ) -> RetrieveStreamQuery: explicit: Final = { key: value for key, value in (("stream", stream), ("starting_after", starting_after)) if value is not None } requested: Final = {**explicit, **(extra_query or {})} - parsed: Final = parse_retrieve_stream_query(requested.get("stream"), requested.get("starting_after")) + parsed: Final = parse_retrieve_stream_query( + requested.get("stream"), requested.get("starting_after"), drop_malformed=drop_malformed + ) if isinstance(parsed, InvalidRetrieveStreamQuery): raise litellm.BadRequestError( message=parsed.message, @@ -2217,6 +2220,7 @@ def get_responses( starting_after=starting_after, extra_query=extra_query, custom_llm_provider=custom_llm_provider, + drop_malformed=drops_unsupported_params(litellm_params.drop_params), ), responses_api_provider_config=responses_api_provider_config, custom_llm_provider=custom_llm_provider, @@ -2248,7 +2252,7 @@ def get_responses( logging_obj=litellm_logging_obj, extra_headers=extra_headers, extra_query={ - key: value for key, value in (extra_query or {}).items() if key not in RETRIEVE_STREAM_QUERY_KEYS + key: value for key, value in (extra_query or {}).items() if key in RETRIEVE_INCLUDE_QUERY_KEYS }, extra_body=extra_body, stream=stream_query.stream, diff --git a/litellm/responses/retrieve_stream_query.py b/litellm/responses/retrieve_stream_query.py index 6ad9bfcbc69..03b6da116f2 100644 --- a/litellm/responses/retrieve_stream_query.py +++ b/litellm/responses/retrieve_stream_query.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from typing import Final, Literal RETRIEVE_STREAM_QUERY_KEYS: Final = frozenset({"stream", "starting_after"}) +RETRIEVE_INCLUDE_QUERY_KEYS: Final = frozenset({"include", "include[]", "include_obfuscation"}) @dataclass(frozen=True, slots=True) @@ -16,6 +17,10 @@ class InvalidRetrieveStreamQuery: message: str +_NO_STREAM: Final = RetrieveStreamQuery(stream=False, starting_after=None) +_FULL_REPLAY: Final = RetrieveStreamQuery(stream=True, starting_after=None) + + def _parsed_stream(value: object) -> bool | InvalidRetrieveStreamQuery: match value: case None: @@ -43,12 +48,14 @@ def _parsed_starting_after(value: object) -> int | None | InvalidRetrieveStreamQ def parse_retrieve_stream_query( - stream: object, starting_after: object + stream: object, starting_after: object, drop_malformed: bool ) -> RetrieveStreamQuery | InvalidRetrieveStreamQuery: parsed_stream: Final = _parsed_stream(stream) if isinstance(parsed_stream, InvalidRetrieveStreamQuery): - return parsed_stream + return _NO_STREAM if drop_malformed else parsed_stream + if not parsed_stream: + return _NO_STREAM parsed_starting_after: Final = _parsed_starting_after(starting_after) if isinstance(parsed_starting_after, InvalidRetrieveStreamQuery): - return parsed_starting_after - return RetrieveStreamQuery(stream=parsed_stream, starting_after=parsed_starting_after) + return _FULL_REPLAY if drop_malformed else parsed_starting_after + return RetrieveStreamQuery(stream=True, starting_after=parsed_starting_after) diff --git a/tests/unit/responses/test_responses_api_lifecycle.py b/tests/unit/responses/test_responses_api_lifecycle.py index 3e01574c6cc..193ba4c89a1 100644 --- a/tests/unit/responses/test_responses_api_lifecycle.py +++ b/tests/unit/responses/test_responses_api_lifecycle.py @@ -219,6 +219,10 @@ async def _event_types(stream: object) -> tuple[str, ...]: {"stream": True, "starting_after": 1, "extra_query": {"starting_after": 3, "include[]": "usage"}}, {"stream": "true", "starting_after": "3", "include[]": "usage"}, ), + ( + {"stream": True, "extra_query": {"include": "usage", "limit": "1", "api-version": "2020-01-01"}}, + {"stream": "true", "include": "usage"}, + ), ), ) @pytest.mark.asyncio @@ -243,13 +247,16 @@ async def test_streamed_retrieve_replays_the_stored_events_after_the_cursor( @pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize("starting_after", (3, -1, "abc")) @pytest.mark.asyncio -async def test_retrieve_without_stream_still_returns_the_response_object(sync_mode: bool) -> None: +async def test_retrieve_without_stream_still_returns_the_response_object( + sync_mode: bool, starting_after: object +) -> None: with respx.mock() as mock_router: route: Final = _retrieve_route(mock_router, "openai").mock( return_value=httpx.Response(200, json=_stored_response("completed")) ) - retrieved: Final = await _retrieve(sync_mode, starting_after=3, **_retrieve_kwargs("openai")) + retrieved: Final = await _retrieve(sync_mode, starting_after=starting_after, **_retrieve_kwargs("openai")) (call,) = route.calls assert isinstance(retrieved, ResponsesAPIResponse) @@ -257,6 +264,21 @@ async def test_retrieve_without_stream_still_returns_the_response_object(sync_mo assert dict(call.request.url.params) == {} +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_streamed_retrieve_answered_with_json_returns_the_response_object(sync_mode: bool) -> None: + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, json=_stored_response("completed")) + ) + retrieved: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai")) + (call,) = route.calls + + assert isinstance(retrieved, ResponsesAPIResponse) + assert retrieved.status == "completed" + assert dict(call.request.url.params) == {"stream": "true", "starting_after": "3"} + + @pytest.mark.parametrize("sync_mode", (True, False)) @pytest.mark.asyncio async def test_resuming_at_the_end_of_a_finished_stream_ends_without_events(sync_mode: bool) -> None: @@ -401,6 +423,35 @@ async def test_a_malformed_resume_query_is_rejected_before_any_upstream_call( assert route.call_count == 0 +@pytest.mark.parametrize( + ("resume_kwargs", "expected_query"), + ( + ({"stream": "maybe", "starting_after": 3}, {}), + ({"extra_query": {"stream": "1"}}, {}), + ({"stream": True, "starting_after": -1}, {"stream": "true"}), + ({"extra_query": {"stream": "true", "starting_after": "abc"}}, {"stream": "true"}), + ), +) +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_malformed_resume_query_is_dropped_under_drop_params( + sync_mode: bool, resume_kwargs: dict[str, object], expected_query: dict[str, str] +) -> None: + upstream: Final = ( + httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3) + if expected_query + else httpx.Response(200, json=_stored_response("completed")) + ) + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock(return_value=upstream) + retrieved: Final = await _retrieve(sync_mode, drop_params=True, **_retrieve_kwargs("openai"), **resume_kwargs) + streamed_types: Final = None if isinstance(retrieved, ResponsesAPIResponse) else await _event_types(retrieved) + (call,) = route.calls + + assert dict(call.request.url.params) == expected_query + assert (streamed_types is not None) == bool(expected_query) + + _UNSUPPORTED_PROVIDERS: Final = ("xai", "manus", "volcengine")