From 7adc05096810ba11f2e88653bcace7eb9db29005 Mon Sep 17 00:00:00 2001
From: wangwei <Wei.Wang@t-systems.com>
Date: Thu, 23 Jul 2026 13:42:33 +0800
Subject: [PATCH] feat: record streaming token usage in
 TrackedLLMClient.stream_chat

Implement manual generator driving using next()/StopIteration to capture
the return value (trailing usage dict) from inner stream_chat() implementations,
enabling token tracking for streaming LLM calls.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 backend/app/services/llm/tracked_client.py    | 21 ++++++++++++-------
 .../observability/test_tracked_client.py      | 20 ++++++++++++++++++
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/backend/app/services/llm/tracked_client.py b/backend/app/services/llm/tracked_client.py
index 1cc249f..d94a128 100644
--- a/backend/app/services/llm/tracked_client.py
+++ b/backend/app/services/llm/tracked_client.py
@@ -56,18 +56,24 @@ class TrackedLLMClient:
         return response
 
     def stream_chat(self, messages: List[Dict[str, str]], *args: Any, **kwargs: Any):
-        """Delegate to the wrapped client's stream_chat(), recording call outcome only.
+        """Delegate to the wrapped client's stream_chat(), recording call outcome and usage.
 
-        Token usage is NOT recorded here: none of the current provider
-        stream_chat() implementations parse a trailing usage chunk from the
-        gateway (see the design doc's Known Limitations), so accumulating a
-        token count here would silently be wrong. Only call success/failure
-        and latency are tracked for streaming calls.
+        Drives the inner generator manually (instead of a plain `for` loop) so
+        it can capture the generator's return value via StopIteration.value —
+        the trailing token-usage dict the inner client captures from a
+        stream_options.include_usage chunk, if the gateway sent one.
         """
         start = time.time()
         error: Optional[str] = None
+        usage: Optional[Dict[str, int]] = None
+        gen = self._inner.stream_chat(messages, *args, **kwargs)
         try:
-            for chunk in self._inner.stream_chat(messages, *args, **kwargs):
+            while True:
+                try:
+                    chunk = next(gen)
+                except StopIteration as stop:
+                    usage = stop.value
+                    break
                 yield chunk
         except Exception as exc:  # noqa: BLE001 - report, then re-raise unchanged
             error = str(exc)
@@ -77,6 +83,7 @@ class TrackedLLMClient:
                 provider=self._inner.config.provider.value,
                 model=self._inner.config.model,
                 success=error is None,
+                usage=usage,
                 latency_ms=int((time.time() - start) * 1000),
                 error=error,
             )
diff --git a/backend/tests/observability/test_tracked_client.py b/backend/tests/observability/test_tracked_client.py
index 3b59009..51718cd 100644
--- a/backend/tests/observability/test_tracked_client.py
+++ b/backend/tests/observability/test_tracked_client.py
@@ -83,3 +83,23 @@ def test_stream_chat_records_call_without_token_usage():
     entry = tracker.get("deepseek", "deepseek-v4-flash")
     assert entry.call_count_ok == 1
     assert entry.total_tokens == 0
+
+
+def test_stream_chat_records_usage_from_generator_return_value():
+    """stream_chat() must forward the inner generator's returned usage dict to record()."""
+    inner = _make_inner()
+
+    def fake_stream(*args, **kwargs):
+        yield "chunk-1"
+        yield "chunk-2"
+        return {"prompt_tokens": 6, "completion_tokens": 2, "total_tokens": 8}
+
+    inner.stream_chat.side_effect = fake_stream
+    tracker = ModelUsageTracker()
+
+    chunks = list(TrackedLLMClient(inner, tracker).stream_chat([{"role": "user", "content": "hi"}]))
+
+    assert chunks == ["chunk-1", "chunk-2"]
+    entry = tracker.get("deepseek", "deepseek-v4-flash")
+    assert entry.total_tokens == 8
+    assert entry.call_count_ok == 1