From 1ef034bff6aaaa48e91c25484e7a0017e3a0d473 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Tue, 28 Apr 2026 21:07:17 +0000
Subject: [PATCH 1/4] fix(passthrough): flush spend tracking on interrupted
 Bedrock streams

When a client disconnects mid-stream from a Bedrock pass-through endpoint,
Starlette calls aclose() on the async generator, raising GeneratorExit
(a BaseException, not Exception) at the suspended yield. The previous
`except Exception` blocks in _async_streaming/_sync_streaming
(litellm/passthrough/main.py) and PassThroughStreamingHandler.chunk_processor
did not catch GeneratorExit, so the post-loop flush that hands collected
raw bytes to async_flush_passthrough_collected_chunks /
_route_streaming_logging_to_handler never ran. All per-chunk usage data
was silently dropped, undercounting spend for interrupted Bedrock invoke
and converse streams.

Move the flush into a finally block in all three sites and guard with a
`flush_scheduled` flag so the success path still flushes exactly once.
Also pull raise_for_status() out of the chunk-collection try block in
_async_streaming so 4xx/5xx responses still raise and don't enter the
flush path with zero bytes (preserving the behavior tested by
test_async_streaming_error_propagation.py).

Add regression coverage:
- test_async_streaming_flushes_on_client_disconnect
- test_async_streaming_flushes_on_upstream_exception_with_partial_data
- test_sync_streaming_flushes_on_early_close
- test_chunk_processor_logs_on_client_disconnect
plus baseline tests for normal completion and the 4xx no-flush path.

Fixes LIT-2642.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
---
 litellm/passthrough/main.py                   |  76 +++--
 .../streaming_handler.py                      |  63 ++--
 ...test_streaming_interrupt_spend_tracking.py | 297 ++++++++++++++++++
 .../test_streaming_handler_interrupt.py       | 144 +++++++++
 4 files changed, 534 insertions(+), 46 deletions(-)
 create mode 100644 tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py
 create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py

diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py
index edee50bdfc..2a80e2bb63 100644
--- a/litellm/passthrough/main.py
+++ b/litellm/passthrough/main.py
@@ -390,19 +390,29 @@ def _sync_streaming(
 ):
     from litellm.utils import executor
 
+    raw_bytes: List[bytes] = []
+    flush_scheduled = False
     try:
-        raw_bytes: List[bytes] = []
         for chunk in response.iter_bytes():  # type: ignore
             raw_bytes.append(chunk)
             yield chunk
