fix(anthropic_messages): log spend for interrupted /v1/messages streams

[devin-ai-integration]

TLDR

Problem this solves:

• interrupted streaming /v1/messages on bedrock invoke logged no spend
• disconnect billing never ran, so partial usage was free

How it solves it:

• log collected chunks from a finally, not only after full consumption
• drop the extra async-generator layer that deferred teardown to GC

Relevant issues

Fixes #35958

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Screenshots / Proof of Fix

Proof was captured against a local proxy on localhost:4000 with a stub bedrock invoke-with-response-stream endpoint (aws_bedrock_runtime_endpoint: http://127.0.0.1:8112) that speaks real AWS event-stream framing and emits real anthropic message events plus amazon-bedrock-invocationMetrics, because this sandbox has no outbound AWS access. Everything else in the path is the real proxy: real auth, real anthropic messages handler, real decoder, real cost calculation, real LiteLLM_SpendLogs writes. Someone with bedrock credentials should re-run the same two commands against a real bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0 deployment before merge

The repro streams /v1/messages, drops the connection after 5 SSE lines, then polls /spend/logs:

curl -sN http://127.0.0.1:4000/v1/messages \
  -H "Authorization: Bearer $KEY" -H 'content-type: application/json' \
  -d '{"model":"stub-bedrock","max_tokens":1024,"stream":true,
       "messages":[{"role":"user","content":"count to 30"}]}' | head -5
sleep 10
curl -s "http://127.0.0.1:4000/spend/logs?api_key=$KEY" -H "Authorization: Bearer sk-1234"

Before, at 24dbd2b2db:

disconnected after 5 chunks
  poll 1: 0 row(s)
  poll 2: 0 row(s)
  poll 3: 0 row(s)
  poll 4: 0 row(s)
  poll 5: 0 row(s)
  poll 6: 0 row(s)
RESULT: FAIL - no spend log row

After, at fce256f584, three consecutive interrupted runs:

