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