-
-        executor.submit(
-            litellm_logging_obj.flush_passthrough_collected_chunks,
-            raw_bytes=raw_bytes,
-            provider_config=provider_config,
-        )
-    except Exception as e:
-        raise e
+    finally:
+        # Always flush collected chunks for spend tracking, even if the
+        # consumer terminates the generator early (GeneratorExit). Without
+        # this, an interrupted stream loses all per-chunk usage data
+        # because the post-loop flush never runs. See LIT-2642.
+        if not flush_scheduled and raw_bytes:
+            flush_scheduled = True
+            try:
+                executor.submit(
+                    litellm_logging_obj.flush_passthrough_collected_chunks,
+                    raw_bytes=raw_bytes,
+                    provider_config=provider_config,
+                )
+            except Exception:
+                # Don't mask the original exception (incl. GeneratorExit)
+                # if scheduling the flush itself fails.
+                pass
 
 
 async def _async_streaming(
@@ -411,23 +421,47 @@ async def _async_streaming(
     provider_config: "BasePassthroughConfig",
 ):
     iter_response = await response
+
+    # Validate response status before consuming the body so 4xx/5xx
+    # responses raise without entering the chunk-collection path.
     try:
         iter_response.raise_for_status()
-        raw_bytes: List[bytes] = []
-
-        async for chunk in iter_response.aiter_bytes():  # type: ignore
-            raw_bytes.append(chunk)
-            yield chunk
-
-        asyncio.create_task(
-            litellm_logging_obj.async_flush_passthrough_collected_chunks(
-                raw_bytes=raw_bytes,
-                provider_config=provider_config,
-            )
-        )
     except Exception:
         try:
             await iter_response.aclose()
         except Exception:
             pass
         raise
+
+    raw_bytes: List[bytes] = []
+    flush_scheduled = False
+    try:
+        async for chunk in iter_response.aiter_bytes():  # type: ignore
+            raw_bytes.append(chunk)
+            yield chunk
+    except Exception:
+        try:
+            await iter_response.aclose()
+        except Exception:
+            pass
+        raise
+    finally:
+        # Always flush collected chunks for spend tracking, even if the
+        # client disconnects mid-stream. On disconnect, Starlette calls
+        # aclose() on this generator, which raises GeneratorExit at the
+        # suspended `yield` — `except Exception` does not catch it, so
+        # the post-loop flush would otherwise be skipped and all
+        # captured per-chunk usage data lost. See LIT-2642.
+        if not flush_scheduled and raw_bytes:
+            flush_scheduled = True
+            try:
+                asyncio.create_task(
+                    litellm_logging_obj.async_flush_passthrough_collected_chunks(
+                        raw_bytes=raw_bytes,
+                        provider_config=provider_config,
+                    )
+                )
+            except Exception:
+                # Don't mask the original exception (incl. GeneratorExit)
+                # if scheduling the flush itself fails.
+                pass
diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py
index 302d7e76ed..fe5bd42a60 100644
--- a/litellm/proxy/pass_through_endpoints/streaming_handler.py
+++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py
@@ -41,16 +41,17 @@ class PassThroughStreamingHandler:
         - Collect non-empty chunks for post-processing (logging)
         - Inject cost into chunks if include_cost_in_streaming_usage is enabled
         """
-        try:
-            raw_bytes: List[bytes] = []
-            # Extract model name for cost injection
-            model_name = PassThroughStreamingHandler._extract_model_for_cost_injection(
-                request_body=request_body,
-                url_route=url_route,
-                endpoint_type=endpoint_type,
-                litellm_logging_obj=litellm_logging_obj,
-            )
+        raw_bytes: List[bytes] = []
+        logging_scheduled = False
+        # Extract model name for cost injection
+        model_name = PassThroughStreamingHandler._extract_model_for_cost_injection(
+            request_body=request_body,
+            url_route=url_route,
+            endpoint_type=endpoint_type,
+            litellm_logging_obj=litellm_logging_obj,
+        )
 
+        try:
             async for chunk in response.aiter_bytes():
                 raw_bytes.append(chunk)
                 if (
@@ -73,25 +74,37 @@ class PassThroughStreamingHandler:
                             chunk = modified_chunk
 
                 yield chunk
-
-            # After all chunks are processed, handle post-processing
-            end_time = datetime.now()
-
-            asyncio.create_task(
-                PassThroughStreamingHandler._route_streaming_logging_to_handler(
-                    litellm_logging_obj=litellm_logging_obj,
-                    passthrough_success_handler_obj=passthrough_success_handler_obj,
-                    url_route=url_route,
-                    request_body=request_body or {},
-                    endpoint_type=endpoint_type,
-                    start_time=start_time,
-                    raw_bytes=raw_bytes,
-                    end_time=end_time,
-                )
-            )
         except Exception as e:
             verbose_proxy_logger.error(f"Error in chunk_processor: {str(e)}")
             raise
+        finally:
+            # Always log collected chunks for spend tracking, even if the
+            # client disconnects mid-stream. On disconnect, Starlette calls
+            # aclose() on this async generator, which raises GeneratorExit
+            # at the suspended `yield` — `except Exception` does not catch
+            # it, so post-loop logging would otherwise be skipped and all
+            # captured per-chunk usage data lost (e.g. for interrupted
+            # Bedrock streams). See LIT-2642.
+            if not logging_scheduled and raw_bytes:
+                logging_scheduled = True
+                try:
+                    end_time = datetime.now()
+                    asyncio.create_task(
+                        PassThroughStreamingHandler._route_streaming_logging_to_handler(
+                            litellm_logging_obj=litellm_logging_obj,
+                            passthrough_success_handler_obj=passthrough_success_handler_obj,
+                            url_route=url_route,
+                            request_body=request_body or {},
+                            endpoint_type=endpoint_type,
+                            start_time=start_time,
+                            raw_bytes=raw_bytes,
+                            end_time=end_time,
+                        )
+                    )
+                except Exception as e:
+                    verbose_proxy_logger.error(
+                        f"Error scheduling chunk_processor logging: {str(e)}"
+                    )
 
     @staticmethod
     async def _route_streaming_logging_to_handler(
diff --git a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py
new file mode 100644
index 0000000000..a1685c4a1d
--- /dev/null
+++ b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py
@@ -0,0 +1,297 @@
+"""
+Regression tests for LIT-2642 — interrupted streaming responses must still
+flush collected chunks so spend is tracked even when the client disconnects
+mid-stream.
+
+Bedrock invoke streaming was the reported reproducer: the proxy passes the
+upstream stream through `_async_streaming` in `litellm/passthrough/main.py`,
+which collects bytes and triggers `async_flush_passthrough_collected_chunks`
+once the loop completes. When a FastAPI client disconnects mid-stream,
+Starlette calls `aclose()` on the async generator and raises `GeneratorExit`
+at the suspended `yield`. The previous `except Exception` branch did not
+catch `GeneratorExit`, so the post-loop flush was skipped and all per-chunk
+usage data was dropped.
+"""
+
+from typing import List
+from unittest.mock import AsyncMock, MagicMock
+
+import httpx
+import pytest
+
+
+def _make_streaming_response(chunks: List[bytes]):
+    """Build a mock httpx.Response that streams the given chunks via aiter_bytes."""
+    mock = MagicMock(spec=httpx.Response)
+    mock.status_code = 200
+    mock.headers = httpx.Headers({"content-type": "application/vnd.amazon.eventstream"})
+    mock.raise_for_status = MagicMock(return_value=None)
+
+    async def _aiter_bytes():
+        for chunk in chunks:
+            yield chunk
+
+    mock.aiter_bytes = _aiter_bytes
+    mock.aclose = AsyncMock()
+    return mock
+
+
+def _make_logging_obj():
+    mock = MagicMock()
+    mock.async_flush_passthrough_collected_chunks = AsyncMock()
+    return mock
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_flushes_on_normal_completion():
+    """Baseline: full stream consumption flushes collected chunks once."""
+    from litellm.passthrough.main import _async_streaming
+
+    chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
+    mock_response = _make_streaming_response(chunks)
+
+    async def response_coro():
+        return mock_response
+
+    mock_logging_obj = _make_logging_obj()
+    provider_config = MagicMock()
+
+    received = []
+    async for chunk in _async_streaming(
+        response=response_coro(),
+        litellm_logging_obj=mock_logging_obj,
+        provider_config=provider_config,
+    ):
+        received.append(chunk)
+
+    assert received == chunks
+
+    # Allow the scheduled task to run.
+    import asyncio
+
+    await asyncio.sleep(0)
+
+    mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once()
+    call_kwargs = (
+        mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs
+    )
+    assert call_kwargs["raw_bytes"] == chunks
+    assert call_kwargs["provider_config"] is provider_config
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_flushes_on_client_disconnect():
+    """
+    LIT-2642 regression: GeneratorExit (raised when the consumer disconnects
+    mid-stream) must still flush whatever chunks we already collected so
+    spend tracking captures the partial usage.
+    """
+    from litellm.passthrough.main import _async_streaming
+
+    chunks = [
+        b'{"chunk": 1, "outputTokens": 10}',
+        b'{"chunk": 2, "outputTokens": 12}',
+        b'{"chunk": 3, "outputTokens": 8}',
+    ]
+    mock_response = _make_streaming_response(chunks)
+
+    async def response_coro():
+        return mock_response
+
+    mock_logging_obj = _make_logging_obj()
+    provider_config = MagicMock()
+
+    gen = _async_streaming(
+        response=response_coro(),
+        litellm_logging_obj=mock_logging_obj,
+        provider_config=provider_config,
+    )
+
+    # Pull one chunk, then close the generator early — mirrors what
+    # Starlette does when the HTTP client disconnects mid-stream.
+    received = [await gen.__anext__()]
+    await gen.aclose()
+
+    assert received == [chunks[0]]
+
+    # Allow the scheduled flush task to run.
+    import asyncio
+
+    await asyncio.sleep(0)
+
+    mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once()
+    call_kwargs = (
+        mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs
+    )
+    # Only the first chunk was consumed before disconnect; that's what we
+    # must hand off to the cost-tracking flush so partial usage isn't
+    # silently dropped.
+    assert call_kwargs["raw_bytes"] == [chunks[0]]
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_does_not_flush_on_4xx():
+    """Error responses must still raise without entering the flush path."""
+    from litellm.passthrough.main import _async_streaming
+
+    err_response = MagicMock(spec=httpx.Response)
+    err_response.status_code = 429
+
+    def _raise():
+        raise httpx.HTTPStatusError(
+            "429",
+            request=httpx.Request("POST", "https://example.com"),
+            response=httpx.Response(
+                429, request=httpx.Request("POST", "https://example.com")
+            ),
+        )
+
+    err_response.raise_for_status = _raise
+    err_response.aclose = AsyncMock()
+
+    async def response_coro():
+        return err_response
+
+    mock_logging_obj = _make_logging_obj()
+
+    with pytest.raises(httpx.HTTPStatusError):
+        async for _ in _async_streaming(
+            response=response_coro(),
+            litellm_logging_obj=mock_logging_obj,
+            provider_config=MagicMock(),
+        ):
+            pass
+
+    # No bytes were collected, so no flush should have been scheduled.
+    mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_flushes_on_upstream_exception_with_partial_data():
+    """
+    If the upstream connection drops mid-stream and aiter_bytes raises,
+    we still surface the exception, but partial chunks already collected
+    are flushed so spend tracking isn't fully lost.
+    """
+    from litellm.passthrough.main import _async_streaming
+
+    partial_chunks = [b"partial-chunk-1", b"partial-chunk-2"]
+
+    mock_response = MagicMock(spec=httpx.Response)
+    mock_response.status_code = 200
+    mock_response.raise_for_status = MagicMock(return_value=None)
+    mock_response.aclose = AsyncMock()
+
+    async def _aiter_bytes_then_raise():
+        for c in partial_chunks:
+            yield c
+        raise httpx.ReadError("upstream disconnected")
+
+    mock_response.aiter_bytes = _aiter_bytes_then_raise
+
+    async def response_coro():
+        return mock_response
+
+    mock_logging_obj = _make_logging_obj()
+    provider_config = MagicMock()
+
+    received = []
+    with pytest.raises(httpx.ReadError):
+        async for chunk in _async_streaming(
+            response=response_coro(),
+            litellm_logging_obj=mock_logging_obj,
+            provider_config=provider_config,
+        ):
+            received.append(chunk)
+
+    assert received == partial_chunks
+
+    import asyncio
+
+    await asyncio.sleep(0)
+
+    mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once()
+    call_kwargs = (
+        mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs
+    )
+    assert call_kwargs["raw_bytes"] == partial_chunks
+
+
+def test_sync_streaming_flushes_on_normal_completion():
+    """Baseline for the sync codepath."""
+    from litellm.passthrough.main import _sync_streaming
+
+    chunks = [b"a", b"b", b"c"]
+
+    mock_response = MagicMock(spec=httpx.Response)
+
+    def _iter_bytes():
+        yield from chunks
+
+    mock_response.iter_bytes = _iter_bytes
+
+    mock_logging_obj = MagicMock()
+    mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()
+    provider_config = MagicMock()
+
+    # Use a synchronous in-process executor so we can assert immediately.
+    class _ImmediateExecutor:
+        def submit(self, fn, *args, **kwargs):
+            fn(*args, **kwargs)
+
+    from unittest.mock import patch
+
+    with patch("litellm.utils.executor", _ImmediateExecutor()):
+        received = list(
+            _sync_streaming(
+                response=mock_response,
+                litellm_logging_obj=mock_logging_obj,
+                provider_config=provider_config,
+            )
+        )
+
+    assert received == chunks
+    mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once()
+
+
+def test_sync_streaming_flushes_on_early_close():
+    """
+    Sync analog of LIT-2642: closing the generator early must still flush
+    so per-chunk usage data is not silently dropped.
+    """
+    from litellm.passthrough.main import _sync_streaming
+
+    chunks = [b"first", b"second", b"third"]
+
+    mock_response = MagicMock(spec=httpx.Response)
+
+    def _iter_bytes():
+        yield from chunks
+
+    mock_response.iter_bytes = _iter_bytes
+
+    mock_logging_obj = MagicMock()
+    mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()
+    provider_config = MagicMock()
+
+    class _ImmediateExecutor:
+        def submit(self, fn, *args, **kwargs):
+            fn(*args, **kwargs)
+
+    from unittest.mock import patch
+
+    with patch("litellm.utils.executor", _ImmediateExecutor()):
+        gen = _sync_streaming(
+            response=mock_response,
+            litellm_logging_obj=mock_logging_obj,
+            provider_config=provider_config,
+        )
+
+        # Consume one chunk, then close — analog of a client disconnect.
+        first = next(gen)
+        gen.close()
+
+    assert first == chunks[0]
+    mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once()
+    call_kwargs = mock_logging_obj.flush_passthrough_collected_chunks.call_args.kwargs
+    assert call_kwargs["raw_bytes"] == [chunks[0]]
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
new file mode 100644
index 0000000000..3edbfef949
--- /dev/null
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
@@ -0,0 +1,144 @@
+"""
+Regression tests for LIT-2642 — interrupted pass-through streams must still
+trigger logging so spend is tracked.
+
+`PassThroughStreamingHandler.chunk_processor` collects bytes from the
+upstream response and schedules `_route_streaming_logging_to_handler` once
+the chunk loop completes. When a FastAPI client disconnects mid-stream,
+Starlette calls `aclose()` on the async generator and raises `GeneratorExit`
+at the suspended `yield`. The previous `except Exception` branch did not
+catch `GeneratorExit`, so the post-loop logging task was never scheduled.
+"""
+
+from datetime import datetime
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import httpx
+import pytest
+
+from litellm.proxy.pass_through_endpoints.streaming_handler import (
+    PassThroughStreamingHandler,
+)
+from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
+
+
+def _make_streaming_response(chunks):
+    mock = MagicMock(spec=httpx.Response)
+    mock.status_code = 200
+
+    async def _aiter_bytes():
+        for c in chunks:
+            yield c
+
+    mock.aiter_bytes = _aiter_bytes
+    return mock
+
+
+@pytest.mark.asyncio
+async def test_chunk_processor_logs_on_normal_completion():
+    """Baseline: full consumption schedules logging exactly once."""
+    chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
+    response = _make_streaming_response(chunks)
+
+    mock_logging_obj = MagicMock()
+    mock_passthrough_handler = MagicMock()
+
+    with patch.object(
+        PassThroughStreamingHandler,
+        "_route_streaming_logging_to_handler",
+        new=AsyncMock(),
+    ) as mock_route:
+        received = []
+        async for chunk in PassThroughStreamingHandler.chunk_processor(
+            response=response,
+            request_body={"model": "claude-3-haiku"},
+            litellm_logging_obj=mock_logging_obj,
+            endpoint_type=EndpointType.GENERIC,
+            start_time=datetime.now(),
+            passthrough_success_handler_obj=mock_passthrough_handler,
+            url_route="/bedrock/model/claude/invoke-with-response-stream",
+        ):
+            received.append(chunk)
+
+        import asyncio
+
+        await asyncio.sleep(0)
+
+    assert received == chunks
+    mock_route.assert_called_once()
+    call_kwargs = mock_route.call_args.kwargs
+    assert call_kwargs["raw_bytes"] == chunks
+
+
+@pytest.mark.asyncio
+async def test_chunk_processor_logs_on_client_disconnect():
+    """
+    LIT-2642 regression: closing the generator early (e.g. client
+    disconnect) must still schedule logging so per-chunk spend data
+    isn't dropped.
+    """
+    chunks = [b"event-1", b"event-2", b"event-3"]
+    response = _make_streaming_response(chunks)
+
+    mock_logging_obj = MagicMock()
+    mock_passthrough_handler = MagicMock()
+
+    with patch.object(
+        PassThroughStreamingHandler,
+        "_route_streaming_logging_to_handler",
+        new=AsyncMock(),
+    ) as mock_route:
+        gen = PassThroughStreamingHandler.chunk_processor(
+            response=response,
+            request_body={"model": "claude-3-haiku"},
+            litellm_logging_obj=mock_logging_obj,
+            endpoint_type=EndpointType.GENERIC,
+            start_time=datetime.now(),
+            passthrough_success_handler_obj=mock_passthrough_handler,
+            url_route="/bedrock/model/claude/invoke-with-response-stream",
+        )
+
+        # Consume one chunk, then close the generator — same path Starlette
+        # takes when the HTTP client disconnects mid-stream.
+        first = await gen.__anext__()
+        await gen.aclose()
+
+        import asyncio
+
+        await asyncio.sleep(0)
+
+    assert first == chunks[0]
+    mock_route.assert_called_once()
+    call_kwargs = mock_route.call_args.kwargs
+    # Only one chunk made it through before disconnect — that is what
+    # the logging handler must be given so partial usage is captured.
+    assert call_kwargs["raw_bytes"] == [chunks[0]]
+
+
+@pytest.mark.asyncio
+async def test_chunk_processor_does_not_schedule_logging_when_no_chunks():
+    """If no chunks were ever received, don't schedule a no-op logging task."""
+    response = _make_streaming_response([])
+
+    mock_logging_obj = MagicMock()
+    mock_passthrough_handler = MagicMock()
+
+    with patch.object(
+        PassThroughStreamingHandler,
+        "_route_streaming_logging_to_handler",
+        new=AsyncMock(),
+    ) as mock_route:
+        received = []
+        async for chunk in PassThroughStreamingHandler.chunk_processor(
+            response=response,
+            request_body={"model": "claude-3-haiku"},
+            litellm_logging_obj=mock_logging_obj,
+            endpoint_type=EndpointType.GENERIC,
+            start_time=datetime.now(),
+            passthrough_success_handler_obj=mock_passthrough_handler,
+            url_route="/bedrock/model/claude/invoke-with-response-stream",
+        ):
+            received.append(chunk)
+
+    assert received == []
+    mock_route.assert_not_called()

From 8759413312c0ef1fe19520e12189a5c215d5146c Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Tue, 28 Apr 2026 22:52:43 +0000
Subject: [PATCH 2/4] refactor: trim explanatory comments from streaming-flush
 fix

Strip module-level docstrings and per-test/per-block prose from the
LIT-2642 fix and tests. Keep one short comment in each streaming site
that flags the GeneratorExit-vs-Exception subtlety, since that's the
non-obvious reason the flush lives in finally rather than after the loop.

Pure cleanup; no behavior change. All 12 regression tests still pass.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
---
 litellm/passthrough/main.py                   | 19 +----
 .../streaming_handler.py                      | 20 ++----
 ...test_streaming_interrupt_spend_tracking.py | 69 +++----------------
 .../test_streaming_handler_interrupt.py       | 28 +-------
 4 files changed, 17 insertions(+), 119 deletions(-)

diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py
index 2a80e2bb63..9b669d1c2c 100644
--- a/litellm/passthrough/main.py
+++ b/litellm/passthrough/main.py
@@ -397,10 +397,6 @@ def _sync_streaming(
             raw_bytes.append(chunk)
             yield chunk
     finally:
-        # Always flush collected chunks for spend tracking, even if the
-        # consumer terminates the generator early (GeneratorExit). Without
-        # this, an interrupted stream loses all per-chunk usage data
-        # because the post-loop flush never runs. See LIT-2642.
         if not flush_scheduled and raw_bytes:
             flush_scheduled = True
             try:
@@ -410,8 +406,6 @@ def _sync_streaming(
                     provider_config=provider_config,
                 )
             except Exception:
-                # Don't mask the original exception (incl. GeneratorExit)
-                # if scheduling the flush itself fails.
                 pass
 
 
@@ -422,8 +416,6 @@ async def _async_streaming(
 ):
     iter_response = await response
 
-    # Validate response status before consuming the body so 4xx/5xx
-    # responses raise without entering the chunk-collection path.
     try:
         iter_response.raise_for_status()
     except Exception:
@@ -446,12 +438,9 @@ async def _async_streaming(
             pass
         raise
     finally:
-        # Always flush collected chunks for spend tracking, even if the
-        # client disconnects mid-stream. On disconnect, Starlette calls
-        # aclose() on this generator, which raises GeneratorExit at the
-        # suspended `yield` — `except Exception` does not catch it, so
-        # the post-loop flush would otherwise be skipped and all
-        # captured per-chunk usage data lost. See LIT-2642.
+        # GeneratorExit (raised on client disconnect) is not caught by
+        # `except Exception`; the finally block ensures partial usage
+        # still gets flushed for spend tracking. See LIT-2642.
         if not flush_scheduled and raw_bytes:
             flush_scheduled = True
             try:
@@ -462,6 +451,4 @@ async def _async_streaming(
                     )
                 )
             except Exception:
-                # Don't mask the original exception (incl. GeneratorExit)
-                # if scheduling the flush itself fails.
                 pass
diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py
index fe5bd42a60..cbfcd34c43 100644
--- a/litellm/proxy/pass_through_endpoints/streaming_handler.py
+++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py
@@ -36,14 +36,8 @@ class PassThroughStreamingHandler:
         passthrough_success_handler_obj: PassThroughEndpointLogging,
         url_route: str,
     ):
-        """
-        - Yields chunks from the response
-        - Collect non-empty chunks for post-processing (logging)
-        - Inject cost into chunks if include_cost_in_streaming_usage is enabled
-        """
         raw_bytes: List[bytes] = []
         logging_scheduled = False
-        # Extract model name for cost injection
         model_name = PassThroughStreamingHandler._extract_model_for_cost_injection(
             request_body=request_body,
             url_route=url_route,
@@ -59,7 +53,6 @@ class PassThroughStreamingHandler:
                     and model_name
                 ):
                     if endpoint_type == EndpointType.VERTEX_AI:
-                        # Only handle streamRawPredict (uses Anthropic format)
                         if "streamRawPredict" in url_route or "rawPredict" in url_route:
                             modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
                                 chunk, model_name
@@ -78,17 +71,12 @@ class PassThroughStreamingHandler:
             verbose_proxy_logger.error(f"Error in chunk_processor: {str(e)}")
             raise
         finally:
-            # Always log collected chunks for spend tracking, even if the
-            # client disconnects mid-stream. On disconnect, Starlette calls
-            # aclose() on this async generator, which raises GeneratorExit
-            # at the suspended `yield` — `except Exception` does not catch
-            # it, so post-loop logging would otherwise be skipped and all
-            # captured per-chunk usage data lost (e.g. for interrupted
-            # Bedrock streams). See LIT-2642.
+            # GeneratorExit (raised on client disconnect) is not caught by
+            # `except Exception`; the finally block ensures partial usage
+            # still gets logged for spend tracking. See LIT-2642.
             if not logging_scheduled and raw_bytes:
                 logging_scheduled = True
                 try:
-                    end_time = datetime.now()
                     asyncio.create_task(
                         PassThroughStreamingHandler._route_streaming_logging_to_handler(
                             litellm_logging_obj=litellm_logging_obj,
@@ -98,7 +86,7 @@ class PassThroughStreamingHandler:
                             endpoint_type=endpoint_type,
                             start_time=start_time,
                             raw_bytes=raw_bytes,
-                            end_time=end_time,
+                            end_time=datetime.now(),
                         )
                     )
                 except Exception as e:
diff --git a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py
index a1685c4a1d..f3fe3ae5c3 100644
--- a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py
+++ b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py
@@ -1,27 +1,14 @@
-"""
-Regression tests for LIT-2642 — interrupted streaming responses must still
-flush collected chunks so spend is tracked even when the client disconnects
-mid-stream.
-
-Bedrock invoke streaming was the reported reproducer: the proxy passes the
-upstream stream through `_async_streaming` in `litellm/passthrough/main.py`,
-which collects bytes and triggers `async_flush_passthrough_collected_chunks`
-once the loop completes. When a FastAPI client disconnects mid-stream,
-Starlette calls `aclose()` on the async generator and raises `GeneratorExit`
-at the suspended `yield`. The previous `except Exception` branch did not
-catch `GeneratorExit`, so the post-loop flush was skipped and all per-chunk
-usage data was dropped.
-"""
+"""Regression tests for LIT-2642 — interrupted streams must still flush usage."""
 
+import asyncio
 from typing import List
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
 
 
 def _make_streaming_response(chunks: List[bytes]):
-    """Build a mock httpx.Response that streams the given chunks via aiter_bytes."""
     mock = MagicMock(spec=httpx.Response)
     mock.status_code = 200
     mock.headers = httpx.Headers({"content-type": "application/vnd.amazon.eventstream"})
@@ -42,9 +29,13 @@ def _make_logging_obj():
     return mock
 
 
+class _ImmediateExecutor:
+    def submit(self, fn, *args, **kwargs):
+        fn(*args, **kwargs)
+
+
 @pytest.mark.asyncio
 async def test_async_streaming_flushes_on_normal_completion():
-    """Baseline: full stream consumption flushes collected chunks once."""
     from litellm.passthrough.main import _async_streaming
 
     chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
@@ -66,9 +57,6 @@ async def test_async_streaming_flushes_on_normal_completion():
 
     assert received == chunks
 
-    # Allow the scheduled task to run.
-    import asyncio
-
     await asyncio.sleep(0)
 
     mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once()
@@ -81,11 +69,6 @@ async def test_async_streaming_flushes_on_normal_completion():
 
 @pytest.mark.asyncio
 async def test_async_streaming_flushes_on_client_disconnect():
-    """
-    LIT-2642 regression: GeneratorExit (raised when the consumer disconnects
-    mid-stream) must still flush whatever chunks we already collected so
-    spend tracking captures the partial usage.
-    """
     from litellm.passthrough.main import _async_streaming
 
     chunks = [
@@ -107,31 +90,22 @@ async def test_async_streaming_flushes_on_client_disconnect():
         provider_config=provider_config,
     )
 
-    # Pull one chunk, then close the generator early — mirrors what
-    # Starlette does when the HTTP client disconnects mid-stream.
     received = [await gen.__anext__()]
     await gen.aclose()
 
     assert received == [chunks[0]]
 
-    # Allow the scheduled flush task to run.
-    import asyncio
-
     await asyncio.sleep(0)
 
     mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once()
     call_kwargs = (
         mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs
     )
-    # Only the first chunk was consumed before disconnect; that's what we
-    # must hand off to the cost-tracking flush so partial usage isn't
-    # silently dropped.
     assert call_kwargs["raw_bytes"] == [chunks[0]]
 
 
 @pytest.mark.asyncio
 async def test_async_streaming_does_not_flush_on_4xx():
-    """Error responses must still raise without entering the flush path."""
     from litellm.passthrough.main import _async_streaming
 
     err_response = MagicMock(spec=httpx.Response)
@@ -162,17 +136,11 @@ async def test_async_streaming_does_not_flush_on_4xx():
         ):
             pass
 
