From c4bcf3227477adb795f34021524c2e36efa4ee0b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:07:15 -0700 Subject: [PATCH] fix(responses): settle a cancelled background call and stop it on delete A cancel that lands before the provider answers now reports one client disconnect through the failure hooks, so the key's parallel slot and budget reservation come back and the spend logs get a failure row. The state guard and the watcher compare a status-first prefix with GETRANGE instead of decoding the stored response inside Redis, so each write and read costs the same whatever the response size. The watcher also stops when the response was deleted, and a cancelled stream closes its upstream response. The TTFT keepalive now cancels the in-flight call when its caller is cancelled before the first ping. --- litellm/proxy/common_request_processing.py | 38 +- .../response_polling/background_streaming.py | 367 +++++++++--------- .../proxy/response_polling/polling_handler.py | 36 +- ...test_background_response_polling_cancel.py | 32 ++ .../test_response_polling_handler.py | 217 ++++++++++- .../proxy/test_common_request_processing.py | 22 ++ 6 files changed, 506 insertions(+), 206 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9bb451495c8..cf411bbc2e1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1351,7 +1351,7 @@ def ttft_keepalive_interval(request_data: Mapping[str, object], llm_router: Rout return resolve_ttft_keepalive_interval(deployments, litellm.sse_keepalive_ping_interval_seconds) -async def _aclose_late_response(produced: Response) -> None: +async def aclose_streaming_response(produced: object) -> None: """Run the cleanup Starlette would have run, for a response it never called. Closing an already-closed async generator is a no-op, so this is safe to call @@ -1374,6 +1374,21 @@ async def _aclose_late_response(produced: Response) -> None: verbose_proxy_logger.debug("error closing relayed streaming generator: %s", exc) +async def _abandon_late_producer(produce_task: "asyncio.Future[_LateResponseT]") -> None: + if not produce_task.done(): + produce_task.cancel() + with anyio.CancelScope(shield=True): + with contextlib.suppress(BaseException): + await produce_task + elif not produce_task.cancelled(): + # The upstream may have answered while nobody was draining this + # relay, e.g. the client vanished first. Nothing else holds that + # response, so its stream only gets closed here. + with anyio.CancelScope(shield=True): + with contextlib.suppress(BaseException): + await aclose_streaming_response(produce_task.result()) + + async def _relay_late_response(produced: Response) -> AsyncGenerator[bytes, None]: """Replay a Response that was built after a keepalive had already opened the wire.""" if not isinstance(produced, StreamingResponse): @@ -1390,7 +1405,7 @@ async def _relay_late_response(produced: Response) -> AsyncGenerator[bytes, None # Starlette never called this response, so the cleanup its __call__ would # have run has to happen here or the upstream LLM connection leaks. with anyio.CancelScope(shield=True): - await _aclose_late_response(produced) + await aclose_streaming_response(produced) async def _sanitized_late_failure( @@ -1446,7 +1461,11 @@ async def open_sse_before_first_byte( # Share admission state with stream cleanup before the producer copies this context. get_or_create_request_stash() produce_task: Final = asyncio.ensure_future(produce_response) - await asyncio.wait((produce_task,), timeout=interval) + try: + await asyncio.wait((produce_task,), timeout=interval) + except BaseException: + await _abandon_late_producer(produce_task) + raise if produce_task.done(): # Fast path: the upstream answered inside one interval, so nothing was # written early and this is byte-identical to not being wrapped at all. @@ -1474,18 +1493,7 @@ async def open_sse_before_first_byte( async for chunk in _relay_late_response(produced): yield chunk finally: - if not produce_task.done(): - produce_task.cancel() - with anyio.CancelScope(shield=True): - with contextlib.suppress(BaseException): - await produce_task - elif not produce_task.cancelled(): - # The upstream may have answered while nobody was draining this - # relay, e.g. the client vanished first. Nothing else holds that - # response, so its stream only gets closed here. - with anyio.CancelScope(shield=True): - with contextlib.suppress(BaseException): - await _aclose_late_response(produce_task.result()) + await _abandon_late_producer(produce_task) verbose_proxy_logger.info( "no upstream response after %ss, opening the SSE response early and sending keepalives", interval diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index e9cad2dbe80..fcdab710206 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -13,14 +13,16 @@ import json from collections.abc import AsyncIterable, Callable, Coroutine, Mapping, Sequence from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable -from fastapi import Request, Response +import anyio +from fastapi import HTTPException, Request, Response from starlette.types import Message from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_CLIENT_DISCONNECT_DETAIL, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED 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 ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, aclose_streaming_response from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from litellm.types.llms.openai import ResponsesAPIStatus @@ -84,7 +86,7 @@ class _StreamsBody(Protocol): class PollingStateReader(Protocol): - async def get_state(self, polling_id: str) -> Mapping[str, object] | None: ... + async def is_cancelled_or_gone(self, polling_id: str) -> bool: ... def _sse_frame_data(frame: str) -> str | None: @@ -105,19 +107,18 @@ def detach_request_from_client(request: Request) -> Request: return Request(request.scope, _never_receive) -async def _reads_cancelled(polling_id: str, polling_state: PollingStateReader) -> bool: +async def _reads_stopped(polling_id: str, polling_state: PollingStateReader) -> bool: try: - state: Final = await polling_state.get_state(polling_id) + return await polling_state.is_cancelled_or_gone(polling_id) except Exception as error: # noqa: BLE001 # a failed read must never stop the upstream call verbose_proxy_logger.warning( "Could not read the polling state of %s, so it keeps streaming: %s", polling_id, error ) return False - return state is not None and state.get("status") == "cancelled" -async def _wait_until_cancelled(polling_id: str, polling_state: PollingStateReader, poll_seconds: float) -> None: - while not await _reads_cancelled(polling_id, polling_state): +async def _wait_until_stopped(polling_id: str, polling_state: PollingStateReader, poll_seconds: float) -> None: + while not await _reads_stopped(polling_id, polling_state): await asyncio.sleep(poll_seconds) @@ -128,7 +129,7 @@ async def stream_until_cancelled( poll_seconds: float = CANCELLATION_POLL_SECONDS, ) -> None: streaming: Final = asyncio.create_task(stream) - watching: Final = asyncio.create_task(_wait_until_cancelled(polling_id, polling_state, poll_seconds)) + watching: Final = asyncio.create_task(_wait_until_stopped(polling_id, polling_state, poll_seconds)) try: await asyncio.wait((streaming, watching), return_when=asyncio.FIRST_COMPLETED) finally: @@ -136,7 +137,9 @@ async def stream_until_cancelled( watching.cancel() await asyncio.gather(streaming, watching, return_exceptions=True) if streaming.cancelled(): - verbose_proxy_logger.info("Stopped the background stream for %s because it was cancelled", polling_id) + verbose_proxy_logger.info( + "Stopped the background stream for %s because it was cancelled or deleted", polling_id + ) return streaming.result() @@ -217,41 +220,53 @@ async def _stream_into_polling_state( try: verbose_proxy_logger.info("Starting background streaming for %s", polling_id) - # Update status to in_progress (OpenAI format) - await polling_handler.update_state( - polling_id=polling_id, - status="in_progress", - ) + try: + # Update status to in_progress (OpenAI format) + await polling_handler.update_state( + polling_id=polling_id, + status="in_progress", + ) - # Force streaming mode and remove background flag - data["stream"] = True - data.pop("background", None) + # Force streaming mode and remove background flag + data["stream"] = True + data.pop("background", None) - # Create processor - processor: Final = ProxyBaseLLMRequestProcessing(data=data) + # Create processor + processor: Final = ProxyBaseLLMRequestProcessing(data=data) - # Make streaming request. - # Pre-call checks (rate limits, guardrails, budget) were already run - # before polling ID creation, so skip them here to avoid double-counting. - response: Final = await processor.base_process_llm_request( - request=detach_request_from_client(request), - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - route_type="aresponses", - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - model=None, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, - skip_pre_call_logic=True, - ) + # Make streaming request. + # Pre-call checks (rate limits, guardrails, budget) were already run + # before polling ID creation, so skip them here to avoid double-counting. + response: Final = await processor.base_process_llm_request( + request=detach_request_from_client(request), + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="aresponses", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + skip_pre_call_logic=True, + ) + except asyncio.CancelledError: + with anyio.CancelScope(shield=True): + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=HTTPException( + status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail=LITELLM_CLIENT_DISCONNECT_DETAIL + ), + request_data=data, + dispatched_call_type="aresponses", + ) + raise # Process streaming response following OpenAI events format # https://platform.openai.com/docs/api-reference/responses-streaming @@ -315,158 +330,162 @@ async def _stream_into_polling_state( ) if isinstance(response, _StreamsBody): - async for chunk in response.body_iterator: - # Parse chunk - if isinstance(chunk, bytes): - chunk = chunk.decode("utf-8") + try: + async for chunk in response.body_iterator: + # Parse chunk + if isinstance(chunk, bytes): + chunk = chunk.decode("utf-8") - if isinstance(chunk, str) and (chunk_data := _sse_frame_data(chunk)) is not None: - if chunk_data == "[DONE]": - break + if isinstance(chunk, str) and (chunk_data := _sse_frame_data(chunk)) is not None: + if chunk_data == "[DONE]": + break - try: - event: _StreamEvent = _StreamEventParser.parse(chunk_data) - event_type = event.get("type", "") + try: + event: _StreamEvent = _StreamEventParser.parse(chunk_data) + event_type = event.get("type", "") - # Process different event types based on OpenAI streaming spec - if event_type == "response.output_item.added": - # New output item added - item = event.get("item", {}) - item_id = item.get("id") - if item_id: - output_items[item_id] = item - state_dirty = True + # Process different event types based on OpenAI streaming spec + if event_type == "response.output_item.added": + # New output item added + item = event.get("item", {}) + item_id = item.get("id") + if item_id: + output_items[item_id] = item + state_dirty = True - elif event_type == "response.content_part.added": - # Content part added to an output item - item_id = event.get("item_id") - content_part = event.get("part", {}) + elif event_type == "response.content_part.added": + # Content part added to an output item + item_id = event.get("item_id") + content_part = event.get("part", {}) - if item_id and item_id in output_items: - # Update the output item with new content - added_item = output_items[item_id] - output_items[item_id] = { - **added_item, - "content": (*added_item.get("content", ()), content_part), - } - state_dirty = True + if item_id and item_id in output_items: + # Update the output item with new content + added_item = output_items[item_id] + output_items[item_id] = { + **added_item, + "content": (*added_item.get("content", ()), content_part), + } + state_dirty = True - elif event_type == "response.output_text.delta": - # Text delta - accumulate text content - # https://platform.openai.com/docs/api-reference/responses-streaming/response-text-delta - item_id = event.get("item_id") - content_index = event.get("content_index", 0) - delta = event.get("delta", "") + elif event_type == "response.output_text.delta": + # Text delta - accumulate text content + # https://platform.openai.com/docs/api-reference/responses-streaming/response-text-delta + item_id = event.get("item_id") + content_index = event.get("content_index", 0) + delta = event.get("delta", "") - if item_id and item_id in output_items: - # Accumulate text delta - key = (item_id, content_index) - if key not in accumulated_text: - accumulated_text[key] = "" - accumulated_text[key] += delta + if item_id and item_id in output_items: + # Accumulate text delta + key = (item_id, content_index) + if key not in accumulated_text: + accumulated_text[key] = "" + accumulated_text[key] += delta - # Update the content in output_items - delta_item = output_items[item_id] - if "content" in delta_item: - content_list = delta_item["content"] - if content_index < len(content_list): - content_entry = content_list[content_index] - if isinstance(content_entry, dict): - content_entry["text"] = accumulated_text[key] - state_dirty = True + # Update the content in output_items + delta_item = output_items[item_id] + if "content" in delta_item: + content_list = delta_item["content"] + if content_index < len(content_list): + content_entry = content_list[content_index] + if isinstance(content_entry, dict): + content_entry["text"] = accumulated_text[key] + state_dirty = True - elif event_type == "response.content_part.done": - # Content part completed - item_id = event.get("item_id") - content_part = event.get("part", {}) - content_index = event.get("content_index", 0) + elif event_type == "response.content_part.done": + # Content part completed + item_id = event.get("item_id") + content_part = event.get("part", {}) + content_index = event.get("content_index", 0) - if item_id and item_id in output_items: - # Update with final content from event - done_item = output_items[item_id] - if "content" in done_item: - content_list = done_item["content"] - if content_index < len(content_list): - output_items[item_id] = { - **done_item, - "content": tuple( - content_part if part_index == content_index else existing_part - for part_index, existing_part in enumerate(content_list) - ), - } - state_dirty = True + if item_id and item_id in output_items: + # Update with final content from event + done_item = output_items[item_id] + if "content" in done_item: + content_list = done_item["content"] + if content_index < len(content_list): + output_items[item_id] = { + **done_item, + "content": tuple( + content_part if part_index == content_index else existing_part + for part_index, existing_part in enumerate(content_list) + ), + } + state_dirty = True - elif event_type == "response.output_item.done": - # Output item completed - use final item data - item = event.get("item", {}) - item_id = item.get("id") - if item_id: - output_items[item_id] = item - state_dirty = True + elif event_type == "response.output_item.done": + # Output item completed - use final item data + item = event.get("item", {}) + item_id = item.get("id") + if item_id: + output_items[item_id] = item + state_dirty = True - elif event_type == "response.in_progress": - # Response is now in progress - # https://platform.openai.com/docs/api-reference/responses-streaming/response-in-progress - await polling_handler.update_state( - polling_id=polling_id, - status="in_progress", - ) + elif event_type == "response.in_progress": + # Response is now in progress + # https://platform.openai.com/docs/api-reference/responses-streaming/response-in-progress + await polling_handler.update_state( + polling_id=polling_id, + status="in_progress", + ) - elif event_type in ( - "response.completed", - "response.failed", - "response.incomplete", - "response.cancelled", - ): - # Terminal event - extract all ResponsesAPIResponse fields - # https://platform.openai.com/docs/api-reference/responses-streaming - response_data = event.get("response", {}) - terminal_status = response_data.get( - "status", - _event_to_status.get(event_type, "completed"), - ) + elif event_type in ( + "response.completed", + "response.failed", + "response.incomplete", + "response.cancelled", + ): + # Terminal event - extract all ResponsesAPIResponse fields + # https://platform.openai.com/docs/api-reference/responses-streaming + response_data = event.get("response", {}) + terminal_status = response_data.get( + "status", + _event_to_status.get(event_type, "completed"), + ) - # Extract error for failed and incomplete responses - if event_type == "response.failed" or event_type == "response.incomplete": - terminal_error = response_data.get("error") + # Extract error for failed and incomplete responses + if event_type == "response.failed" or event_type == "response.incomplete": + terminal_error = response_data.get("error") - # Core response fields - usage_data = response_data.get("usage") - reasoning_data = response_data.get("reasoning") - tool_choice_data = response_data.get("tool_choice") - tools_data = response_data.get("tools") + # Core response fields + usage_data = response_data.get("usage") + reasoning_data = response_data.get("reasoning") + tool_choice_data = response_data.get("tool_choice") + tools_data = response_data.get("tools") - # Additional ResponsesAPIResponse fields - model_data = response_data.get("model") - instructions_data = response_data.get("instructions") - temperature_data = response_data.get("temperature") - top_p_data = response_data.get("top_p") - max_output_tokens_data = response_data.get("max_output_tokens") - previous_response_id_data = response_data.get("previous_response_id") - text_data = response_data.get("text") - truncation_data = response_data.get("truncation") - parallel_tool_calls_data = response_data.get("parallel_tool_calls") - user_data = response_data.get("user") - store_data = response_data.get("store") - incomplete_details_data = response_data.get("incomplete_details") + # Additional ResponsesAPIResponse fields + model_data = response_data.get("model") + instructions_data = response_data.get("instructions") + temperature_data = response_data.get("temperature") + top_p_data = response_data.get("top_p") + max_output_tokens_data = response_data.get("max_output_tokens") + previous_response_id_data = response_data.get("previous_response_id") + text_data = response_data.get("text") + truncation_data = response_data.get("truncation") + parallel_tool_calls_data = response_data.get("parallel_tool_calls") + user_data = response_data.get("user") + store_data = response_data.get("store") + incomplete_details_data = response_data.get("incomplete_details") - # Also update output from final response if available - if "output" in response_data: - final_output = response_data.get("output", []) - for item in final_output: - item_id = item.get("id") - if item_id: - output_items[item_id] = item - state_dirty = True + # Also update output from final response if available + if "output" in response_data: + final_output = response_data.get("output", []) + for item in final_output: + item_id = item.get("id") + if item_id: + output_items[item_id] = item + state_dirty = True - # Flush state to Redis if interval elapsed - await flush_state_if_needed() + # Flush state to Redis if interval elapsed + await flush_state_if_needed() - except json.JSONDecodeError as e: - verbose_proxy_logger.warning("Failed to parse streaming chunk: %s", e) + except json.JSONDecodeError as e: + verbose_proxy_logger.warning("Failed to parse streaming chunk: %s", e) - # Final flush to ensure all accumulated state is saved - await flush_state_if_needed(force=True) + # Final flush to ensure all accumulated state is saved + await flush_state_if_needed(force=True) + finally: + with anyio.CancelScope(shield=True): + await aclose_streaming_response(response) # Use the terminal status from the stream, default to "completed" final_status: Final = terminal_status or "completed" diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index 9c8c8e27aed..e4ea6b873da 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -21,22 +21,32 @@ _JSON_LIST: Final = TypeAdapter(list[object]) _SCRIPT_REPLY: Final = TypeAdapter(int) _WRITTEN: Final = 1 + +def _encoded_state(state: Mapping[str, object]) -> str: + return json.dumps(json.dumps({"status": state["status"], **state})) + + +_CANCELLED_PREFIX: Final = _encoded_state({"status": "cancelled"}).removesuffix('}"') + _WRITE_UNLESS_CANCELLED_SCRIPT: Final = """ -local current = redis.call('GET', KEYS[1]) -if current == false then +if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end -local decoded, stored = pcall(cjson.decode, current) -if decoded and type(stored) == 'string' then - decoded, stored = pcall(cjson.decode, stored) -end -if decoded and type(stored) == 'table' and stored['status'] == 'cancelled' then +if redis.call('GETRANGE', KEYS[1], 0, string.len(ARGV[3]) - 1) == ARGV[3] then return 0 end redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) return 1 """ +_CANCELLED_OR_GONE_SCRIPT: Final = """ +local head = redis.call('GETRANGE', KEYS[1], 0, string.len(ARGV[1]) - 1) +if head == '' or head == ARGV[1] then + return 1 +end +return 0 +""" + class ResponsePollingHandler: """Handles polling-based responses with Redis cache""" @@ -233,7 +243,9 @@ class ResponsePollingHandler: write_unless_cancelled: Final = self.redis_cache.async_register_script(_WRITE_UNLESS_CANCELLED_SCRIPT) try: written: Final = _SCRIPT_REPLY.validate_python( - await write_unless_cancelled(keys=(cache_key,), args=(json.dumps(json.dumps(state)), int(self.ttl))) + await write_unless_cancelled( + keys=(cache_key,), args=(_encoded_state(state), int(self.ttl), _CANCELLED_PREFIX) + ) ) except Exception as write_error: # noqa: BLE001 # a Redis fault must not end the background stream verbose_proxy_logger.error("Failed to write the polling state for %s: %s", polling_id, write_error) @@ -263,6 +275,14 @@ class ResponsePollingHandler: return None + @with_service_target(_RESPONSE_POLLING_TARGET) + async def is_cancelled_or_gone(self, polling_id: str) -> bool: + if not self.redis_cache: + return False + cancelled_or_gone: Final = self.redis_cache.async_register_script(_CANCELLED_OR_GONE_SCRIPT) + reply: Final = await cancelled_or_gone(keys=(self.get_cache_key(polling_id),), args=(_CANCELLED_PREFIX,)) + return _SCRIPT_REPLY.validate_python(reply) == 1 + async def cancel_polling(self, polling_id: str) -> bool: """ Cancel a polling request diff --git a/tests/integration/streaming/test_background_response_polling_cancel.py b/tests/integration/streaming/test_background_response_polling_cancel.py index b1c125005be..ee59d0708d7 100644 --- a/tests/integration/streaming/test_background_response_polling_cancel.py +++ b/tests/integration/streaming/test_background_response_polling_cancel.py @@ -167,3 +167,35 @@ def test_a_background_response_cancelled_just_before_its_stream_finishes_stays_c _cancel(proxy, polling_id) cancel_sent.set() held(lambda: _status(proxy, polling_id), lambda status: status == "cancelled", holding=3) + + +@pytest.mark.timeout(2 * _OWNED_PROXY_SECONDS + 3 * _STREAM_SECONDS + 60) +def test_a_delete_sent_to_another_proxy_stops_the_background_stream_and_the_response_stays_gone( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + requested: Final = threading.Event() + slow_stream: Final = Reply( + content_type="text/event-stream", + chunks=( + frame(created(identity)), + *(frame(delta(identity, "word ")) for _ in range(_DELTAS)), + frame(completed(identity, "word " * _DELTAS)), + ), + pause_between_chunks=_DELTA_PAUSE_SECONDS, + ) + with ( + wire_server(_serving(slow_stream, requested)) as wire, + _proxies_sharing_one_redis(gateway, tmp_path, wire, 2) as (streaming, deleting), + ): + polling_id: Final = _create_background_response(streaming) + eventually(requested.is_set, lambda seen: seen, seconds=30) + deleted: Final = deleting.request("DELETE", f"/v1/responses/{polling_id}") + assert deleted.status_code == 200, deleted.text + eventually(lambda: wire.disconnected.qsize(), lambda dropped: dropped == 1, seconds=_STREAM_SECONDS) + held( + lambda: streaming.request("GET", f"/v1/responses/{polling_id}").status_code, + lambda status_code: status_code == 404, + holding=_STREAM_SECONDS, + seconds=2 * _STREAM_SECONDS, + ) diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 8d33a7c2761..9195447816c 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -46,19 +46,21 @@ class _RedisStandIn: self.stored[key] = json.dumps(value) self.ttls[key] = ttl - def _status(self, key: str) -> Any: - stored = json.loads(self.stored[key]) - return json.loads(stored)["status"] if isinstance(stored, str) else stored["status"] + async def async_delete_cache(self, key: str) -> None: + self.stored.pop(key, None) def async_register_script(self, script: str): async def write_unless_cancelled(keys, args): - if keys[0] not in self.stored or self._status(keys[0]) == "cancelled": + if keys[0] not in self.stored or self.stored[keys[0]].startswith(args[2]): return 0 self.stored[keys[0]] = args[0] self.ttls[keys[0]] = int(args[1]) return 1 - return write_unless_cancelled + async def cancelled_or_gone(keys, args): + return int(self.stored.get(keys[0], "")[: len(args[0])] in ("", args[0])) + + return write_unless_cancelled if "'SET'" in script else cancelled_or_gone class _UnreachableRedisStandIn(_RedisStandIn): @@ -537,6 +539,12 @@ class TestResponsePollingHandler: assert result is None + @pytest.mark.asyncio + async def test_a_response_never_reads_as_stopped_without_redis(self): + handler = ResponsePollingHandler(redis_cache=None) + + assert await handler.is_cancelled_or_gone("litellm_poll_test") is False + # ==================== Cancel Polling Tests ==================== @pytest.mark.asyncio @@ -589,6 +597,31 @@ class TestResponsePollingHandler: "created_at": 1234567890, } + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("stop", "expected"), + [ + (None, False), + ("cancel", True), + ("delete", True), + ], + ) + async def test_a_cancelled_or_deleted_response_reads_as_stopped(self, stop, expected): + handler = ResponsePollingHandler( + redis_cache=_RedisStandIn.holding( + {"id": "litellm_poll_test", "object": "response", "output": [], "status": "in_progress"} + ) + ) + await handler.update_state(polling_id="litellm_poll_test", output=[{"id": "item_1"}]) + + if stop == "cancel": + await handler.cancel_polling("litellm_poll_test") + if stop == "delete": + await handler.delete_polling("litellm_poll_test") + + assert await handler.is_cancelled_or_gone("litellm_poll_test") is expected + assert await handler.is_cancelled_or_gone("litellm_poll_never_created") is True + @pytest.mark.asyncio async def test_a_failed_state_write_keeps_the_stored_state_and_does_not_raise(self): state = { @@ -1499,7 +1532,7 @@ def _make_background_streaming_kwargs( general_settings={}, llm_router=None, proxy_config=Mock(), - proxy_logging_obj=Mock(), + proxy_logging_obj=AsyncMock(), select_data_generator=Mock(), user_model=None, user_temperature=None, @@ -1513,6 +1546,7 @@ def _make_background_streaming_kwargs( def _streaming_polling_handler(status: str = "in_progress") -> AsyncMock: handler = AsyncMock(spec=ResponsePollingHandler) handler.get_state.return_value = {"status": status} + handler.is_cancelled_or_gone.return_value = status == "cancelled" return handler @@ -1876,16 +1910,16 @@ class _CancelledAfterReads: self.reads = 0 self.reads_before_cancel = reads_before_cancel - async def get_state(self, polling_id: str) -> Dict[str, Any]: + async def is_cancelled_or_gone(self, polling_id: str) -> bool: self.reads += 1 - return {"id": polling_id, "status": "cancelled" if self.reads > self.reads_before_cancel else "in_progress"} + return self.reads > self.reads_before_cancel class _UnreadableState: def __init__(self) -> None: self.reads = 0 - async def get_state(self, polling_id: str) -> Dict[str, Any]: + async def is_cancelled_or_gone(self, polling_id: str) -> bool: self.reads += 1 raise ConnectionError("redis is unreachable") @@ -2023,7 +2057,172 @@ class TestStreamUntilCancelled: assert stopped.is_set() assert {call.kwargs.get("status") for call in handler.update_state.await_args_list} <= {None, "in_progress"} + kwargs["proxy_logging_obj"].post_call_failure_hook.assert_not_awaited() + @pytest.mark.asyncio + async def test_a_cancel_while_the_stream_waits_for_its_first_event_leaves_the_failure_report_to_the_stream(self): + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + stopped = asyncio.Event() + + async def body_before_its_first_event(): + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + stopped.set() + raise + yield "data: [DONE]" + + silent_response = Mock() + silent_response.body_iterator = body_before_its_first_event() + handler = _streaming_polling_handler(status="cancelled") + kwargs = _make_background_streaming_kwargs("poll_cancelled_before_first_event", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=silent_response) + await asyncio.wait_for(background_streaming_task(**kwargs), timeout=5) + + assert stopped.is_set() + kwargs["proxy_logging_obj"].post_call_failure_hook.assert_not_awaited() + + @pytest.mark.asyncio + async def test_a_cancel_before_the_upstream_answers_reports_one_client_disconnect(self): + from fastapi import HTTPException + + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + upstream_cancelled = asyncio.Event() + + async def upstream_that_never_answers(**_: object) -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + + handler = _streaming_polling_handler(status="cancelled") + kwargs = _make_background_streaming_kwargs("poll_cancelled_before_upstream", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = upstream_that_never_answers + await asyncio.wait_for(background_streaming_task(**kwargs), timeout=5) + + report = kwargs["proxy_logging_obj"].post_call_failure_hook + report.assert_awaited_once() + reported = report.await_args.kwargs + assert upstream_cancelled.is_set() + assert isinstance(reported["original_exception"], HTTPException) + assert reported["original_exception"].status_code == 499 + assert reported["request_data"] is kwargs["data"] + assert reported["user_api_key_dict"] is kwargs["user_api_key_dict"] + assert reported["dispatched_call_type"] == "aresponses" + + @pytest.mark.asyncio + async def test_a_cancel_during_a_blocked_state_write_closes_the_upstream_stream(self): + from starlette.responses import StreamingResponse + + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + closed = asyncio.Event() + + async def body(): + try: + yield f"data: {json.dumps({'type': 'response.output_item.added', 'item': {'id': 'item_1'}})}" + yield "data: [DONE]" + yield f"data: {json.dumps({'type': 'response.in_progress'})}" + finally: + closed.set() + + async def output_writes_block(**fields: object) -> None: + if "output" in fields: + await asyncio.Event().wait() + + handler = _streaming_polling_handler(status="cancelled") + handler.update_state.side_effect = output_writes_block + kwargs = _make_background_streaming_kwargs("poll_cancelled_mid_write", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=StreamingResponse(body())) + await asyncio.wait_for(background_streaming_task(**kwargs), timeout=5) + + assert closed.is_set() + assert {call.kwargs.get("status") for call in handler.update_state.await_args_list} <= {None, "in_progress"} + + + +_MESSAGE_ADDED: Final = { + "type": "response.output_item.added", + "item": {"id": "msg_1", "type": "message", "role": "assistant", "content": []}, +} +_PART_ADDED: Final = { + "type": "response.content_part.added", + "item_id": "msg_1", + "content_index": 0, + "part": {"type": "output_text", "text": ""}, +} +_DELTAS: Final = tuple( + {"type": "response.output_text.delta", "item_id": "msg_1", "content_index": 0, "delta": delta} + for delta in ("Hel", "lo") +) +_PART_DONE: Final = { + "type": "response.content_part.done", + "item_id": "msg_1", + "content_index": 0, + "part": {"type": "output_text", "text": "Hello!"}, +} +_MESSAGE_DONE: Final = { + "type": "response.output_item.done", + "item": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello, world"}], + }, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("events", "expected_content"), + [ + ((_MESSAGE_ADDED, _PART_ADDED, *_DELTAS), [{"type": "output_text", "text": "Hello"}]), + ((_MESSAGE_ADDED, _PART_ADDED, *_DELTAS, _PART_DONE), [{"type": "output_text", "text": "Hello!"}]), + ( + (_MESSAGE_ADDED, _PART_ADDED, *_DELTAS, _PART_DONE, _MESSAGE_DONE), + [{"type": "output_text", "text": "Hello, world"}], + ), + ], +) +async def test_a_background_stream_writes_the_message_its_events_build(events, expected_content): + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + async def body(): + yield "data: {not json" + for event in events: + yield f"data: {json.dumps(event)}" + yield "data: [DONE]" + + streamed = Mock() + streamed.body_iterator = body() + handler = _streaming_polling_handler() + kwargs = _make_background_streaming_kwargs("poll_message", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=streamed) + await background_streaming_task(**kwargs) + + writes = handler.update_state.await_args_list + written_outputs = [call.kwargs["output"] for call in writes if "output" in call.kwargs] + assert [list(item["content"]) for item in written_outputs[-1]] == [expected_content] + assert handler.update_state.await_args_list[-1].kwargs["status"] == "completed" class TestEdgeCases: """Test edge cases and error scenarios""" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fdcbcbfcd05..078990c7d98 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -10587,6 +10587,28 @@ async def test_ttft_keepalive_cancels_the_in_flight_call_when_the_client_gives_u assert upstream_cancelled.is_set() +@pytest.mark.asyncio +async def test_ttft_keepalive_cancels_the_in_flight_call_when_its_caller_is_cancelled_before_the_first_ping(): + upstream_started = asyncio.Event() + upstream_cancelled = asyncio.Event() + + async def never_answers(): + upstream_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + + caller = asyncio.create_task(open_sse_before_first_byte(never_answers(), ping_interval_seconds=60)) + await upstream_started.wait() + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + + assert upstream_cancelled.is_set() + + @pytest.mark.parametrize( "request_data, global_interval, expected", [