diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c9b6fff3a00..2c76ff7a99b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1445,7 +1445,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 @@ -1468,6 +1468,28 @@ 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()) + + +def _abandoned_producer_failure(produce_task: "asyncio.Future[_LateResponseT]") -> Exception: + failure: Final = None if produce_task.cancelled() else produce_task.exception() + if isinstance(failure, Exception): + return failure + return HTTPException(status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail=LITELLM_CLIENT_DISCONNECT_DETAIL) + + 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): @@ -1484,7 +1506,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( @@ -1540,7 +1562,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. @@ -1551,6 +1577,12 @@ async def open_sse_before_first_byte( while not produce_task.done(): yield SSE_COMMENT_PING_BYTES await asyncio.wait((produce_task,), timeout=interval) + except BaseException: + await _abandon_late_producer(produce_task) + with anyio.CancelScope(shield=True): + await _sanitized_late_failure(_abandoned_producer_failure(produce_task), on_late_failure) + raise + try: try: produced: Final = produce_task.result() except Exception as exc: # noqa: BLE001 # the status line is already sent; surface it as a frame @@ -1568,18 +1600,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 1eb76f8f558..fcdab710206 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,17 +10,19 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from collections.abc import AsyncIterable, Callable, Mapping, Sequence +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 @@ -33,6 +35,8 @@ if TYPE_CHECKING: _JsonDict: TypeAlias = dict[str, object] _JsonList: TypeAlias = list[object] +CANCELLATION_POLL_SECONDS: Final = 1.0 + class _OutputItem(TypedDict, total=False): id: ReadOnly[str] @@ -81,6 +85,10 @@ class _StreamsBody(Protocol): def body_iterator(self) -> AsyncIterable[str | bytes | memoryview]: ... +class PollingStateReader(Protocol): + async def is_cancelled_or_gone(self, polling_id: str) -> bool: ... + + def _sse_frame_data(frame: str) -> str | None: return next((line[6:].strip() for line in split_sse_lines(frame) if line.startswith("data: ")), None) @@ -99,6 +107,43 @@ def detach_request_from_client(request: Request) -> Request: return Request(request.scope, _never_receive) +async def _reads_stopped(polling_id: str, polling_state: PollingStateReader) -> bool: + try: + 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 + + +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) + + +async def stream_until_cancelled( + polling_id: str, + stream: Coroutine[object, object, None], + polling_state: PollingStateReader, + poll_seconds: float = CANCELLATION_POLL_SECONDS, +) -> None: + streaming: Final = asyncio.create_task(stream) + 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: + streaming.cancel() + 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 or deleted", polling_id + ) + return + streaming.result() + + async def background_streaming_task( polling_id: str, data: dict[str, object], @@ -117,7 +162,51 @@ async def background_streaming_task( user_max_tokens: int | None, user_api_base: str | None, version: str | None, -): +) -> None: + await stream_until_cancelled( + polling_id, + _stream_into_polling_state( + polling_id=polling_id, + data=data, + polling_handler=polling_handler, + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + general_settings=general_settings, + llm_router=llm_router, + proxy_config=proxy_config, + proxy_logging_obj=proxy_logging_obj, + select_data_generator=select_data_generator, + 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, + ), + polling_handler, + ) + + +async def _stream_into_polling_state( + polling_id: str, + data: dict[str, object], + polling_handler: ResponsePollingHandler, + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth, + general_settings: dict[str, object], + llm_router: "Router | None", + proxy_config: "ProxyConfig", + proxy_logging_obj: "ProxyLogging", + select_data_generator: Callable[..., object] | None, + user_model: str | None, + user_temperature: float | None, + user_request_timeout: float | None, + user_max_tokens: int | None, + user_api_base: str | None, + version: str | None, +) -> None: """ Background task to stream response and update cache @@ -131,41 +220,53 @@ async def background_streaming_task( 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 @@ -229,158 +330,162 @@ async def background_streaming_task( ) 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 fe224a52cfa..7b67fcbc746 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -3,7 +3,7 @@ Response Polling Handler for Background Responses with Cache """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from typing import Final @@ -18,6 +18,34 @@ from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStatus _RESPONSE_POLLING_TARGET: Final = "response_polling" _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _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 = """ +if redis.call('EXISTS', KEYS[1]) == 0 then + return 0 +end +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: @@ -97,13 +125,13 @@ class ResponsePollingHandler: self, polling_id: str, status: ResponsesAPIStatus | None = None, - usage: dict | None = None, - error: dict | None = None, - incomplete_details: dict | None = None, - reasoning: dict | None = None, + usage: Mapping[str, object] | None = None, + error: Mapping[str, object] | None = None, + incomplete_details: Mapping[str, object] | None = None, + reasoning: Mapping[str, object] | None = None, tool_choice: object | None = None, - tools: list | None = None, - output: list | None = None, + tools: Sequence[object] | None = None, + output: Sequence[object] | None = None, # Additional ResponsesAPIResponse fields model: str | None = None, instructions: str | None = None, @@ -111,7 +139,7 @@ class ResponsePollingHandler: top_p: float | None = None, max_output_tokens: int | None = None, previous_response_id: str | None = None, - text: dict | None = None, + text: Mapping[str, object] | None = None, truncation: str | None = None, parallel_tool_calls: bool | None = None, user: str | None = None, @@ -212,12 +240,22 @@ class ResponsePollingHandler: if store is not None: state["store"] = store - # Update cache with configured TTL - await self.redis_cache.async_set_cache( - key=cache_key, - value=json.dumps(state), - ttl=self.ttl, - ) + 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=(_encoded_state(state), self.redis_cache.get_ttl(ttl=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) + return + if written != _WRITTEN: + verbose_proxy_logger.debug( + "Skipped the polling state write for %s: the response was cancelled or has expired", polling_id + ) + return output_count: Final = len(_JSON_LIST.validate_python(state.get("output", []))) verbose_proxy_logger.debug( @@ -238,6 +276,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 @@ -268,7 +314,7 @@ def should_use_polling_for_request( redis_cache, # RedisCache or None model: str, llm_router, # Router instance or None - native_background_mode: list[str] | None = None, # List of models that should use native background mode + native_background_mode: Sequence[str] | None = None, # List of models that should use native background mode ) -> bool: """ Determine if polling via cache should be used for a request. diff --git a/tests/integration/streaming/test_background_response_polling_cancel.py b/tests/integration/streaming/test_background_response_polling_cancel.py new file mode 100644 index 00000000000..88d89ab8bd5 --- /dev/null +++ b/tests/integration/streaming/test_background_response_polling_cancel.py @@ -0,0 +1,256 @@ +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import ExitStack, contextmanager +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, eventually, held, object_value, string_value +from integration._support.process import DB_PUSH, graceful_stop_seconds, owned_proxy +from integration._support.redis_process import owned_redis +from integration._support.responses_stream import OPENAI_TARGET, completed, created, delta, frame, response_object +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +_MODEL: Final = "background-polling-cancel" +_NO_DATABASE: Final = ("DATABASE_URL", "DATABASE_URL_READ_REPLICA") +_DELTAS: Final = 40 +_DELTA_PAUSE_SECONDS: Final = 0.25 +_STREAM_SECONDS: Final = _DELTAS * _DELTA_PAUSE_SECONDS +_OWNED_PROXY_SECONDS: Final = 2 * graceful_stop_seconds() + 240 +_MODEL_DISCOVERY: Final = ("GET", "/v1/models") +_NO_MODELS: Final = Reply(body=b'{"object":"list","data":[]}') +_POLLING_VIA_CACHE: Final = MappingProxyType({"responses": {"background_mode": {"polling_via_cache": "all"}}}) +_KEEPALIVE_SECONDS: Final = 0.05 + + +def _config(directory: Path, wire: Wire, litellm_settings: Mapping[str, JsonValue] = _POLLING_VIA_CACHE) -> Path: + path: Final = directory / f"background-polling-cancel-{uuid.uuid4().hex}.yaml" + path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": _MODEL, + "litellm_params": { + "model": "openai/gpt-5.4-mini", + "api_base": wire.url, + "api_key": "synthetic-openai-key", + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "coordination_redis": {"host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, + }, + "litellm_settings": dict(litellm_settings), + } + ) + ) + return path + + +@contextmanager +def _proxies_sharing_one_redis( + gateway: Gateway, directory: Path, config: Path, count: int, *, with_database: bool = False +) -> Iterator[tuple[Gateway, ...]]: + with owned_redis(directory / "redis") as cache, ExitStack() as proxies: + overrides: Final = { + "REDIS_HOST": cache.host, + "REDIS_PORT": str(cache.port), + "STORE_MODEL_IN_DB": "False", + "DISABLE_SCHEMA_UPDATE": str(not with_database).lower(), + } + yield tuple( + proxies.enter_context( + owned_proxy( + gateway, + directory / f"proxy-{index}", + overrides, + config=config, + database_setup=DB_PUSH if with_database else (), + remove_environment=() if with_database else _NO_DATABASE, + ) + ) + for index in range(count) + ) + + +def _serving(reply: Reply, requested: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if (request.method, request.target) == _MODEL_DISCOVERY: + return _NO_MODELS + assert request.method == "POST" and request.target.startswith(OPENAI_TARGET), request.target + requested.set() + return reply + + return respond + + +def _create_background_response(proxy: Gateway, key: str | None = None) -> str: + created_response: Final = proxy.request( + "POST", "/v1/responses", {"model": _MODEL, "input": uuid.uuid4().hex, "background": True}, key=key + ) + assert created_response.status_code == 200, created_response.text + body: Final = object_value(created_response.json()) + assert body["status"] == "queued", body + return string_value(body["id"]) + + +def _cancel(proxy: Gateway, polling_id: str) -> None: + cancelled: Final = proxy.request("POST", f"/v1/responses/{polling_id}/cancel") + assert cancelled.status_code == 200, cancelled.text + assert object_value(cancelled.json())["status"] == "cancelled", cancelled.text + + +def _status(proxy: Gateway, polling_id: str) -> JsonValue: + polled: Final = proxy.request("GET", f"/v1/responses/{polling_id}") + assert polled.status_code == 200, polled.text + return object_value(polled.json())["status"] + + +@pytest.mark.timeout(2 * _OWNED_PROXY_SECONDS + 3 * _STREAM_SECONDS + 60) +def test_a_cancel_sent_to_another_proxy_stops_the_background_stream_and_the_response_stays_cancelled( + 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, _config(tmp_path, wire), 2) as (streaming, cancelling), + ): + polling_id: Final = _create_background_response(streaming) + eventually(requested.is_set, lambda seen: seen, seconds=30) + _cancel(cancelling, polling_id) + eventually(lambda: wire.disconnected.qsize(), lambda dropped: dropped == 1, seconds=_STREAM_SECONDS) + held( + lambda: _status(cancelling, polling_id), + lambda status: status == "cancelled", + holding=_STREAM_SECONDS, + seconds=2 * _STREAM_SECONDS, + ) + + +@pytest.mark.timeout(_OWNED_PROXY_SECONDS + 60) +def test_a_background_response_cancelled_just_before_its_stream_finishes_stays_cancelled( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + requested: Final = threading.Event() + cancel_sent: Final = threading.Event() + held_stream: Final = Reply( + content_type="text/event-stream", + chunks=( + frame(created(identity)), + frame({"type": "response.in_progress", "response": response_object(identity, "in_progress")}), + frame(delta(identity, "done")), + frame(completed(identity, "done")), + ), + gate_after_first=cancel_sent, + ) + with ( + wire_server(_serving(held_stream, requested)) as wire, + _proxies_sharing_one_redis(gateway, tmp_path, _config(tmp_path, wire), 1) as (proxy,), + ): + polling_id: Final = _create_background_response(proxy) + eventually(requested.is_set, lambda seen: seen, seconds=30) + _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, _config(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, + ) + + +def _streamed_response_status(proxy: Gateway, key: str) -> int: + return proxy.request( + "POST", "/v1/responses", {"model": _MODEL, "input": uuid.uuid4().hex, "stream": True}, key=key + ).status_code + + +@pytest.mark.timeout(_OWNED_PROXY_SECONDS + 90) +def test_a_background_response_cancelled_during_keepalive_pings_frees_its_parallel_slot( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + requested: Final = threading.Event() + answer: Final = threading.Event() + quick_stream: Final = Reply( + content_type="text/event-stream", + chunks=(frame(created(identity)), frame(delta(identity, "done")), frame(completed(identity, "done"))), + ) + + def respond(request: Request) -> Reply: + if (request.method, request.target) == _MODEL_DISCOVERY: + return _NO_MODELS + assert request.method == "POST" and request.target.startswith(OPENAI_TARGET), request.target + if not requested.is_set(): + requested.set() + answer.wait() + return quick_stream + + with ( + wire_server(respond) as wire, + _proxies_sharing_one_redis( + gateway, + tmp_path, + _config(tmp_path, wire, {**_POLLING_VIA_CACHE, "sse_keepalive_ping_interval_seconds": _KEEPALIVE_SECONDS}), + 1, + with_database=True, + ) as (proxy,), + proxy.scenario() as scenario, + ): + try: + key: Final = scenario.key(max_parallel_requests=1) + polling_id: Final = _create_background_response(proxy, key) + eventually(requested.is_set, lambda seen: seen, seconds=30) + held(lambda: _status(proxy, polling_id), lambda status: status == "in_progress", holding=1) + assert _streamed_response_status(proxy, key) == 429 + _cancel(proxy, polling_id) + eventually( + lambda: _streamed_response_status(proxy, key), lambda status_code: status_code == 200, seconds=20 + ) + assert _status(proxy, polling_id) == "cancelled" + finally: + answer.set() diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index a50aa577a79..4427bb32eb3 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -23,7 +23,7 @@ import litellm from litellm import Router from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.litellm_core_utils import ptu_pricing from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling @@ -4840,6 +4840,83 @@ async def test_sse_keepalive_releases_parallel_capacity( assert handler._gauge_in_flight_from_cache_value(await cache.async_get_cache(counter)) == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("leaves_by", ["close", "cancel"]) +@pytest.mark.parametrize("preadmitted", [False, True], ids=["producer-admission", "existing-admission"]) +async def test_sse_keepalive_releases_parallel_capacity_when_the_reader_leaves_during_pings( + leaves_by: str, preadmitted: bool +) -> None: + from fastapi.responses import StreamingResponse + + from litellm.proxy.common_request_processing import open_sse_before_first_byte + from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES + + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + api_key: Final = hash_token("sk-keepalive-reader-leaves") + auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + counter: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + + async def request(call_id: str) -> None: + data: Final = {"model": "test-model", "litellm_call_id": call_id} + admitted: Final = asyncio.Event() + never_answers: Final = asyncio.Event() + upstream_cancelled: Final = asyncio.Event() + reported: Final = asyncio.Queue[Exception]() + + async def admit() -> None: + await handler.async_pre_call_hook( + user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion" + ) + + async def failure(exc: Exception) -> None: + await reported.put(exc) + await handler.async_post_call_failure_hook( + request_data=data, original_exception=exc, user_api_key_dict=auth + ) + + async def produce() -> StreamingResponse: + if not preadmitted: + await admit() + admitted.set() + try: + await never_answers.wait() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + return StreamingResponse(iter(())) + + if preadmitted: + await admit() + response: Final = await open_sse_before_first_byte(produce(), 0.001, on_late_failure=failure) + assert isinstance(response, StreamingResponse) + assert await anext(response.body_iterator) == SSE_COMMENT_PING_BYTES + await admitted.wait() + + async def read_until_cancelled() -> None: + async for _ in response.body_iterator: + continue + + if leaves_by == "close": + body: Final = response.body_iterator + assert isinstance(body, AsyncGenerator) + await body.aclose() + else: + reader: Final = asyncio.ensure_future(read_until_cancelled()) + await asyncio.sleep(0.01) + reader.cancel() + with pytest.raises(asyncio.CancelledError): + await reader + assert upstream_cancelled.is_set() + left: Final = reported.get_nowait() + assert isinstance(left, HTTPException) and left.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED + assert reported.empty() + + for index in range(3): + await contextvars.Context().run(asyncio.create_task, request(f"keepalive-left-{index}")) + assert handler._gauge_in_flight_from_cache_value(await cache.async_get_cache(counter)) == 0 + + @pytest.mark.asyncio async def test_sse_keepalive_rejects_overlap_and_reuses_completed_capacity() -> None: from fastapi.responses import StreamingResponse 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 087b273548a..38a6a11ef74 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -16,19 +16,69 @@ following the OpenAI Response API format. import asyncio import json +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Dict, Final, Optional +from typing import Dict, Final, NoReturn, Optional from unittest.mock import AsyncMock, Mock, patch import pytest from fastapi import Request from pydantic import ValidationError +from litellm.caching.base_cache import BaseCache from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs +class _RedisStandIn(BaseCache): + def __init__(self) -> None: + super().__init__() + self.stored: Dict[str, str] = {} + self.ttls: Dict[str, Optional[int]] = {} + + @classmethod + def holding(cls, state: Mapping[str, object], polling_id: str = "litellm_poll_test") -> "_RedisStandIn": + cache: Final = cls() + cache.stored[ResponsePollingHandler.get_cache_key(polling_id)] = json.dumps(json.dumps(state)) + return cache + + async def async_get_cache(self, key: str) -> object: + raw: Final = self.stored.get(key) + return None if raw is None else json.loads(raw) + + async def async_set_cache(self, key: str, value: object, ttl: Optional[int] = None) -> None: + self.stored[key] = json.dumps(value) + self.ttls[key] = ttl + + async def async_delete_cache(self, key: str) -> None: + self.stored.pop(key, None) + + async def async_set_cache_pipeline(self, cache_list: Sequence[tuple[str, object]], **kwargs: object) -> NoReturn: + raise NotImplementedError + + def async_register_script(self, script: str) -> Callable[..., Awaitable[int]]: + async def write_unless_cancelled(keys: Sequence[str], args: Sequence[object]) -> int: + 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 + + async def cancelled_or_gone(keys: Sequence[str], args: Sequence[object]) -> int: + 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): + def async_register_script(self, script: str) -> Callable[..., Awaitable[int]]: + async def write_fails(keys: Sequence[str], args: Sequence[object]) -> NoReturn: + raise ConnectionError("redis is unreachable") + + return write_fails + + class TestResponsePollingHandler: """Test cases for ResponsePollingHandler""" @@ -153,8 +203,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_changes_status_to_in_progress(self): """Test that update_state can change status to in_progress""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -171,18 +220,14 @@ class TestResponsePollingHandler: status="in_progress", ) - # Verify the update was saved - mock_redis.async_set_cache.assert_called_once() - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["status"] == "in_progress" @pytest.mark.asyncio async def test_update_state_replaces_full_output_list(self): """Test that update_state replaces the full output list""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -212,8 +257,7 @@ class TestResponsePollingHandler: output=new_output, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert len(stored["output"]) == 2 assert stored["output"][0]["id"] == "item_1" @@ -222,8 +266,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_usage(self): """Test that update_state correctly stores usage data""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -243,8 +286,7 @@ class TestResponsePollingHandler: usage=usage_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["status"] == "completed" assert stored["usage"] == usage_data @@ -252,8 +294,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_reasoning_tools_tool_choice(self): """Test that update_state stores reasoning, tools, and tool_choice from response.completed""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -279,8 +320,7 @@ class TestResponsePollingHandler: tools=tools_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["reasoning"] == reasoning_data assert stored["tool_choice"] == tool_choice_data @@ -289,8 +329,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_all_responses_api_fields(self): """Test that update_state stores all ResponsesAPIResponse fields from response.completed""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -324,8 +363,7 @@ class TestResponsePollingHandler: incomplete_details={"reason": "max_output_tokens"}, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") # Verify all fields are stored correctly assert stored["status"] == "completed" @@ -353,8 +391,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_preserves_existing_fields(self): """Test that update_state preserves fields not being updated""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -374,8 +411,7 @@ class TestResponsePollingHandler: status="completed", ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") # Verify existing fields are preserved assert stored["status"] == "completed" @@ -386,8 +422,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_error_sets_failed_status(self): """Test that providing an error automatically sets status to failed""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -410,8 +445,7 @@ class TestResponsePollingHandler: error=error_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["status"] == "failed" assert stored["error"] == error_data @@ -419,8 +453,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_incomplete_details(self): """Test that update_state stores incomplete_details""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -440,8 +473,7 @@ class TestResponsePollingHandler: incomplete_details=incomplete_details, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["status"] == "incomplete" assert stored["incomplete_details"] == incomplete_details @@ -460,8 +492,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_handles_missing_cached_state(self): """Test that update_state handles case when cached state doesn't exist""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = None # Cache miss + mock_redis = _RedisStandIn() handler = ResponsePollingHandler(redis_cache=mock_redis) @@ -471,8 +502,7 @@ class TestResponsePollingHandler: status="in_progress", ) - # Should not try to set cache if nothing was found - mock_redis.async_set_cache.assert_not_called() + assert mock_redis.stored == {} # ==================== Get State Tests ==================== @@ -517,13 +547,18 @@ class TestResponsePollingHandler: assert result is None + @pytest.mark.asyncio + async def test_a_response_never_reads_as_stopped_without_redis(self) -> None: + handler: Final = ResponsePollingHandler(redis_cache=None) + + assert await handler.is_cancelled_or_gone("litellm_poll_test") is False + # ==================== Cancel Polling Tests ==================== @pytest.mark.asyncio async def test_cancel_polling_updates_status_to_cancelled(self): """Test that cancel_polling sets status to cancelled""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -539,10 +574,92 @@ class TestResponsePollingHandler: assert result is True - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["status"] == "cancelled" + @pytest.mark.asyncio + async def test_an_update_after_a_cancel_leaves_the_response_cancelled(self) -> None: + mock_redis: Final = _RedisStandIn.holding( + { + "id": "litellm_poll_test", + "object": "response", + "status": "in_progress", + "output": [], + "created_at": 1234567890, + } + ) + handler: Final = ResponsePollingHandler(redis_cache=mock_redis) + + await handler.cancel_polling("litellm_poll_test") + await handler.update_state( + polling_id="litellm_poll_test", + status="completed", + output=[{"id": "item_1", "type": "message"}], + ) + + assert await handler.get_state("litellm_poll_test") == { + "id": "litellm_poll_test", + "object": "response", + "status": "cancelled", + "output": [], + "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: str | None, expected: bool) -> None: + handler: Final = 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) -> None: + state: Final = { + "id": "litellm_poll_test", + "object": "response", + "status": "in_progress", + "output": [], + "created_at": 1234567890, + } + handler: Final = ResponsePollingHandler(redis_cache=_UnreachableRedisStandIn.holding(state)) + + await handler.update_state(polling_id="litellm_poll_test", status="completed") + + assert await handler.get_state("litellm_poll_test") == state + + @pytest.mark.asyncio + async def test_a_queued_response_reads_back_after_an_update(self) -> None: + handler: Final = ResponsePollingHandler(redis_cache=_RedisStandIn()) + + queued: Final = await handler.create_initial_state( + polling_id="litellm_poll_test", request_data={"model": "gpt-4o", "metadata": {"trace": "t-1"}} + ) + await handler.update_state(polling_id="litellm_poll_test", status="in_progress", model="gpt-4o") + + assert await handler.get_state("litellm_poll_test") == { + **queued.model_dump(), + "status": "in_progress", + "model": "gpt-4o", + } + # ==================== Delete Polling Tests ==================== @pytest.mark.asyncio @@ -596,8 +713,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_uses_configured_ttl(self): """Test that update_state uses the configured TTL""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -614,8 +730,21 @@ class TestResponsePollingHandler: status="in_progress", ) - call_args = mock_redis.async_set_cache.call_args - assert call_args.kwargs["ttl"] == 1800 + assert mock_redis.ttls == {ResponsePollingHandler.get_cache_key("litellm_poll_test"): 1800} + + @pytest.mark.asyncio + @pytest.mark.parametrize("configured_ttl", [None, "1h"]) + async def test_an_unusable_configured_ttl_still_writes_the_state_with_the_cache_default( + self, configured_ttl: object + ) -> None: + cache: Final = _RedisStandIn.holding({"id": "litellm_poll_test", "status": "queued", "output": []}) + handler: Final = ResponsePollingHandler(redis_cache=cache, ttl=configured_ttl) + + await handler.update_state(polling_id="litellm_poll_test", status="in_progress") + + key: Final = ResponsePollingHandler.get_cache_key("litellm_poll_test") + assert json.loads(json.loads(cache.stored[key]))["status"] == "in_progress" + assert cache.ttls == {key: cache.default_ttl} class TestStreamingEventProcessing: @@ -1425,7 +1554,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, @@ -1436,6 +1565,13 @@ def _make_background_streaming_kwargs( ) +def _streaming_polling_handler(status: str = "in_progress") -> AsyncMock: + handler: Final = AsyncMock(spec=ResponsePollingHandler) + handler.get_state.return_value = {"status": status} + handler.is_cancelled_or_gone.return_value = status == "cancelled" + return handler + + @pytest.mark.xdist_group("heavy_imports") class TestBackgroundStreamingTerminalEvents: """ @@ -1470,7 +1606,7 @@ class TestBackgroundStreamingTerminalEvents: }, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_1", handler) with patch( @@ -1509,7 +1645,7 @@ class TestBackgroundStreamingTerminalEvents: mock_response = Mock() mock_response.body_iterator = _body_iterator() - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_named_event", handler) with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests @@ -1552,7 +1688,7 @@ class TestBackgroundStreamingTerminalEvents: }, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_2", handler) with patch( @@ -1591,7 +1727,7 @@ class TestBackgroundStreamingTerminalEvents: }, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_3", handler) with patch( @@ -1626,7 +1762,7 @@ class TestBackgroundStreamingTerminalEvents: }, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_4", handler) with patch( @@ -1655,7 +1791,7 @@ class TestBackgroundStreamingTerminalEvents: events: Final = [ {"type": "response.completed", "response": {"id": "resp_123", "status": "completed", "output": [item]}} ] - handler: Final = AsyncMock(spec=ResponsePollingHandler) + handler: Final = _streaming_polling_handler() with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" @@ -1693,7 +1829,7 @@ class TestBackgroundStreamingTerminalEvents: }, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_5", handler) with patch( @@ -1719,7 +1855,7 @@ class TestBackgroundStreamingTerminalEvents: {"type": "response.in_progress"}, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_6", handler) with patch( @@ -1767,7 +1903,7 @@ class TestBackgroundStreamingTerminalEvents: slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"] ) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_7", handler) kwargs["request"] = Request( { @@ -1791,6 +1927,331 @@ class TestBackgroundStreamingTerminalEvents: assert final_call.kwargs["usage"] == {"input_tokens": 13, "output_tokens": 10} +class _CancelledAfterReads: + def __init__(self, reads_before_cancel: int) -> None: + self.reads = 0 + self.reads_before_cancel = reads_before_cancel + + async def is_cancelled_or_gone(self, polling_id: str) -> bool: + self.reads += 1 + return self.reads > self.reads_before_cancel + + +class _UnreadableState: + def __init__(self) -> None: + self.reads = 0 + + async def is_cancelled_or_gone(self, polling_id: str) -> bool: + self.reads += 1 + raise ConnectionError("redis is unreachable") + + +async def _let_the_loop_run() -> None: + for _ in range(20): + await asyncio.sleep(0) + + +@pytest.mark.xdist_group("heavy_imports") +class TestStreamUntilCancelled: + @pytest.mark.asyncio + async def test_a_cancelled_state_stops_the_stream(self) -> None: + from litellm.proxy.response_polling.background_streaming import stream_until_cancelled + + state: Final = _CancelledAfterReads(reads_before_cancel=2) + stopped: Final = asyncio.Event() + + async def endless_stream() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + stopped.set() + raise + + await asyncio.wait_for( + stream_until_cancelled("litellm_poll_test", endless_stream(), state, poll_seconds=0), timeout=5 + ) + + assert stopped.is_set() + assert state.reads == 3 + + @pytest.mark.asyncio + async def test_failing_state_reads_leave_the_stream_running_to_its_end(self) -> None: + from litellm.proxy.response_polling.background_streaming import stream_until_cancelled + + state: Final = _UnreadableState() + finished: Final = asyncio.Event() + + async def stream_outliving_three_reads() -> None: + while state.reads < 3: + await asyncio.sleep(0) + finished.set() + + await asyncio.wait_for( + stream_until_cancelled("litellm_poll_test", stream_outliving_three_reads(), state, poll_seconds=0), + timeout=5, + ) + + assert finished.is_set() + + @pytest.mark.asyncio + async def test_a_stream_that_finishes_first_leaves_no_watcher_reading(self) -> None: + from litellm.proxy.response_polling.background_streaming import stream_until_cancelled + + state: Final = _CancelledAfterReads(reads_before_cancel=1_000_000) + + async def short_stream() -> None: + await _let_the_loop_run() + + await asyncio.wait_for( + stream_until_cancelled("litellm_poll_test", short_stream(), state, poll_seconds=0), timeout=5 + ) + reads_when_done: Final = state.reads + await _let_the_loop_run() + + assert reads_when_done > 0 + assert state.reads == reads_when_done + + @pytest.mark.asyncio + async def test_a_stream_error_reaches_the_caller(self) -> None: + from litellm.proxy.response_polling.background_streaming import stream_until_cancelled + + async def broken_stream() -> None: + raise ValueError("upstream broke") + + with pytest.raises(ValueError, match="upstream broke"): + await stream_until_cancelled( + "litellm_poll_test", broken_stream(), _CancelledAfterReads(1_000_000), poll_seconds=0 + ) + + @pytest.mark.asyncio + async def test_cancelling_the_caller_stops_both_the_stream_and_the_watcher(self) -> None: + from litellm.proxy.response_polling.background_streaming import stream_until_cancelled + + state: Final = _CancelledAfterReads(reads_before_cancel=1_000_000) + started: Final = asyncio.Event() + stopped: Final = asyncio.Event() + + async def endless_stream() -> None: + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + stopped.set() + raise + + caller: Final = asyncio.create_task( + stream_until_cancelled("litellm_poll_test", endless_stream(), state, poll_seconds=0) + ) + await started.wait() + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + reads_when_cancelled: Final = state.reads + await _let_the_loop_run() + + assert stopped.is_set() + assert state.reads == reads_when_cancelled + + @pytest.mark.asyncio + async def test_a_cancelled_background_stream_writes_no_state_after_it_stops(self) -> None: + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + stopped: Final = asyncio.Event() + + async def endless_body() -> AsyncIterator[str]: + yield f"data: {json.dumps({'type': 'response.in_progress'})}" + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + stopped.set() + raise + + endless_response: Final = Mock() + endless_response.body_iterator = endless_body() + handler: Final = _streaming_polling_handler(status="cancelled") + kwargs: Final = _make_background_streaming_kwargs("poll_cancelled", 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=endless_response) + await asyncio.wait_for(background_streaming_task(**kwargs), timeout=5) + + 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, + ) -> None: + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + stopped: Final = asyncio.Event() + + async def body_before_its_first_event() -> AsyncIterator[str]: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + stopped.set() + raise + yield "data: [DONE]" + + silent_response: Final = Mock() + silent_response.body_iterator = body_before_its_first_event() + handler: Final = _streaming_polling_handler(status="cancelled") + kwargs: Final = _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) -> None: + from fastapi import HTTPException + + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + upstream_cancelled: Final = asyncio.Event() + + async def upstream_that_never_answers(**_: object) -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + + handler: Final = _streaming_polling_handler(status="cancelled") + kwargs: Final = _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: Final = kwargs["proxy_logging_obj"].post_call_failure_hook + assert isinstance(report, AsyncMock) + report.assert_awaited_once() + assert report.await_args is not None + reported: Final = 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) -> None: + from starlette.responses import StreamingResponse + + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + closed: Final = asyncio.Event() + + async def body() -> AsyncIterator[str]: + 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: Final = _streaming_polling_handler(status="cancelled") + handler.update_state.side_effect = output_writes_block + kwargs: Final = _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: tuple[Mapping[str, object], ...], expected_content: list[dict[str, str]] +) -> None: + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + async def body() -> AsyncIterator[str]: + yield "data: {not json" + for event in events: + yield f"data: {json.dumps(event)}" + yield "data: [DONE]" + + streamed: Final = Mock() + streamed.body_iterator = body() + handler: Final = _streaming_polling_handler() + kwargs: Final = _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: Final = handler.update_state.await_args_list + written_outputs: Final = [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""" @@ -1851,8 +2312,7 @@ class TestEdgeCases: @pytest.mark.asyncio async def test_update_state_with_none_output_clears_output(self): """Test that output=[] explicitly sets empty output""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -1869,8 +2329,7 @@ class TestEdgeCases: output=[], # Explicitly set empty ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["output"] == [] diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fe6b894b368..18c34c32141 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -1,4 +1,5 @@ import asyncio +import collections.abc import contextvars import copy import datetime @@ -10656,6 +10657,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", [ @@ -10827,6 +10850,33 @@ async def test_a_late_failure_is_reported_to_the_failure_hook(): assert collected[-1] == b"data: [DONE]\n\n" +@pytest.mark.asyncio +async def test_a_failure_that_lands_as_the_reader_leaves_is_reported_instead_of_a_disconnect(): + may_fail: Final = asyncio.Event() + audited: Final = asyncio.Queue[Exception]() + + async def fails_between_pings() -> StreamingResponse: + await may_fail.wait() + raise HTTPException(status_code=429, detail="rate limited") + + async def record(exc: Exception) -> None: + await audited.put(exc) + + response: Final = await open_sse_before_first_byte( + fails_between_pings(), ping_interval_seconds=0.05, on_late_failure=record + ) + body: Final = response.body_iterator + assert isinstance(body, collections.abc.AsyncGenerator) + assert await body.__anext__() == TTFT_PING + may_fail.set() + await asyncio.sleep(0.01) + await body.aclose() + + reported: Final = audited.get_nowait() + assert isinstance(reported, HTTPException) and reported.status_code == 429 + assert audited.empty() + + @pytest.mark.asyncio async def test_a_failing_audit_hook_never_costs_the_client_its_error_frame(): async def slow_failure():