-    # No bytes were collected, so no flush should have been scheduled.
     mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called()
 
 
 @pytest.mark.asyncio
 async def test_async_streaming_flushes_on_upstream_exception_with_partial_data():
-    """
-    If the upstream connection drops mid-stream and aiter_bytes raises,
-    we still surface the exception, but partial chunks already collected
-    are flushed so spend tracking isn't fully lost.
-    """
     from litellm.passthrough.main import _async_streaming
 
     partial_chunks = [b"partial-chunk-1", b"partial-chunk-2"]
@@ -206,8 +174,6 @@ async def test_async_streaming_flushes_on_upstream_exception_with_partial_data()
 
     assert received == partial_chunks
 
-    import asyncio
-
     await asyncio.sleep(0)
 
     mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once()
@@ -218,7 +184,6 @@ async def test_async_streaming_flushes_on_upstream_exception_with_partial_data()
 
 
 def test_sync_streaming_flushes_on_normal_completion():
-    """Baseline for the sync codepath."""
     from litellm.passthrough.main import _sync_streaming
 
     chunks = [b"a", b"b", b"c"]
@@ -234,13 +199,6 @@ def test_sync_streaming_flushes_on_normal_completion():
     mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()
     provider_config = MagicMock()
 
-    # Use a synchronous in-process executor so we can assert immediately.
-    class _ImmediateExecutor:
-        def submit(self, fn, *args, **kwargs):
-            fn(*args, **kwargs)
-
-    from unittest.mock import patch
-
     with patch("litellm.utils.executor", _ImmediateExecutor()):
         received = list(
             _sync_streaming(
@@ -255,10 +213,6 @@ def test_sync_streaming_flushes_on_normal_completion():
 
 
 def test_sync_streaming_flushes_on_early_close():
-    """
-    Sync analog of LIT-2642: closing the generator early must still flush
-    so per-chunk usage data is not silently dropped.
-    """
     from litellm.passthrough.main import _sync_streaming
 
     chunks = [b"first", b"second", b"third"]
@@ -274,12 +228,6 @@ def test_sync_streaming_flushes_on_early_close():
     mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()
     provider_config = MagicMock()
 
-    class _ImmediateExecutor:
-        def submit(self, fn, *args, **kwargs):
-            fn(*args, **kwargs)
-
-    from unittest.mock import patch
-
     with patch("litellm.utils.executor", _ImmediateExecutor()):
         gen = _sync_streaming(
             response=mock_response,
@@ -287,7 +235,6 @@ def test_sync_streaming_flushes_on_early_close():
             provider_config=provider_config,
         )
 
-        # Consume one chunk, then close — analog of a client disconnect.
         first = next(gen)
         gen.close()
 
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
index 3edbfef949..f73aee77cc 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
@@ -1,15 +1,6 @@
-"""
-Regression tests for LIT-2642 — interrupted pass-through streams must still
-trigger logging so spend is tracked.
-
-`PassThroughStreamingHandler.chunk_processor` collects bytes from the
-upstream response and schedules `_route_streaming_logging_to_handler` once
-the chunk loop completes. When a FastAPI client disconnects mid-stream,
-Starlette calls `aclose()` on the async generator and raises `GeneratorExit`
-at the suspended `yield`. The previous `except Exception` branch did not
-catch `GeneratorExit`, so the post-loop logging task was never scheduled.
-"""
+"""Regression tests for LIT-2642 — interrupted pass-through streams must still log usage."""
 
+import asyncio
 from datetime import datetime
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -36,7 +27,6 @@ def _make_streaming_response(chunks):
 
 @pytest.mark.asyncio
 async def test_chunk_processor_logs_on_normal_completion():
-    """Baseline: full consumption schedules logging exactly once."""
     chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
     response = _make_streaming_response(chunks)
 
@@ -60,8 +50,6 @@ async def test_chunk_processor_logs_on_normal_completion():
         ):
             received.append(chunk)
 
-        import asyncio
-
         await asyncio.sleep(0)
 
     assert received == chunks
@@ -72,11 +60,6 @@ async def test_chunk_processor_logs_on_normal_completion():
 
 @pytest.mark.asyncio
 async def test_chunk_processor_logs_on_client_disconnect():
-    """
-    LIT-2642 regression: closing the generator early (e.g. client
-    disconnect) must still schedule logging so per-chunk spend data
-    isn't dropped.
-    """
     chunks = [b"event-1", b"event-2", b"event-3"]
     response = _make_streaming_response(chunks)
 
@@ -98,26 +81,19 @@ async def test_chunk_processor_logs_on_client_disconnect():
             url_route="/bedrock/model/claude/invoke-with-response-stream",
         )
 
