fix(exa): account for interrupted search streams

cookbook/exa_search/README.md

@@ -86,7 +86,7 @@ async with stream:
 print(stream.final_response.model_dump())
 ```
 
-The synchronous equivalent uses `with litellm.search(..., stream=True) as stream` and `for event in stream`. Use the context manager or call `close()` / `await aclose()` when abandoning a stream early.
+The synchronous equivalent uses `with litellm.search(..., stream=True) as stream` and `for event in stream`. Use the context manager or call `close()` / `await aclose()` when abandoning a stream early. Early close and cancellation record LiteLLM's per-query charge once, using the results received so far; they do not populate `final_response` without a `done` event.
 
 Events preserve their native `type` and payload. The documented types are `text-delta`, `grounding`, `results`, `stream-reset`, `done`, and `error`. On `stream-reset`, discard previously accumulated synthesized text. `done.output` is authoritative. `final_response` becomes available after `done`, combining the most recent results with the terminal output and metadata. Provider error events, malformed events, and premature disconnects raise an exception and close the response. Heartbeats are ignored. SSE framing is decoded incrementally, including fragmented UTF-8 and multiline data.
 

litellm/llms/base_llm/search/streaming_iterator.py

@@ -90,6 +90,8 @@ class SearchStreamWrapper:
         self._results: object = []
         self._request_metadata: Mapping[str, object] = {}
         self._finished = False
+        self._logging_started = False
+        self._success_task: asyncio.Task[None] | None = None
         self._iterator = self._iter_events()
         self._async_iterator = self._aiter_events()
         self.final_response: SearchResponse | None = None
@@ -162,7 +164,11 @@ class SearchStreamWrapper:
                 if self._finished:
                     return
             raise self._error("Search stream disconnected before its done event")
+        except GeneratorExit:
+            self._log_success()
+            raise
         except Exception as error:
+            self._logging_started = True
             failure: Final = self._stream_error(error)
             self.logging_obj.failure_handler(
                 failure, str(failure), self.logging_obj.start_time, datetime.now(timezone.utc).replace(tzinfo=None)
@@ -180,18 +186,17 @@ class SearchStreamWrapper:
                     continue
                 event = self._parse(sse.data)
                 if self._finished:
-                    await self.logging_obj.dispatch_success_handlers(
-                        self.final_response,
-                        start_time=self.logging_obj.start_time,
-                        end_time=datetime.now(timezone.utc).replace(tzinfo=None),
-                        prefer_async_handlers=True,
-                    )
+                    await self._alog_success()
                     await self.response.aclose()
                 yield event
                 if self._finished:
                     return
             raise self._error("Search stream disconnected before its done event")
+        except (GeneratorExit, asyncio.CancelledError):
+            await self._alog_success()
+            raise
         except Exception as error:
+            self._logging_started = True
             failure: Final = self._stream_error(error)
             await self.logging_obj.async_failure_handler(
                 failure, str(failure), self.logging_obj.start_time, datetime.now(timezone.utc).replace(tzinfo=None)
@@ -202,12 +207,38 @@ class SearchStreamWrapper:
         finally:
             await self.response.aclose()
 
+    def _billing_response(self) -> SearchResponse:
+        if self.final_response is not None:
+            return self.final_response
+        return self.provider_config.transform_search_response(
+            httpx.Response(200, json={**self._request_metadata, "results": self._results}),
+            self._transform_logging_obj,
+        )
+
+    async def _alog_success(self) -> None:
+        if self._success_task is None:
+            if self._logging_started:
+                return
+            self._logging_started = True
+            self._success_task = asyncio.create_task(
+                self.logging_obj.dispatch_success_handlers(
+                    self._billing_response(),
+                    start_time=self.logging_obj.start_time,
+                    end_time=datetime.now(timezone.utc).replace(tzinfo=None),
+                    prefer_async_handlers=True,
+                )
+            )
+        await asyncio.shield(self._success_task)
+
     def _log_success(self) -> None:
+        if self._logging_started:
+            return
+        self._logging_started = True
         end_time: Final = datetime.now(timezone.utc).replace(tzinfo=None)
         executor.submit(
             asyncio.run,
             self.logging_obj.dispatch_success_handlers(
-                self.final_response, start_time=self.logging_obj.start_time, end_time=end_time
+                self._billing_response(), start_time=self.logging_obj.start_time, end_time=end_time
             ),
         )
 
@@ -228,10 +259,12 @@ class SearchStreamWrapper:
     def close(self) -> None:
         self._iterator.close()
         self.response.close()
+        self._log_success()
 
     async def aclose(self) -> None:
         await self._async_iterator.aclose()
         await self.response.aclose()
+        await self._alog_success()
 
     def __enter__(self) -> SearchStreamWrapper:
         return self

litellm/proxy/common_request_processing.py

@@ -446,6 +446,11 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
     logging_obj: Final = request_data.get("litellm_logging_obj")
     if not isinstance(logging_obj, LiteLLMLoggingObj):
         return False
+    from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
+
+    if isinstance(response, SearchStreamWrapper):
+        await response.aclose()
+        return bool(logging_obj.model_call_details.get("has_dispatched_final_stream_success"))
     if logging_obj.model_call_details.get("has_dispatched_final_stream_success"):
         # A natural end-of-stream success event already fired and released the
         # slot; do not bill again, and let the caller skip the slot release.

litellm/utils.py

@@ -1735,10 +1735,6 @@ def client(original_function):
                 result=result,
             ):
                 if kwargs.get("complete_response") is True and not isinstance(result, Coroutine):
-                    from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
-
-                    if isinstance(result, SearchStreamWrapper):
-                        return result.get_final_response()
                     chunks = []
                     for idx, chunk in enumerate(result):
                         chunks.append(chunk)

tests/integration/sdk/test_exa_search_wire.py

@@ -145,3 +145,91 @@ async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, s
         assert isinstance(details["standard_logging_object"], Mapping)
         assert isinstance(details["response_cost"], (int, float))
         assert details["stream"] is sse
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("asynchronous", "consume"),
+    [
+        (False, "none"),
+        (True, "none"),
+        (False, "results"),
+        (True, "results"),
+        (False, "done"),
+        (True, "done"),
+        (True, "cancel"),
+    ],
+)
+async def test_search_disconnect_bills_once(asynchronous: bool, consume: str) -> None:
+    import asyncio
+
+    from litellm.integrations.custom_logger import CustomLogger
+
+    class Recorder(CustomLogger):
+        def __init__(self) -> None:
+            self.finished = threading.Event()
+            self.records = ()
+
+        def log_success_event(self, kwargs, response_obj, start_time, end_time):
+            self.records = (*self.records, (kwargs, response_obj))
+            self.finished.set()
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            self.log_success_event(kwargs, response_obj, start_time, end_time)
+
+    recorder: Final = Recorder()
+    release: Final = threading.Event()
+    events: Final = (
+        {"type": "results", "results": PAYLOAD["results"], "requestId": "disconnect"},
+        {"type": "done", "output": PAYLOAD["output"], "searchTime": 14},
+    )
+
+    def respond(request: Request) -> Reply:
+        assert json.loads(request.body)["stream"] is True
+        return Reply(
+            content_type="text/event-stream",
+            chunks=tuple(("data: " + json.dumps(event) + "\n\n").encode() for event in events),
+            gate_after_first=release if consume == "cancel" else None,
+        )
+
+    with wire_server(respond) as server:
+        params: Final = {
+            "query": "question",
+            "search_provider": "exa_ai",
+            "api_base": server.url,
+            "api_key": "test-key",
+            "stream": True,
+            "success_callback": [recorder.async_log_success_event if asynchronous else recorder],
+            **PARAMS,
+        }
+        stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
+        try:
+            if consume != "none":
+                first: Final = await anext(stream) if asynchronous else next(stream)
+                assert first.type == "results"
+            if consume == "done":
+                if asynchronous:
+                    await stream.aget_final_response()
+                else:
+                    stream.get_final_response()
+            if consume == "cancel":
+                pending: Final = asyncio.create_task(anext(stream))
+                await asyncio.sleep(0)
+                pending.cancel()
+                with pytest.raises(asyncio.CancelledError):
+                    await pending
+            if asynchronous:
+                await stream.aclose()
+                await stream.aclose()
+            else:
+                stream.close()
+                stream.close()
+            assert await asyncio.to_thread(recorder.finished.wait, 5)
+            assert len(recorder.records) == 1
+            details, logged = recorder.records[0]
+            assert details["response_cost"] > 0
+            assert details["standard_logging_object"]["response_cost"] == details["response_cost"]
+            assert len(logged.results) == (0 if consume == "none" else 1)
+            assert (stream.final_response is not None) is (consume == "done")
+        finally:
+            release.set()

tests/unit/llms/base_llm/search/test_streaming_iterator.py

@@ -56,7 +56,8 @@ class FragmentedStream(httpx.SyncByteStream, httpx.AsyncByteStream):
 
 def wire(*events: object) -> bytes:
     return (
-        ": heartbeat\r\n\r\n" + "".join("event: message\r\ndata: " + json.dumps(event) + "\r\n\r\n" for event in events)
+        ": heartbeat\r\n\r\nevent: ping\r\ndata:\r\n\r\n"
+        + "".join("event: message\r\ndata: " + json.dumps(event) + "\r\n\r\n" for event in events)
     ).encode()
 
 
@@ -91,6 +92,36 @@ async def test_search_stream_preserves_events_and_final_response(respx_mock: res
     assert stream.final_response.model_dump()["choices"] == EVENTS[-1]["choices"]
 
 
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.asyncio
+async def test_abandoned_stream_cannot_return_incomplete_final_response(
+    respx_mock: respx.MockRouter, asynchronous: bool
+):
+    body: Final = FragmentedStream(wire(*EVENTS))
+    respx_mock.post(URL).mock(
+        return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=body)
+    )
+    params: Final = {
+        "query": "question",
+        "search_provider": "exa_ai",
+        "api_key": "test",
+        "stream": True,
+        "outputSchema": SCHEMA,
+    }
+    stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
+    if asynchronous:
+        await anext(stream)
+        await stream.aclose()
+        with pytest.raises(litellm.APIError, match="no final response"):
+            await stream.aget_final_response()
+    else:
+        next(stream)
+        stream.close()
+        with pytest.raises(litellm.APIError, match="no final response"):
+            stream.get_final_response()
+    assert body.closed
+
+
 @pytest.mark.parametrize("asynchronous", [False, True])
 @pytest.mark.parametrize(
     "body",
@@ -123,6 +154,10 @@ async def test_stream_failure_is_not_silent_and_closes_connection(
     else:
         with pytest.raises(litellm.APIError):
             list(stream)
+    if asynchronous:
+        await stream.aclose()
+    else:
+        stream.close()
     assert transport.closed
 
 

tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py

@@ -3,9 +3,12 @@ import json
 from collections.abc import Callable
 from datetime import datetime
 from types import SimpleNamespace
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import httpx
 import pytest
+import respx
 from fastapi import HTTPException
 from fastapi.testclient import TestClient
 
@@ -25,6 +28,30 @@ from litellm.types.search import SearchToolInfoResponse
 client = TestClient(app)
 
 
+@pytest.mark.parametrize("has_results", [False, True])
+@pytest.mark.asyncio
+async def test_exa_connection_reports_returned_result_count(
+    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, has_results: bool
+) -> None:
+    import litellm
+    from litellm.proxy.search_endpoints.search_tool_management import (
+        TestSearchToolConnectionRequest,
+        test_search_tool_connection,
+    )
+
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    results: Final = [{"title": "Source", "url": "https://example.com"}] if has_results else []
+    route: Final = respx_mock.post("https://api.exa.ai/search").mock(
+        return_value=httpx.Response(200, json={"results": results})
+    )
+    response: Final = await test_search_tool_connection(
+        TestSearchToolConnectionRequest(litellm_params={"search_provider": "exa_ai", "api_key": "test"})
+    )
+    assert response["status"] == "success"
+    assert response["results_count"] == len(results)
+    assert json.loads(route.calls.last.request.content)["numResults"] == 1
+
+
 @pytest.mark.asyncio
 async def test_list_search_tools_db_only(monkeypatch):
     """Test listing search tools when only DB tools exist"""

tests/unit/proxy/test_common_request_processing.py

@@ -10979,3 +10979,60 @@ class TestResponseCostFromLoggingObjForBridgedMessages:
         assert logging_obj.cost_breakdown is not None
         assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(expected_output)
         assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(expected_reasoning)
+
+
+@pytest.mark.parametrize("cancel_logging", [False, True])
+@pytest.mark.asyncio
+async def test_search_disconnect_billing_owns_slot_release(respx_mock, monkeypatch, cancel_logging):
+    import asyncio
+
+    import httpx
+
+    from litellm.integrations.custom_logger import CustomLogger
+
+    started: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+    if not cancel_logging:
+        release.set()
+
+    class Recorder(CustomLogger):
+        def __init__(self):
+            self.costs = ()
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            started.set()
+            await release.wait()
+            self.costs = (*self.costs, kwargs["response_cost"])
+
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    recorder: Final = Recorder()
+    respx_mock.post("https://api.exa.ai/search").mock(
+        return_value=httpx.Response(
+            200,
+            headers={"content-type": "text/event-stream"},
+            content=b'data: {"type":"results","results":[{"url":"https://example.com","title":"Source"}]}\n\n',
+        )
+    )
+    stream: Final = await litellm.asearch(
+        query="question",
+        search_provider="exa_ai",
+        api_key="test",
+        stream=True,
+        outputSchema={"type": "text"},
+        success_callback=[recorder.async_log_success_event],
+    )
+    assert (await anext(stream)).type == "results"
+    if cancel_logging:
+        closing: Final = asyncio.create_task(stream.aclose())
+        await started.wait()
+        closing.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await closing
+        release.set()
+    request_data: Final = {"litellm_logging_obj": stream.logging_obj}
+    assert await _bill_partial_streamed_spend_on_disconnect(request_data, stream)
+    assert await _bill_partial_streamed_spend_on_disconnect(request_data, stream)
+    assert len(recorder.costs) == 1
+    assert recorder.costs[0] > 0
+    assert stream.response.is_closed
+    assert stream.final_response is None