From ef0bcfb4a0b9189eda0f8f43619949c4d50accd2 Mon Sep 17 00:00:00 2001 From: Zachary Lyon Date: Tue, 15 Sep 2026 15:22:51 -0700 Subject: [PATCH] fix(proxy): defer billing for disconnected tinyfish SSE runs to the background poller --- .../tinyfish_passthrough_logging_handler.py | 36 ++++++++++++++++++- .../streaming_handler.py | 3 ++ ...st_tinyfish_passthrough_logging_handler.py | 33 +++++++++++++++++ 3 files changed, 71 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py index 70fcf921616..610f98f4f91 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py @@ -156,6 +156,25 @@ class TinyFishPassthroughLoggingHandler: verbose_proxy_logger.warning( "TinyFish passthrough: run-async response carried no run_id; logging the request without cost" ) + TinyFishPassthroughLoggingHandler._spawn_billing_poller( + run_id=run_id, + logging_obj=logging_obj, + result=result, + start_time=start_time, + cache_hit=cache_hit, + kwargs=kwargs, + ) + + @staticmethod + def _spawn_billing_poller( + run_id: str | None, + logging_obj: LiteLLMLoggingObj, + result: str, + start_time: datetime, + cache_hit: bool, + kwargs: Mapping[str, object], + client: AsyncHTTPHandler | None = None, + ) -> None: task: Final = asyncio.create_task( TinyFishPassthroughLoggingHandler._poll_and_log( run_id=run_id, @@ -164,6 +183,7 @@ class TinyFishPassthroughLoggingHandler: start_time=start_time, cache_hit=cache_hit, kwargs=kwargs, + client=client, ) ) _BACKGROUND_BILLING_TASKS.add(task) @@ -286,7 +306,9 @@ class TinyFishPassthroughLoggingHandler: client: AsyncHTTPHandler | None = None, ) -> PassThroughEndpointLoggingTypedDict: """Bill a POST /v1/automation/run-sse stream: SSE events carry no num_of_steps, so the - run_id parsed from the buffered events prices the run via one GET /v1/runs/{id}.""" + run_id parsed from the buffered events prices the run via one GET /v1/runs/{id}. A run + still live at stream end (client disconnect) is handed to the background poller instead, + signalled by a None result, so its steps are still billed when it finishes.""" try: run_id: Final = _run_id_from_sse_chunks(all_chunks) if run_id is None: @@ -294,6 +316,18 @@ class TinyFishPassthroughLoggingHandler: "TinyFish passthrough: no run_id in SSE stream; logging the request without cost" ) run: Final = await TinyFishPassthroughLoggingHandler._fetch_run(run_id, client) if run_id else None + if run_id is not None and (run is None or (run.get("status") or "") not in TINYFISH_TERMINAL_RUN_STATUSES): + TinyFishPassthroughLoggingHandler._spawn_billing_poller( + run_id=run_id, + logging_obj=litellm_logging_obj, + result="", + start_time=start_time, + cache_hit=False, + kwargs=_EMPTY_KWARGS, + client=client, + ) + deferred_payload: Final[PassThroughEndpointLoggingTypedDict] = {"result": None, "kwargs": {}} + return deferred_payload payload: Final = TinyFishPassthroughLoggingHandler._build_logging_payload( run=run, logging_obj=litellm_logging_obj, diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 69a8a7afc20..a922a75f7cb 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -286,6 +286,9 @@ class PassThroughStreamingHandler: end_time=end_time, ) ) + if tinyfish_payload["result"] is None: + # run still live upstream (client disconnect); the background poller owns the spend log + return await litellm_logging_obj.dispatch_success_handlers( result=tinyfish_payload["result"], start_time=start_time, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py index 5f5dde06c53..053911ed3c8 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py @@ -248,6 +248,39 @@ class TestSseBilling: assert payload["kwargs"]["model"] == "tinyfish/automation-run" assert fake_client.requested_urls == ["https://agent.tinyfish.ai/v1/runs/run-7?screenshots=none"] + def test_disconnected_stream_defers_billing_to_poller(self, tinyfish_env): + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + _BACKGROUND_BILLING_TASKS, + ) + + logging_obj = _make_logging_obj() + logging_obj.dispatch_success_handlers = AsyncMock() + chunks = ['data: {"type": "STARTED", "run_id": "run-11", "status": "RUNNING"}'] + fake_client = _FakeClient( + payloads=[ + {"run_id": "run-11", "status": "RUNNING"}, + {"run_id": "run-11", "status": "COMPLETED", "num_of_steps": 4}, + ] + ) + + async def _run(): + payload = await TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks( + litellm_logging_obj=logging_obj, + url_route="https://agent.tinyfish.ai/v1/automation/run-sse", + start_time=datetime.now(), + all_chunks=chunks, + end_time=datetime.now(), + client=fake_client, + ) + assert payload["result"] is None + await asyncio.gather(*list(_BACKGROUND_BILLING_TASKS)) + + asyncio.run(_run()) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + assert logging_obj.dispatch_success_handlers.await_args.kwargs["response_cost"] == pytest.approx(0.064) + assert len(fake_client.requested_urls) == 2 + def test_stream_without_run_id_logs_without_cost(self, tinyfish_env): fake_client = _FakeClient(payloads=[{}])