-        # Consume one chunk, then close the generator — same path Starlette
-        # takes when the HTTP client disconnects mid-stream.
         first = await gen.__anext__()
         await gen.aclose()
 
-        import asyncio
-
         await asyncio.sleep(0)
 
     assert first == chunks[0]
     mock_route.assert_called_once()
     call_kwargs = mock_route.call_args.kwargs
-    # Only one chunk made it through before disconnect — that is what
-    # the logging handler must be given so partial usage is captured.
     assert call_kwargs["raw_bytes"] == [chunks[0]]
 
 
 @pytest.mark.asyncio
 async def test_chunk_processor_does_not_schedule_logging_when_no_chunks():
-    """If no chunks were ever received, don't schedule a no-op logging task."""
     response = _make_streaming_response([])
 
     mock_logging_obj = MagicMock()

From 3791abf4bbc9153fdfc9d8a3d5cab849cb754393 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Tue, 28 Apr 2026 22:54:58 +0000
Subject: [PATCH 3/4] fix(passthrough): log when streaming spend-tracking flush
 fails to schedule

Address Greptile feedback: the bare `except Exception: pass` in the
finally blocks of _sync_streaming / _async_streaming silently dropped
errors from executor.submit() / asyncio.create_task() (e.g. saturated
thread pool, closed event loop). Since the entire point of the fix is
that spend tracking should not silently lose data, mirror the peer
streaming_handler.py logging pattern so any scheduling failure is
diagnosable in production.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
---
 litellm/passthrough/main.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py
