mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix background polling request being cancelled on disconnect
POST /v1/responses with background:true under polling_via_cache was returning empty results because the client disconnects immediately and _buffer_first_chunk_honoring_disconnect cancelled the upstream call. Pass request=None from the background task so the disconnect guard is skipped, and widen base_process_llm_request to accept a missing request. Fixes BerriAI#36275
This commit is contained in:
parent
ade805ef0c
commit
81a579a4f9
3 changed files with 41 additions and 10 deletions
|
|
@ -86,6 +86,12 @@ from litellm.types.utils import (
|
|||
# is off (the default).
|
||||
_DD_STREAMING_TRACE_ENABLED: Final = not isinstance(tracer, NullTracer)
|
||||
|
||||
# Frozen empty header mapping used when no FastAPI request is available
|
||||
# (e.g. detached background polling tasks).
|
||||
_EMPTY_REQUEST_HEADERS: Final[
|
||||
dict[str, str]
|
||||
] = {} # mutable-ok: immutable default never mutated; hook API requires dict[str, str]
|
||||
|
||||
|
||||
_CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInformation] = {
|
||||
"error_code": str(LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED),
|
||||
|
|
@ -1162,7 +1168,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
data=request_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_headers=dict(request.headers),
|
||||
request_headers=dict(request.headers) if request is not None else _EMPTY_REQUEST_HEADERS,
|
||||
)
|
||||
if callback_headers:
|
||||
custom_headers.update(callback_headers)
|
||||
|
|
@ -1171,7 +1177,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
async def common_processing_pre_call_logic(
|
||||
self,
|
||||
request: Request,
|
||||
request: Request | None,
|
||||
general_settings: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
|
|
@ -1378,7 +1384,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
if alias_target is not None:
|
||||
self.data["model"] = alias_target
|
||||
|
||||
self.data["litellm_call_id"] = request.headers.get("x-litellm-call-id", str(uuid.uuid4()))
|
||||
self.data["litellm_call_id"] = (
|
||||
request.headers.get("x-litellm-call-id", str(uuid.uuid4())) if request is not None else str(uuid.uuid4())
|
||||
)
|
||||
DDSpanTagger.tag_call_id(self.data.get("litellm_call_id"))
|
||||
DDSpanTagger.tag_request(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -1439,7 +1447,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
async def _pre_call_with_fallbacks(
|
||||
self,
|
||||
request: Request,
|
||||
request: Request | None,
|
||||
general_settings: dict,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -1622,7 +1630,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
async def base_process_llm_request(
|
||||
self,
|
||||
request: Request,
|
||||
request: Request | None,
|
||||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
route_type: Literal[
|
||||
|
|
@ -1873,7 +1881,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
data=self.data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_headers=dict(request.headers),
|
||||
request_headers=dict(request.headers) if request is not None else _EMPTY_REQUEST_HEADERS,
|
||||
)
|
||||
if callback_headers:
|
||||
custom_headers.update(callback_headers)
|
||||
|
|
@ -1958,7 +1966,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
custom_headers=custom_headers,
|
||||
request_headers=dict(request.headers),
|
||||
request_headers=dict(request.headers) if request is not None else _EMPTY_REQUEST_HEADERS,
|
||||
)
|
||||
if _early is not None:
|
||||
return _early
|
||||
|
|
@ -2050,7 +2058,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
custom_headers=_non_streaming_custom_headers,
|
||||
request_headers=dict(request.headers),
|
||||
request_headers=dict(request.headers) if request is not None else _EMPTY_REQUEST_HEADERS,
|
||||
)
|
||||
if _early is not None:
|
||||
return _early
|
||||
|
|
@ -2140,7 +2148,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
data=self.data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_headers=dict(request.headers),
|
||||
request_headers=dict(request.headers) if request is not None else _EMPTY_REQUEST_HEADERS,
|
||||
)
|
||||
if callback_headers:
|
||||
fastapi_response.headers.update(callback_headers)
|
||||
|
|
|
|||
|
|
@ -70,7 +70,7 @@ async def background_streaming_task(
|
|||
# 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=request,
|
||||
request=None,
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route_type="aresponses",
|
||||
|
|
|
|||
|
|
@ -1749,3 +1749,26 @@ class TestEdgeCases:
|
|||
stored = json.loads(call_args.kwargs["value"])
|
||||
|
||||
assert stored["output"] == []
|
||||
|
||||
|
||||
class TestBackgroundStreamingRequestNone:
|
||||
"""Regression tests for background polling passing request=None."""
|
||||
|
||||
def test_background_streaming_passes_none_request(self):
|
||||
"""background_streaming_task must pass request=None to base_process_llm_request.
|
||||
|
||||
The client disconnects immediately on background requests, so the
|
||||
disconnect buffer must not arm and cancel the upstream LLM call.
|
||||
"""
|
||||
import inspect
|
||||
import textwrap
|
||||
|
||||
from litellm.proxy.response_polling.background_streaming import (
|
||||
background_streaming_task,
|
||||
)
|
||||
|
||||
source = textwrap.dedent(inspect.getsource(background_streaming_task))
|
||||
assert "request=None," in source or "request = None" in source, (
|
||||
"background_streaming_task should pass request=None to base_process_llm_request"
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue