fix(stream): estimate input tokens when upstream reports prompt_tokens=0 (#4615)

Integrated into release/v3.8.34 (rebuilt onto tip)

open-sse/utils/stream.ts

@@ -1545,6 +1545,26 @@ export function createSSEStream(options: StreamOptions = {}) {
                   if (Array.isArray(parsed.choices) && parsed.choices.length === 0) {
                     const emptyChoicesUsage = extractUsage(parsed) ?? parsed.usage;
                     if (hasValidUsage(emptyChoicesUsage)) {
+                      // Some upstreams (e.g. Ollama Cloud) emit prompt_tokens: 0
+                      // even when input was sent — they simply don't count input
+                      // tokens.  When we have a non-zero output but zero input,
+                      // estimate the real input token count from the request body.
+                      if (
+                        emptyChoicesUsage &&
+                        typeof emptyChoicesUsage === "object" &&
+                        !Array.isArray(emptyChoicesUsage) &&
+                        emptyChoicesUsage.completion_tokens > 0
+                      ) {
+                        const pt = emptyChoicesUsage.prompt_tokens ?? 0;
+                        if (pt === 0) {
+                          const estimated = estimateUsage(body, totalContentLength, FORMATS.OPENAI);
+                          if (estimated?.prompt_tokens > 0) {
+                            emptyChoicesUsage.prompt_tokens = estimated.prompt_tokens;
+                            emptyChoicesUsage.total_tokens =
+                              (emptyChoicesUsage.total_tokens ?? 0) + estimated.prompt_tokens;
+                          }
+                        }
+                      }
                       usage = emptyChoicesUsage;
                       output = `data: ${JSON.stringify(parsed)}\n\n`;
                       injectedUsage = true;

tests/unit/stream-prompt-tokens-zero-upstream.test.ts

@@ -0,0 +1,114 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+import fs from "node:fs";
+import os from "node:os";
+import path from "node:path";
+
+const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-stream-zero-prompt-"));
+process.env.DATA_DIR = TEST_DATA_DIR;
+const core = await import("../../src/lib/db/core.ts");
+
+const { createSSEStream } = await import("../../open-sse/utils/stream.ts");
+const { FORMATS } = await import("../../open-sse/translator/formats.ts");
+
+const textEncoder = new TextEncoder();
+
+async function readTransformed(chunks, options) {
+  const source = new ReadableStream({
+    start(controller) {
+      for (const chunk of chunks) {
+        controller.enqueue(textEncoder.encode(chunk));
+      }
+      controller.close();
+    },
+  });
+
+  return new Response(source.pipeThrough(createSSEStream(options))).text();
+}
+
+test("createSSEStream passthrough estimates input tokens when upstream reports prompt_tokens=0 (Ollama Cloud pattern)", async () => {
+  let onCompletePayload = null;
+  const body = { messages: [{ role: "user", content: "Write a 500-line research report" }] };
+  const text = await readTransformed(
+    [
+      `data: ${JSON.stringify({
+        id: "chatcmpl_ollama",
+        object: "chat.completion.chunk",
+        created: 1,
+        model: "minimax-m3",
+        system_fingerprint: "fp_ollama",
+        choices: [{ index: 0, delta: { content: "Here" } }],
+      })}\n\n`,
+      `data: ${JSON.stringify({
+        id: "chatcmpl_ollama",
+        object: "chat.completion.chunk",
+        created: 1,
+        model: "minimax-m3",
+        system_fingerprint: "fp_ollama",
+        choices: [{ index: 0, delta: { content: " is" } }],
+      })}\n\n`,
+      `data: ${JSON.stringify({
+        id: "chatcmpl_ollama",
+        object: "chat.completion.chunk",
+        created: 1,
+        model: "minimax-m3",
+        system_fingerprint: "fp_ollama",
+        choices: [{ index: 0, delta: { content: " the" } }],
+      })}\n\n`,
+      `data: ${JSON.stringify({
+        id: "chatcmpl_ollama",
+        object: "chat.completion.chunk",
+        created: 1,
+        model: "minimax-m3",
+        system_fingerprint: "fp_ollama",
+        choices: [{ index: 0, delta: { content: " response" } }],
+      })}\n\n`,
+      // Ollama Cloud returns usage with prompt_tokens: 0 — it doesn't count input tokens
+      `data: ${JSON.stringify({
+        id: "chatcmpl_ollama",
+        object: "chat.completion.chunk",
+        created: 1,
+        model: "minimax-m3",
+        system_fingerprint: "fp_ollama",
+        choices: [],
+        usage: { prompt_tokens: 0, completion_tokens: 10, total_tokens: 0 },
+      })}\n\n`,
+      `data: [DONE]\n\n`,
+    ],
+    {
+      mode: "passthrough",
+      sourceFormat: FORMATS.OPENAI,
+      provider: "ollamacloud",
+      model: "minimax-m3",
+      body,
+      onComplete(payload) {
+        onCompletePayload = payload;
+      },
+    }
+  );
+
+  // The fix should estimate input tokens from the request body instead of showing 0
+  assert.equal(
+    onCompletePayload.responseBody.usage.completion_tokens,
+    10,
+    "completion tokens should be unchanged"
+  );
+  assert.ok(
+    onCompletePayload.responseBody.usage.prompt_tokens > 0,
+    `prompt_tokens should be estimated (> 0), got ${onCompletePayload.responseBody.usage.prompt_tokens}`
+  );
+  assert.equal(
+    onCompletePayload.responseBody.usage.total_tokens,
+    onCompletePayload.responseBody.usage.prompt_tokens + 10,
+    "total_tokens should equal prompt_tokens + completion_tokens"
+  );
+});
+
+test.after(() => {
+  core.resetDbInstance();
+  if (fs.existsSync(TEST_DATA_DIR)) {
+    for (const entry of fs.readdirSync(TEST_DATA_DIR)) {
+      fs.rmSync(path.join(TEST_DATA_DIR, entry), { recursive: true, force: true });
+    }
+  }
+});