index 9b669d1c2c..c4c9aea6f6 100644
--- a/litellm/passthrough/main.py
+++ b/litellm/passthrough/main.py
@@ -21,6 +21,7 @@ import httpx
 from httpx._types import CookieTypes, QueryParamTypes, RequestFiles
 
 import litellm
+from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
@@ -405,8 +406,13 @@ def _sync_streaming(
                     raw_bytes=raw_bytes,
                     provider_config=provider_config,
                 )
-            except Exception:
-                pass
+            except Exception as e:
+                verbose_logger.exception(
+                    "Failed to schedule passthrough spend-tracking flush "
+                    "in _sync_streaming; %d buffered chunks dropped: %s",
+                    len(raw_bytes),
+                    e,
+                )
 
 
 async def _async_streaming(
@@ -450,5 +456,10 @@ async def _async_streaming(
                         provider_config=provider_config,
                     )
                 )
-            except Exception:
-                pass
+            except Exception as e:
+                verbose_logger.exception(
+                    "Failed to schedule passthrough spend-tracking flush "
+                    "in _async_streaming; %d buffered chunks dropped: %s",
+                    len(raw_bytes),
+                    e,
+                )

From 793a35dfe2406c803a24a2b2174c6cc8dff970ae Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Thu, 30 Apr 2026 03:09:52 +0000
Subject: [PATCH 4/4] test(prometheus): update master-key hash assertions to
 alias

