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", [