mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
test(responses): cover the polling disconnect regression through background_streaming_task
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus
parent
9c8caae1e4
commit
e589a0ed82
5 changed files with 59 additions and 54 deletions
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -155,7 +155,6 @@ jobs:
|
|||
tests/test_litellm/proxy/vector_store_files_endpoints
|
||||
tests/test_litellm/proxy/video_endpoints
|
||||
tests/test_litellm/proxy/response_api_endpoints
|
||||
tests/test_litellm/proxy/response_polling
|
||||
tests/test_litellm/proxy/image_endpoints
|
||||
tests/test_litellm/proxy/ocr_endpoints
|
||||
tests/test_litellm/proxy/vector_store_endpoints
|
||||
|
|
|
|||
2
Makefile
2
Makefile
|
|
@ -298,7 +298,7 @@ test-unit-proxy-core: install-test-deps
|
|||
$(UV_RUN) pytest tests/test_litellm/proxy/auth tests/test_litellm/proxy/client tests/test_litellm/proxy/db tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine --tb=short -vv -n 4 --durations=20
|
||||
|
||||
test-unit-proxy-misc: install-test-deps
|
||||
$(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/response_polling tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20
|
||||
$(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20
|
||||
|
||||
test-unit-integrations: install-test-deps
|
||||
$(UV_RUN) pytest tests/test_litellm/integrations --tb=short -vv -n 4 --durations=20
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ These tests ensure the polling handler correctly manages response state
|
|||
following the OpenAI Response API format.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, Optional
|
||||
|
|
@ -1663,6 +1664,63 @@ class TestBackgroundStreamingTerminalEvents:
|
|||
final_call = handler.update_state.call_args_list[-1]
|
||||
assert final_call.kwargs["status"] == "completed"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_polling_client_disconnect_does_not_cancel_upstream_call(self):
|
||||
"""The polling client hangs up right after getting its polling id. The detached task
|
||||
must still stream the upstream response through the client-disconnect guards."""
|
||||
from litellm.proxy.common_request_processing import create_response
|
||||
from litellm.proxy.response_polling.background_streaming import (
|
||||
background_streaming_task,
|
||||
)
|
||||
|
||||
async def client_already_left():
|
||||
return {"type": "http.disconnect"}
|
||||
|
||||
async def slow_upstream_stream():
|
||||
await asyncio.sleep(0.05)
|
||||
for event in (
|
||||
{"type": "response.in_progress"},
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_123",
|
||||
"status": "completed",
|
||||
"usage": {"input_tokens": 13, "output_tokens": 10},
|
||||
"model": "gpt-4o",
|
||||
"output": [{"id": "item_1", "type": "message"}],
|
||||
},
|
||||
},
|
||||
):
|
||||
yield f"data: {json.dumps(event)}\n\n"
|
||||
|
||||
async def upstream_call_behind_disconnect_guard(**kwargs):
|
||||
return await create_response(
|
||||
slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"]
|
||||
)
|
||||
|
||||
handler = AsyncMock(spec=ResponsePollingHandler)
|
||||
kwargs = _make_background_streaming_kwargs("poll_7", handler)
|
||||
kwargs["request"] = Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/v1/responses",
|
||||
"headers": [(b"x-litellm-call-id", b"call-123")],
|
||||
"query_string": b"",
|
||||
},
|
||||
client_already_left,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing"
|
||||
) as MockProcessor:
|
||||
MockProcessor.return_value.base_process_llm_request = upstream_call_behind_disconnect_guard
|
||||
await background_streaming_task(**kwargs)
|
||||
|
||||
final_call = handler.update_state.call_args_list[-1]
|
||||
assert final_call.kwargs["status"] == "completed"
|
||||
assert final_call.kwargs["usage"] == {"input_tokens": 13, "output_tokens": 10}
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Test edge cases and error scenarios"""
|
||||
|
|
|
|||
|
|
@ -1,52 +0,0 @@
|
|||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from fastapi import Request
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
from starlette.types import Message
|
||||
|
||||
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
|
||||
from litellm.proxy.common_request_processing import create_response
|
||||
from litellm.proxy.response_polling.background_streaming import detach_request_from_client
|
||||
|
||||
|
||||
def _request_whose_client_already_left() -> Request:
|
||||
async def receive() -> Message:
|
||||
return {"type": "http.disconnect"}
|
||||
|
||||
scope: Final = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/v1/responses",
|
||||
"headers": [(b"x-litellm-call-id", b"call-123")],
|
||||
"query_string": b"",
|
||||
}
|
||||
return Request(scope, receive)
|
||||
|
||||
|
||||
async def _slow_first_chunk() -> AsyncGenerator[str, None]:
|
||||
await asyncio.sleep(0.05)
|
||||
yield 'data: {"type": "response.created"}\n\n'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_detached_request_survives_client_disconnect_before_first_chunk():
|
||||
original: Final = _request_whose_client_already_left()
|
||||
|
||||
cancelled: Final = await create_response(_slow_first_chunk(), "text/event-stream", {}, request=original)
|
||||
assert isinstance(cancelled, JSONResponse)
|
||||
assert cancelled.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
|
||||
|
||||
detached: Final = detach_request_from_client(original)
|
||||
kept_alive: Final = await create_response(_slow_first_chunk(), "text/event-stream", {}, request=detached)
|
||||
assert isinstance(kept_alive, StreamingResponse)
|
||||
assert kept_alive.status_code == 200
|
||||
|
||||
|
||||
def test_detached_request_keeps_scope():
|
||||
original: Final = _request_whose_client_already_left()
|
||||
detached: Final = detach_request_from_client(original)
|
||||
assert detached.headers["x-litellm-call-id"] == "call-123"
|
||||
assert detached.scope is original.scope
|
||||
Loading…
Add table
Reference in a new issue