PR #26484 substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for
hash_token(master_key) in UserAPIKeyAuth so the master key (or its
hash) never reaches spend logs / metrics. The otel prometheus tests
still hardcoded the SHA-256 of "sk-1234"
("88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"),
so the metric labels no longer matched and test_proxy_failure_metrics
failed. Reference the alias constant directly.

https://claude.ai/code/session_01UkzyZKiADEkZDbZFwB98yV

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
---
 tests/otel_tests/test_prometheus.py | 40 +++++++++++++++++++----------
 1 file changed, 26 insertions(+), 14 deletions(-)

diff --git a/tests/otel_tests/test_prometheus.py b/tests/otel_tests/test_prometheus.py
index 2d772c4a63..75061dda94 100644
--- a/tests/otel_tests/test_prometheus.py
+++ b/tests/otel_tests/test_prometheus.py
@@ -115,6 +115,13 @@ async def test_proxy_failure_metrics():
             "litellm_llm_api_failed_requests_metric_total{",  # Deprecated but may still be used
         ]
 
+        # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for
+        # hash_token(master_key) so the master key (or its hash) never
+        # propagates into metrics. See PR #26484.
+        from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
+
+        expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS
+
         # Check if either pattern is in metrics and contains required fields
         found_metric = False
         for pattern in expected_patterns:
