fix(responses): keep deployment stream off retrieves and reject oversized cursors

This commit is contained in:
mateo-berri 2026-10-28 19:55:45 -07:00
parent 9d6f615392
commit 76bc3c39f3
3 changed files with 66 additions and 3 deletions

View file

@ -7,6 +7,7 @@ from pydantic import TypeAdapter
RETRIEVE_STREAM_QUERY_KEYS: Final = frozenset({"stream", "starting_after"})
RETRIEVE_INCLUDE_QUERY_KEYS: Final = frozenset({"include", "include[]", "include_obfuscation"})
_QUERY_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_MAX_STARTING_AFTER_DIGITS: Final = len(str(2**63 - 1))
@dataclass(frozen=True, slots=True)
@ -43,7 +44,7 @@ def _parsed_starting_after(value: object) -> int | None | InvalidRetrieveStreamQ
return None
case int() if not isinstance(value, bool) and value >= 0:
return value
case str() if value.isascii() and value.isdigit():
case str() if value.isascii() and value.isdigit() and len(value) <= _MAX_STARTING_AFTER_DIGITS:
return int(value)
case _:
return InvalidRetrieveStreamQuery(
@ -72,3 +73,11 @@ def requests_streamed_retrieval(request: Mapping[str, object]) -> bool:
and isinstance(extra_query, Mapping)
and _parsed_stream(_QUERY_ADAPTER.validate_python(extra_query).get("stream")) is True
)
def without_retrieve_stream_defaults(
deployment_params: Mapping[str, object], request: Mapping[str, object]
) -> Mapping[str, object]:
if "response_id" not in request:
return deployment_params
return {key: value for key, value in deployment_params.items() if key not in RETRIEVE_STREAM_QUERY_KEYS}

View file

@ -148,7 +148,7 @@ from litellm.llms.openai_like.model_info import (
get_openai_compatible_model_info,
)
from litellm.responses.bridge_check import keep_registered_mode
from litellm.responses.retrieve_stream_query import requests_streamed_retrieval
from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults
from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.router_strategy.complexity_router.context_compaction import (
@ -6372,7 +6372,7 @@ class Router:
generic_function_name: Final[str] = getattr(original_generic_function, "__name__", "")
response_kwargs: Final = {
**data,
**without_retrieve_stream_defaults(data, kwargs),
"caching": self.cache_responses,
**kwargs,
"model": model_name,

View file

@ -412,6 +412,7 @@ async def test_a_resume_that_times_out_raises_a_timeout(sync_mode: bool) -> None
{"stream": True, "starting_after": -1},
{"extra_query": {"stream": "true", "starting_after": "abc"}},
{"extra_query": {"stream": "1"}},
{"extra_query": {"stream": "true", "starting_after": "9" * 5000}},
),
)
@pytest.mark.parametrize("sync_mode", (True, False))
@ -435,6 +436,7 @@ async def test_a_malformed_resume_query_is_rejected_before_any_upstream_call(
({"extra_query": {"stream": "1"}}, {}),
({"stream": True, "starting_after": -1}, {"stream": "true"}),
({"extra_query": {"stream": "true", "starting_after": "abc"}}, {"stream": "true"}),
({"extra_query": {"stream": "true", "starting_after": "9" * 5000}}, {"stream": "true"}),
),
)
@pytest.mark.parametrize("sync_mode", (True, False))
@ -538,6 +540,58 @@ async def test_a_routed_resume_requested_in_extra_query_gets_the_deployment_stre
assert call.request.extensions["timeout"]["read"] == 60
@pytest.mark.asyncio
async def test_a_routed_plain_retrieve_on_a_deployment_that_forces_stream_answers_the_response_object() -> None:
router: Final = litellm.Router(
model_list=[
{
"model_name": "background-model",
"litellm_params": {
"model": "openai/gpt-5.5",
"api_key": "sk-test",
"api_base": "https://api.openai.com/v1",
"stream": True,
},
}
]
)
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 router.aget_responses(model="background-model", response_id=_STORED_RESPONSE_ID)
(call,) = route.calls
assert isinstance(retrieved, ResponsesAPIResponse)
assert dict(call.request.url.params) == {}
@pytest.mark.asyncio
async def test_a_routed_create_on_a_deployment_that_forces_stream_still_streams() -> None:
router: Final = litellm.Router(
model_list=[
{
"model_name": "background-model",
"litellm_params": {
"model": "openai/gpt-5.5",
"api_key": "sk-test",
"api_base": "https://api.openai.com/v1",
"stream": True,
},
}
]
)
with respx.mock() as mock_router:
route: Final = mock_router.post("https://api.openai.com/v1/responses").mock(
return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3)
)
stream: Final = await router.aresponses(model="background-model", input="hi", background=True)
event_types: Final = await _event_types(stream)
(call,) = route.calls
assert event_types == ("response.output_text.delta", "response.completed")
assert json.loads(call.request.content)["stream"] is True
class _SuccessLog(CustomLogger):
def __init__(self) -> None:
super().__init__()