Merge ddb7409151 into ec3f8183c3

litellm/llms/anthropic/chat/handler.py

@@ -118,6 +118,7 @@ async def make_call(
         speed=speed,
         tool_name_reverse_map=tool_name_reverse_map,
     )
+    completion_stream.http_response = response
 
     # LOGGING
     logging_obj.post_call(
@@ -186,6 +187,7 @@ def make_sync_call(
         speed=speed,
         tool_name_reverse_map=tool_name_reverse_map,
     )
+    completion_stream.http_response = response
 
     # LOGGING
     logging_obj.post_call(
@@ -638,6 +640,8 @@ class ModelResponseIterator:
     ):
         self.streaming_response = streaming_response
         self.response_iterator = self.streaming_response
+        self.sync_stream = sync_stream
+        self.http_response: httpx.Response | None = None
         self.content_blocks: list[ContentBlockDelta] = []
         self.tool_index = -1
         self.json_mode = json_mode
@@ -691,6 +695,16 @@ class ModelResponseIterator:
     def accumulated_json(self, value: str) -> None:
         self._json_buffer.set(value)
 
+    async def aclose(self) -> None:
+        response: Final = self.http_response
+        self.http_response = None
+        if response is None:
+            return
+        if self.sync_stream:
+            response.close()
+            return
+        await response.aclose()
+
     def check_empty_tool_call_args(self) -> bool:
         """
         Check if the tool call block so far has been an empty string

tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py

@@ -1,13 +1,22 @@
 import json
 import threading
+from collections.abc import AsyncIterator, Iterator
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import anyio
+import httpx
 import pytest
 
 import litellm
 from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
-from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call
+from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+from litellm.llms.anthropic.chat.handler import (
+    ModelResponseIterator,
+    make_call,
+    make_sync_call,
+)
 from litellm.types.llms.openai import (
     ChatCompletionToolCallChunk,
     ChatCompletionToolCallFunctionChunk,
@@ -15,6 +24,29 @@ from litellm.types.llms.openai import (
 from litellm.types.responses.main import OutputCodeInterpreterCall
 
 
+class _RecordingAsyncByteStream(httpx.AsyncByteStream):
+    def __init__(self) -> None:
+        self.aclose_calls: int = 0
+
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        yield b'data: {"type":"message_start"}\n\n'
+
+    async def aclose(self) -> None:
+        await anyio.sleep(0)
+        self.aclose_calls += 1
+
+
+class _RecordingSyncByteStream(httpx.SyncByteStream):
+    def __init__(self) -> None:
+        self.close_calls: int = 0
+
+    def __iter__(self) -> Iterator[bytes]:
+        yield b'data: {"type":"message_start"}\n\n'
+
+    def close(self) -> None:
+        self.close_calls += 1
+
+
 @pytest.mark.asyncio
 async def test_make_call_passes_logging_obj_to_client_post():
     """make_call must pass logging_obj to client.post so track_llm_api_timing can set llm_api_duration_ms for litellm_overhead_time_ms."""
@@ -46,6 +78,78 @@ async def test_make_call_passes_logging_obj_to_client_post():
     assert call_kwargs.get("logging_obj") is logging_obj
 
 
+@pytest.mark.asyncio
+async def test_make_call_stream_cleanup_closes_http_response_once_under_cancellation():
+    stream: Final = _RecordingAsyncByteStream()
+    response: Final = httpx.Response(200, stream=stream)
+    mock_client: Final = AsyncMock()
+    mock_client.post.return_value = response
+    logging_obj: Final = MagicMock(model_call_details={"litellm_params": {}})
+
+    completion_stream, _ = await make_call(
+        client=mock_client,
+        api_base="https://api.anthropic.com/v1/messages",
+        headers={},
+        data="{}",
+        model="claude-haiku-4-5",
+        messages=[{"role": "user", "content": "Hi"}],
+        logging_obj=logging_obj,
+        timeout=60.0,
+        json_mode=False,
+    )
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=completion_stream,
+        model="claude-haiku-4-5",
+        custom_llm_provider="anthropic",
+        logging_obj=logging_obj,
+    )
+
+    with anyio.CancelScope() as cancel_scope:
+        cancel_scope.cancel()
+        await wrapper.aclose()
+    await wrapper.aclose()
+    await completion_stream.aclose()
+
+    assert response.is_closed is True
+    assert completion_stream.http_response is None
+    assert stream.aclose_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_make_sync_call_stream_cleanup_closes_http_response_once():
+    stream: Final = _RecordingSyncByteStream()
+    response: Final = httpx.Response(200, stream=stream)
+    mock_client: Final = MagicMock()
+    mock_client.post.return_value = response
+    logging_obj: Final = MagicMock(model_call_details={"litellm_params": {}})
+
+    completion_stream, _ = make_sync_call(
+        client=mock_client,
+        api_base="https://api.anthropic.com/v1/messages",
+        headers={},
+        data="{}",
+        model="claude-haiku-4-5",
+        messages=[{"role": "user", "content": "Hi"}],
+        logging_obj=logging_obj,
+        timeout=60.0,
+        json_mode=False,
+    )
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=completion_stream,
+        model="claude-haiku-4-5",
+        custom_llm_provider="anthropic",
+        logging_obj=logging_obj,
+    )
+
+    await wrapper.aclose()
+    await wrapper.aclose()
+    await completion_stream.aclose()
+
+    assert response.is_closed is True
+    assert completion_stream.http_response is None
+    assert stream.close_calls == 1
+
+
 def test_redacted_thinking_content_block_delta():
     chunk = {
         "type": "content_block_start",