@@ -125,8 +132,7 @@ async def test_proxy_failure_metrics():
                         'api_key_alias="None"' in line
                         and 'exception_class="Openai.RateLimitError"' in line
                         and 'exception_status="429"' in line
-                        and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
-                        in line
+                        and f'hashed_api_key="{expected_hashed_api_key}"' in line
                         and 'requested_model="fake-azure-endpoint"' in line
                         and 'route="/chat/completions"' in line
                     ):
@@ -135,8 +141,7 @@ async def test_proxy_failure_metrics():
                 # For deprecated llm_api metric, check llm-specific fields
                 elif "litellm_llm_api_failed_requests_metric_total{" in line:
                     if (
-                        'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
-                        in line
+                        f'hashed_api_key="{expected_hashed_api_key}"' in line
                         and 'model="429"' in line
                     ):  # The deprecated metric uses the actual model from the request
                         found_metric = True
@@ -156,8 +161,7 @@ async def test_proxy_failure_metrics():
         for line in metrics.split("\n"):
             if (
                 total_requests_pattern in line
-                and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
-                in line
+                and f'hashed_api_key="{expected_hashed_api_key}"' in line
                 and 'requested_model="fake-azure-endpoint"' in line
                 and 'status_code="429"' in line
             ):
@@ -195,6 +199,12 @@ async def test_proxy_success_metrics():
 
         assert END_USER_ID not in metrics
 