RESULT: PASS - {"request_id": "70f4970a-5e6a-45aa-88fb-06aae9f3e173", "call_type": "anthropic_messages", "spend": 0.000198, "prompt_tokens": 45, "completion_tokens": 3, "total_tokens": 48, "model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "status": "success"}
RESULT: PASS - {"request_id": "a838b080-a5e3-4b4b-b429-1bd99f213bf5", "call_type": "anthropic_messages", "spend": 0.000198, "prompt_tokens": 45, "completion_tokens": 3, "total_tokens": 48, "model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "status": "success"}
RESULT: PASS - {"request_id": "77f7fe78-41af-4f8c-a5fe-66c4d0fea084", "call_type": "anthropic_messages", "spend": 0.000198, "prompt_tokens": 45, "completion_tokens": 3, "total_tokens": 48, "model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "status": "success"}

A fully consumed stream at fce256f584 still writes exactly one row with the provider's own metrics, so nothing is double billed:

RESULT: PASS - {"request_id": "bcaf46e4-3bc0-4ec2-98fe-5ff75806b79c", "call_type": "anthropic_messages", "spend": 0.00858, "prompt_tokens": 45, "completion_tokens": 511, "total_tokens": 556, "model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "status": "success"}

$ psql $DATABASE_URL -c 'select request_id, count(*) from "LiteLLM_SpendLogs" group by request_id order by min("startTime") desc limit 4'
              request_id              | count
--------------------------------------+-------
 bcaf46e4-3bc0-4ec2-98fe-5ff75806b79c |     1
 77f7fe78-41af-4f8c-a5fe-66c4d0fea084 |     1
 a838b080-a5e3-4b4b-b429-1bd99f213bf5 |     1
 70f4970a-5e6a-45aa-88fb-06aae9f3e173 |     1

Type

🐛 Bug Fix

Changes

BaseAnthropicMessagesStreamingIterator.async_sse_wrapper collected every SSE chunk it forwarded and then, once the upstream stream ended, dispatched spend logging for them. A client that hangs up mid-stream gets that generator closed instead, which raises GeneratorExit at the yield, so the dispatch line was simply never reached and the request was billed nothing. The proxy-level disconnect billing added in #33736 cannot cover this path either, since _bill_partial_streamed_spend_on_disconnect reads response.chunks off a CustomStreamWrapper and the anthropic messages response is a bare async generator. Moving the dispatch into a finally makes the chunks already streamed get logged on interruption, matching what PassThroughStreamingHandler.chunk_processor already does for the native anthropic passthrough (which is why the same interruption on anthropic/... does log):

try:
    async for chunk in completion_stream:
        collected_chunks.append(encoded_chunk)
        yield encoded_chunk
    ...
finally:
    if collected_chunks:
        await self._handle_streaming_logging(collected_chunks)

That alone was not enough on bedrock invoke. bedrock_sse_wrapper was itself an async generator re-yielding the wrapper's chunks, and closing the outer generator does not deterministically close the inner one, it only drops the reference and leaves the finally to whenever the event loop finalizes the abandoned asyncgen. In practice the spend log landed sometimes and was lost most of the time. It now returns the wrapper's iterator directly, so the aclose() the proxy already performs in _finalize_streaming_generator_cleanup reaches the finally right away:

-    async def bedrock_sse_wrapper(...):
-        async for chunk in handler.async_sse_wrapper(patched_stream):
-            yield chunk
+    def bedrock_sse_wrapper(...) -> AsyncIterator[bytes]:
+        return handler.async_sse_wrapper(patched_stream)

Logging dispatch also moves from a bare asyncio.create_task to GLOBAL_LOGGING_WORKER, the same queue chunk_processor uses, since a task created while the request is being torn down has no owner keeping it alive

Usage on an interrupted stream is whatever the provider reported in the chunks the client did receive, so message_start input tokens plus the output tokens seen so far, not the tokens the model would have gone on to generate. Bedrock only sends amazon-bedrock-invocationMetrics on the final frame, and the proxy stops reading upstream on disconnect by design (_UpstreamClosingStreamingResponse), so billing the full generation would mean draining the abandoned stream to the end. This PR deliberately does not do that; it is the same partial-usage semantics #33736 chose for /chat/completions

QA runbook

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[greptile-apps]

Greptile Summary

This PR moves Anthropic messages streaming spend dispatch into generator teardown and removes Bedrock's extra async-generator layer so interrupted streams are billed promptly.

• Enqueues collected streaming chunks through the global logging worker.
• Adds interruption tests for the shared Anthropic wrapper and Bedrock invoke path.
• Preserves direct async-iterator cleanup through the Bedrock transformation.

Confidence Score: 4/5

The upstream-exception path needs to be fixed before merging because it dispatches both success and failure logging for the same failed stream.

The teardown change correctly covers client disconnects, but its unconditional success dispatch also runs for ordinary upstream exceptions after partial output and conflicts with the proxy's subsequent failure handling.

Files Needing Attention: litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py

Important Files Changed

Filename	Overview
litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py	Moves logging into finally and onto the global worker, but regular upstream exceptions now trigger success logging before the proxy's failure logging.
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py	Returns the inner async iterator directly so explicit closure reaches its teardown deterministically.
tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py	Adds coverage for explicit closure after partial consumption and closure before the first chunk, but not an upstream exception after partial output.
tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py	Verifies that closing the Bedrock iterator reaches logging setup, while leaving the exception-exit behavior uncovered.

_{Reviews (1): Last reviewed commit: "fix(anthropic_messages): log spend for i..." | Re-trigger Greptile}

[greptile-apps]

Exceptions trigger success logging

When a Bedrock stream yields chunks and then raises an upstream exception, this finally dispatches the partial response through the success handlers before the exception reaches the proxy failure handler, causing the same failed request to produce contradictory success and failure logging and a spend record marked as successful.

Knowledge Base Used: Cost Tracking and Budget Enforcement

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing devin_ai_fix_interrupted_messages_stream_spend_logging (fce256f) with litellm_internal_staging (2792887)}

[nuernber]

This does not preserve with full spend logs, just tested it. I manually checked the recorded spend log in LiteLLM vs. the Bedrock invocation log. The input tokens matched but the Bedrock invocation logs had more output tokens. The solution in PR #36008 works with full spend log preservation.

[nuernber]

Review notes (AI generated)

Main concern: this fix makes an interrupted stream produce a spend row again, but it bills only the chunks the client drained before disconnecting, not the full response Bedrock actually generated and billed. So it fixes the "zero rows" half of the regression while leaving the other half, the output tokens reading lower than the AWS Bedrock invocation logs, unfixed. That output undercount is the specific symptom being reported.

Why it happens: on a client disconnect, GeneratorExit is thrown into async_sse_wrapper at the yield, the finally runs, and it logs collected_chunks, which only holds what was streamed so far. The upstream completion_stream is never drained past that point, so the terminal message_delta / message_stop carrying the real output_tokens is never read. Cost then falls back to re-tokenizing the truncated partial text, which is always lower than what Bedrock billed. Bedrock keeps generating server-side regardless of the client, so the only way to bill the true output is to keep reading the upstream to its terminal usage event even after the client is gone.

Concrete repro on an early disconnect (client reads a 3 event prefix, real stream ends with output_tokens=64):

this branch:   billed output_tokens = None  (falls back to re-tokenized partial, ~1-15), no message_stop in billed chunks
detached pump: billed output_tokens = 64,    message_stop present in billed chunks

The added test test_async_sse_wrapper_logs_collected_chunks_when_client_disconnects_mid_stream asserts iterator.logged_chunks == streamed, i.e. it asserts that only the 3 chunks the client read get billed. That locks in the truncated-billing behavior as expected, so the suite passes even though the output count is wrong. A test that gates the stream tail behind an event released only after the client disconnects, then asserts the billed output_tokens equals the full stream's terminal value, would surface the gap.

Secondary robustness note: _handle_streaming_logging is awaited inside the finally while GeneratorExit is unwinding the generator during aclose(). It happens to only enqueue synchronously today, so it does not suspend, but awaiting in a finally on the disconnect teardown path is fragile: any future await that yields to the loop there risks "async generator ignored GeneratorExit" or a hung aclose. Doing the drain and billing in a detached task keeps that work off the teardown path entirely.

Alternative approach for comparison: run the upstream read in a detached background task rooted so it is not garbage collected. The pump drains the provider stream to its terminal usage event and bills there, and the client facing generator only relays chunks off a queue, so a disconnect tears down the relay but not the drain. A client_detached event stops enqueueing after disconnect so the queue cannot grow unbounded while the pump finishes. That bills the real output tokens on interrupted streams, which is the part this branch still misses.