scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 245

Merge a9d3c7cf32 into afec145edc

Browse source
This commit is contained in:
tavren27 2026-10-31 05:59:10 +00:00 • committed by GitHub
commit a14151c452
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 1274 additions and 260 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/proxy/common_request_processing.py
View file

@ -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

441
litellm/proxy/response_polling/background_streaming.py
View file

@ -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"

76
litellm/proxy/response_polling/polling_handler.py
View file

@ -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.

256
tests/integration/streaming/test_background_response_polling_cancel.py Normal file
View file

@ -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()

79
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -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

581
tests/unit/proxy/response_polling/test_response_polling_handler.py
View file

@ -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"] == []

50
tests/unit/proxy/test_common_request_processing.py
View file

@ -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():