+        # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for
+        # hash_token(master_key) (PR #26484).
+        from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
+
+        expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS
+
         # Check if the success metric is present and correct - use flexible matching
         # Check for request_total_latency_metric with required fields
         # Note: The model can be "gpt-3.5-turbo-0301" or similar depending on what's returned
@@ -203,8 +213,7 @@ async def test_proxy_success_metrics():
             if (
                 "litellm_request_total_latency_metric_bucket{" in line
                 and 'api_key_alias="None"' in line
-                and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
-                in line
+                and f'hashed_api_key="{expected_hashed_api_key}"' in line
                 and 'requested_model="fake-openai-endpoint"' in line
                 and 'le="0.005"' in line
             ):
@@ -221,8 +230,7 @@ async def test_proxy_success_metrics():
             if (
                 "litellm_llm_api_latency_metric_bucket{" in line
                 and 'api_key_alias="None"' in line
-                and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
-                in line
+                and f'hashed_api_key="{expected_hashed_api_key}"' in line
                 and 'requested_model="fake-openai-endpoint"' in line
                 and 'le="0.005"' in line
             ):
@@ -298,6 +306,12 @@ async def test_proxy_fallback_metrics():
 
         print("/metrics", metrics)
 
+        # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for
+        # hash_token(master_key) (PR #26484).
+        from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
+
+        expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS
+
         # Check if successful fallback metric is incremented - use flexible matching
         found_successful_fallback = False
         for line in metrics.split("\n"):
@@ -307,8 +321,7 @@ async def test_proxy_fallback_metrics():
                 and 'exception_class="Openai.RateLimitError"' in line
                 and 'exception_status="429"' in line
                 and 'fallback_model="fake-openai-endpoint"' in line
-                and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
-                in line
+                and f'hashed_api_key="{expected_hashed_api_key}"' in line
                 and 'requested_model="fake-azure-endpoint"' in line
                 and "1.0" in line
             ):
@@ -328,8 +341,7 @@ async def test_proxy_fallback_metrics():
                 and 'exception_class="Openai.RateLimitError"' in line
                 and 'exception_status="429"' in line
                 and 'fallback_model="unknown-model"' in line
-                and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"'
-                in line
+                and f'hashed_api_key="{expected_hashed_api_key}"' in line
                 and 'requested_model="fake-azure-endpoint"' in line
                 and "1.0" in line
             ):