From 3cf4bbae30f43d6c07980eba3f1692483f040f8a Mon Sep 17 00:00:00 2001
From: Jon Saad-Falcon <41205309+jonsaadfalcon@users.noreply.github.com>
Date: Sat, 14 Mar 2026 14:55:36 -0700
Subject: [PATCH] feat: display thinking tokens in X-Ray footer and stream
 usage data
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Ollama engine: capture eval_count and prompt_eval_count from the
  streaming final chunk (includes thinking/reasoning tokens).
- Routes: include usage data in the SSE finish chunk so the frontend
  gets accurate token counts even in streaming mode.
- X-Ray footer: show estimated thinking tokens when completion_tokens
  significantly exceeds visible output (e.g. "1695 generated · 11 prompt
  · ~1680 thinking").
---
 frontend/src/components/Chat/XRayFooter.tsx | 14 +++++++++++++-
 src/openjarvis/engine/ollama.py             | 11 +++++++++++
 src/openjarvis/server/routes.py             | 21 ++++++++++++++++++---
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/frontend/src/components/Chat/XRayFooter.tsx b/frontend/src/components/Chat/XRayFooter.tsx
index 2b1dde95..564369b9 100644
--- a/frontend/src/components/Chat/XRayFooter.tsx
+++ b/frontend/src/components/Chat/XRayFooter.tsx
@@ -35,7 +35,19 @@ export function XRayFooter({ usage, telemetry }: Props) {
     rows.push({ label: 'Engine', value: `${telemetry.engine}${modelDetail ? ` (${modelDetail})` : ''}` });
   }
   if (usage) {
-    rows.push({ label: 'Tokens', value: `${usage.completion_tokens} generated \u00B7 ${usage.prompt_tokens} prompt` });
+    const tokenParts = [`${usage.completion_tokens} generated`, `${usage.prompt_tokens} prompt`];
+    // Estimate thinking tokens: if total generated >> visible output, the
+    // difference is internal reasoning (e.g. Qwen3.5 thinking mode).
+    if (telemetry?.tokens_per_sec && telemetry.total_ms && usage.completion_tokens > 50) {
+      const visibleEstimate = Math.ceil(
+        (telemetry.total_ms / 1000) * telemetry.tokens_per_sec,
+      );
+      if (usage.completion_tokens > visibleEstimate * 1.5) {
+        const thinking = usage.completion_tokens - visibleEstimate;
+        tokenParts.push(`~${thinking} thinking`);
+      }
+    }
+    rows.push({ label: 'Tokens', value: tokenParts.join(' \u00B7 ') });
   }
   if (telemetry?.tokens_per_sec) {
     rows.push({ label: 'Speed', value: `${Math.round(telemetry.tokens_per_sec)} tok/s` });
diff --git a/src/openjarvis/engine/ollama.py b/src/openjarvis/engine/ollama.py
index 049c489d..5e95ef9a 100644
--- a/src/openjarvis/engine/ollama.py
+++ b/src/openjarvis/engine/ollama.py
@@ -41,6 +41,8 @@ class OllamaEngine(InferenceEngine):
             host = env_host or self._DEFAULT_HOST
         self._host = host.rstrip("/")
         self._client = httpx.Client(base_url=self._host, timeout=timeout)
+        # Last stream usage — captured from Ollama's final chunk
+        self._last_stream_usage: Dict[str, int] = {}
 
     def generate(
         self,
@@ -174,6 +176,15 @@ class OllamaEngine(InferenceEngine):
                     if content:
                         yield content
                     if chunk.get("done", False):
+                        # Capture usage from final chunk
+                        self._last_stream_usage = {
+                            "prompt_tokens": chunk.get("prompt_eval_count", 0),
+                            "completion_tokens": chunk.get("eval_count", 0),
+                            "total_tokens": (
+                                chunk.get("prompt_eval_count", 0)
+                                + chunk.get("eval_count", 0)
+                            ),
+                        }
                         break
         except (httpx.ConnectError, httpx.TimeoutException) as exc:
             raise EngineConnectionError(
diff --git a/src/openjarvis/server/routes.py b/src/openjarvis/server/routes.py
index 5fc4e62c..30dab503 100644
--- a/src/openjarvis/server/routes.py
+++ b/src/openjarvis/server/routes.py
@@ -222,8 +222,9 @@ async def _handle_stream(engine, model: str, req: ChatCompletionRequest):
             yield "data: [DONE]\n\n"
             return
 
-        # Send finish chunk
-        finish_chunk = ChatCompletionChunk(
+        # Send finish chunk with usage data if available
+        import json as _json
+        finish_data = ChatCompletionChunk(
             id=chunk_id,
             model=model,
             choices=[StreamChoice(
@@ -231,7 +232,21 @@ async def _handle_stream(engine, model: str, req: ChatCompletionRequest):
                 finish_reason="stop",
             )],
         )
-        yield f"data: {finish_chunk.model_dump_json()}\n\n"
+        finish_dict = _json.loads(finish_data.model_dump_json())
+
+        # Pull usage from the engine if it tracked it during streaming
+        raw_engine = engine
+        # Unwrap InstrumentedEngine if present
+        if hasattr(raw_engine, "_inner"):
+            raw_engine = raw_engine._inner
+        # Unwrap MultiEngine if present
+        if hasattr(raw_engine, "_engine_for"):
+            raw_engine = raw_engine._engine_for(model)
+        stream_usage = getattr(raw_engine, "_last_stream_usage", None)
+        if isinstance(stream_usage, dict) and stream_usage.get("total_tokens", 0) > 0:
+            finish_dict["usage"] = stream_usage
+
+        yield f"data: {_json.dumps(finish_dict)}\n\n"
         yield "data: [DONE]\n\n"
 
     return StreamingResponse(