From d4296d6087b0fcaf02ba4f3c5b83000b27db22db Mon Sep 17 00:00:00 2001
From: 9qeklajc <tthefux@gmail.com>
Date: Wed, 3 Jun 2026 17:04:48 +0200
Subject: [PATCH] fix usage caturing

---
 routstr/upstream/base.py                   | 27 +++++++++++++
 tests/unit/test_streaming_sse_providers.py | 45 ++++++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py
index e335e1bc..51af3a05 100644
--- a/routstr/upstream/base.py
+++ b/routstr/upstream/base.py
@@ -790,6 +790,33 @@ class BaseUpstreamProvider:
                             self._current_stream_id = f"chatcmpl-{uuid.uuid4()}"
                         obj["id"] = self._current_stream_id
                     if isinstance(obj.get("usage"), dict):
+                        # Capture usage for end-of-stream cost reconciliation.
+                        # Some models (e.g. Gemini thinking models over the
+                        # OpenAI-compat endpoint) attach ``usage`` to the SAME
+                        # chunk that carries the final content/finish_reason
+                        # rather than sending a separate ``choices: []`` usage
+                        # chunk. Only swallow the chunk when it is a pure usage
+                        # chunk (no choices); otherwise the content would be
+                        # silently dropped and the client would receive no
+                        # assistant message at all.
+                        if obj.get("choices"):
+                            # Capture usage (with model) for the cost trailer,
+                            # but with choices stripped so the trailer never
+                            # re-emits this chunk's content.
+                            usage_chunk_data = {
+                                k: v for k, v in obj.items() if k != "choices"
+                            }
+                            usage_chunk_data["choices"] = []
+                            # Forward the content now, without usage, so token
+                            # usage is reported exactly once (in the trailer).
+                            forward = {k: v for k, v in obj.items() if k != "usage"}
+                            yield (
+                                prefix
+                                + b"data: "
+                                + json.dumps(forward).encode()
+                                + b"\n\n"
+                            )
+                            return
                         usage_chunk_data = obj
                         return
                     yield prefix + b"data: " + json.dumps(obj).encode() + b"\n\n"
diff --git a/tests/unit/test_streaming_sse_providers.py b/tests/unit/test_streaming_sse_providers.py
index 17cd632e..82d909fd 100644
--- a/tests/unit/test_streaming_sse_providers.py
+++ b/tests/unit/test_streaming_sse_providers.py
@@ -254,6 +254,51 @@ async def test_openrouter_mid_stream_error_event() -> None:
     assert any("error" in o for o in objs), "error event must be forwarded intact"
 
 
+@pytest.mark.asyncio
+async def test_gemini_combined_content_and_usage_chunk() -> None:
+    """Gemini thinking models pack usage into the final *content* chunk.
+
+    Regression: the parser swallowed any chunk carrying a ``usage`` dict, so
+    when content + usage arrived together the assistant text was dropped and
+    the client saw "no assistant messages" despite a 200 + token accounting.
+    """
+    chunks = [
+        b'data: {"id":"g","choices":[{"delta":{"content":"the answer"},'
+        b'"finish_reason":"stop"}],"usage":{"prompt_tokens":3,'
+        b'"completion_tokens":2,"total_tokens":5}}\n\n',
+        b"data: [DONE]\n\n",
+    ]
+    out = await _drive(chunks)
+    objs = _assert_clean(out)
+    contents = [
+        c["delta"]["content"]
+        for o in objs
+        for c in o.get("choices", [])
+        if "delta" in c
+    ]
+    # Content delivered exactly once (not dropped, not duplicated by the trailer).
+    assert contents == ["the answer"]
+
+
+@pytest.mark.asyncio
+async def test_separate_usage_chunk_not_forwarded_as_content() -> None:
+    """A pure usage chunk (choices: []) is still swallowed, content intact."""
+    chunks = [
+        b'data: {"id":"x","choices":[{"delta":{"content":"hello"}}]}\n\n',
+        b'data: {"id":"x","choices":[],"usage":{"total_tokens":4}}\n\n',
+        b"data: [DONE]\n\n",
+    ]
+    out = await _drive(chunks)
+    objs = _assert_clean(out)
+    contents = [
+        c["delta"]["content"]
+        for o in objs
+        for c in o.get("choices", [])
+        if "delta" in c
+    ]
+    assert contents == ["hello"]
+
+
 @pytest.mark.asyncio
 async def test_requested_model_override_applied() -> None:
     """Model rewriting still works through the buffered parser."""