fix(proxy): keep the reservation when a disconnect happens while provider output is held back

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py

@@ -207,6 +207,11 @@ class AgenticAnthropicStreamingIterator:
         self._replay_index = 0
         self._error_emitted = False
 
+    @property
+    def has_buffered_provider_output(self) -> bool:
+        """Whether provider output was received but withheld from the client behind keepalive pings."""
+        return self._hold_back and bool(self._collected_bytes)
+
     def __aiter__(self):
         return self
 

litellm/proxy/common_request_processing.py

@@ -40,6 +40,9 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import (
     get_response_headers,
 )
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
+    AgenticAnthropicStreamingIterator,
+)
 from litellm.proxy._types import ProxyException, UserAPIKeyAuth
 from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
 from litellm.proxy.auth.auth_utils import check_response_size_is_safe
@@ -95,6 +98,10 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInforma
 }
 
 
+def _withheld_provider_output(response: object) -> bool:
+    return isinstance(response, AgenticAnthropicStreamingIterator) and response.has_buffered_provider_output
+
+
 def _should_return_raw_model_name(request_data: dict[str, object]) -> bool:
     return any(
         isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True
@@ -2970,7 +2977,7 @@ class ProxyBaseLLMRequestProcessing:
             # only sees GeneratorExit on GC) cannot own the refund.
             if not stream_completed:
                 client_disconnected = True
-            if not delivered_chunk:
+            if not delivered_chunk and not _withheld_provider_output(response):
                 from litellm.proxy.spend_tracking.budget_reservation import (
                     release_budget_reservation_on_cancel,
                 )

tests/test_litellm/proxy/test_budget_reservation.py

@@ -8,6 +8,9 @@ from fastapi import HTTPException
 import litellm
 from litellm.caching.dual_cache import DualCache
 from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES
+from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
+    AgenticAnthropicStreamingIterator,
+)
 from litellm.proxy._types import (
     LiteLLM_BudgetTable,
     LiteLLM_EndUserTable,
@@ -2376,6 +2379,11 @@ async def _reserve_for_stream(counter_cache, key_cache, proxy_logging_obj, token
     return valid_token, reservation
 
 
+async def _never_ending_stream():
+    yield b'event: message_start\ndata: {"type": "message_start"}\n\n'
+    await asyncio.sleep(30)
+
+
 def _drive_streaming_cancel(valid_token, iterator_hook):
     streaming_logging_obj = MagicMock()
     streaming_logging_obj.async_post_call_streaming_iterator_hook = iterator_hook
@@ -2481,6 +2489,61 @@ async def test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_c
     assert reservation["finalized"] is True
 
 
+@pytest.mark.asyncio
+async def test_streaming_cancel_while_holding_back_provider_output_keeps_reservation(
+    spend_counter_state,
+):
+    counter_cache, key_cache = spend_counter_state
+    proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
+    valid_token, reservation = await _reserve_for_stream(
+        counter_cache, key_cache, proxy_logging_obj, "key-cancel-held-back"
+    )
+
+    held_back = AgenticAnthropicStreamingIterator(
+        completion_stream=_never_ending_stream(),
+        http_handler=MagicMock(),
+        model="claude-haiku-4-5",
+        messages=[],
+        anthropic_messages_provider_config=MagicMock(),
+        anthropic_messages_optional_request_params={},
+        logging_obj=MagicMock(),
+        custom_llm_provider="anthropic",
+        kwargs={},
+        hold_back=True,
+        server_fulfilled_tool_names=frozenset({"headroom_retrieve"}),
+        ping_interval_seconds=0.01,
+    )
+
+    async def ping_then_cancel(user_api_key_dict, response, request_data):
+        yield await response.__anext__()
+        while not response.has_buffered_provider_output:
+            yield await response.__anext__()
+        raise asyncio.CancelledError()
+
+    streaming_logging_obj = MagicMock()
+    streaming_logging_obj.async_post_call_streaming_iterator_hook = ping_then_cancel
+    streaming_logging_obj._arelease_max_parallel_requests_on_disconnect = AsyncMock()
+    generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
+        response=held_back,
+        user_api_key_dict=valid_token,
+        request_data=_request_body(),
+        proxy_logging_obj=streaming_logging_obj,
+        serialize_chunk=lambda chunk: chunk,
+        serialize_error=lambda exc: str(exc),
+    )
+
+    received = []
+    with pytest.raises(asyncio.CancelledError):
+        async for chunk in generator:
+            received.append(chunk)
+
+    assert received and received == [STREAM_SSE_KEEPALIVE_PING_BYTES] * len(received)
+    assert counter_cache.in_memory_cache.get_cache(
+        key="spend:key:key-cancel-held-back"
+    ) == pytest.approx(2.0)
+    assert reservation.get("finalized") is not True
+
+
 @pytest.mark.asyncio
 async def test_release_budget_reservation_on_cancel_swallows_release_errors():
     # If the release itself fails (e.g. Redis unavailable) it must not escape