From cbb50bb37ef0b0ac10ac38ff62de8f289480e65e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:05:11 -0700 Subject: [PATCH] fix(responses): flush streaming cache write cancelled at event loop shutdown --- litellm/caching/caching_handler.py | 9 +-- litellm/responses/streaming_iterator.py | 11 +-- .../responses/test_streaming_iterator.py | 70 +++++++++++++++++++ 3 files changed, 82 insertions(+), 8 deletions(-) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 5a55318f8dc..8fe60876b4e 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -141,10 +141,11 @@ async def _complete_cache_write_despite_cancellation(write_factory: Callable[[], raise -def _create_cache_write_task(write_factory: Callable[[], Awaitable[None]]) -> None: +def create_cache_write_task(write_factory: Callable[[], Awaitable[None]]) -> "asyncio.Task[None]": task: Final = asyncio.create_task(_complete_cache_write_despite_cancellation(write_factory)) _PENDING_CACHE_WRITES.add(task) task.add_done_callback(_PENDING_CACHE_WRITES.discard) + return task def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None: @@ -1030,14 +1031,14 @@ class LLMCachingHandler: isinstance(result, EmbeddingResponse) and not isinstance(cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude. ): - _create_cache_write_task( + create_cache_write_task( lambda: cache.async_add_cache_pipeline( result, dynamic_cache_object=self.dual_cache, **new_kwargs ) ) else: result_json: Final = result.model_dump_json() - _create_cache_write_task( + create_cache_write_task( lambda: cache.async_add_cache( result_json, dynamic_cache_object=self.dual_cache, @@ -1045,7 +1046,7 @@ class LLMCachingHandler: ) ) else: - _create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs)) + create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs)) def sync_set_cache( self, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 5c0d6fc536e..368fd481e63 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -573,13 +573,16 @@ class BaseResponsesAPIStreamingIterator: ): return - if litellm.cache is None: + cache: Final = litellm.cache + if cache is None: return cached_response: Final = response_obj.model_dump_json() if is_async: - cache_write_task: Final = asyncio.create_task( - litellm.cache.async_add_cache( + from litellm.caching.caching_handler import create_cache_write_task + + cache_write_task: Final = create_cache_write_task( + lambda: cache.async_add_cache( cached_response, dynamic_cache_object=getattr(caching_handler, "dual_cache", None), **request_kwargs, @@ -592,7 +595,7 @@ class BaseResponsesAPIStreamingIterator: ) ) else: - litellm.cache.add_cache( + cache.add_cache( cached_response, dynamic_cache_object=getattr(caching_handler, "dual_cache", None), **request_kwargs, diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 5b0f40fdf27..38407c94fe7 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -235,3 +235,73 @@ def test_sync_transport_error_before_completed_event_raises(): with pytest.raises(httpx.ReadError): for _ in iterator: pass + + +def test_stream_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch): + """ + Regression test for LIT-6184 on the /v1/responses streaming surface: the + completed-stream cache write was dispatched as a bare fire-and-forget task, + so asyncio.run cancelled it at loop close before the write landed. The + write must survive loop shutdown just like the chat-completions one. + """ + import asyncio + from types import SimpleNamespace + + import litellm + from litellm.types.utils import CallTypes + + writes = [] + + class _SlowWriteCache: + async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs): + await asyncio.sleep(0.2) + writes.append(result) + + def add_cache(self, *args, **kwargs): + raise AssertionError("sync write must not run on the async path") + + caching_handler = SimpleNamespace( + request_kwargs={ + "model": "test-model", + "input": "hello", + "stream": True, + "caching": True, + "metadata": None, + "custom_llm_provider": "openai", + }, + preset_cache_key="responses-stream-cache-key", + original_function=litellm.aresponses, + dual_cache=None, + _should_store_result_in_cache=lambda original_function, kwargs: True, + ) + logging_obj = SimpleNamespace( + model_call_details={"litellm_params": {}}, + _llm_caching_handler=caching_handler, + ) + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=Mock(spec=BaseResponsesAPIConfig), + logging_obj=logging_obj, + request_data=caching_handler.request_kwargs, + call_type=CallTypes.aresponses.value, + ) + iterator.completed_response = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_lit6184", + created_at=int(datetime.now().timestamp()), + status="completed", + model="test-model", + object="response", + output=[], + ), + ) + monkeypatch.setattr(litellm, "cache", _SlowWriteCache()) + + async def _short_lived_script(): + iterator._persist_completed_response_to_cache(is_async=True) + + asyncio.run(_short_lived_script()) + + assert len(writes) == 1