fix: preserve streaming tool arguments and complete usage totals

Consolidates #2, #3 and #4 with split-frame regression tests. Uses Pi's streaming JSON parser and preserves usage across empty trailers.

Co-authored-by: Max19970 <maxshubert@yandex.ru>
Co-authored-by: SannaMarcoDev <sannamarco.dev@gmail.com>
Co-authored-by: Jeremias Volker <jeremias.volker@uniklinik-freiburg.de>

src/stream.ts

@@ -9,6 +9,7 @@ import {
   type TranscriptContext,
   calculateCost,
   createAssistantMessageEventStream,
+  parseStreamingJson,
 } from "@earendil-works/pi-ai";
 import { mapContextToChat, type ChatHistoryItem, type ContentPart, type ToolDef } from "./context-map.js";
 import { getCachedUserJwt } from "./jwt.js";
@@ -211,12 +212,12 @@ function decodeUsage(buf: Buffer): CloudChatEvent | null {
     else if (metric.includes("cached") || metric.includes("cache_read")) cachedInputTokens = n;
     else if (metric.includes("cache_creation")) cacheCreationInputTokens = n;
   }
-  if (promptTokens === undefined && completionTokens === undefined) return null;
+  if ([promptTokens, completionTokens, cachedInputTokens, cacheCreationInputTokens].every((n) => n === undefined)) return null;
   return {
     kind: "usage",
     promptTokens,
     completionTokens,
-    totalTokens: (promptTokens ?? 0) + (completionTokens ?? 0),
+    totalTokens: (promptTokens ?? 0) + (completionTokens ?? 0) + (cachedInputTokens ?? 0) + (cacheCreationInputTokens ?? 0),
     cachedInputTokens,
     cacheCreationInputTokens,
   };
@@ -424,11 +425,7 @@ export function streamDevin(
       if (toolIndex < 0) return;
       const block = output.content[toolIndex];
       if (block.type === "toolCall") {
-        try {
+        block.arguments = parseStreamingJson(partialJson);
-          block.arguments = JSON.parse(partialJson);
-        } catch {
-          // keep last parsed object
-        }
         stream.push({
           type: "toolcall_end",
           contentIndex: toolIndex,
@@ -498,11 +495,7 @@ export function streamDevin(
           partialJson += event.argsDelta;
           const block = output.content[toolIndex];
           if (block.type === "toolCall") {
-            try {
+            block.arguments = parseStreamingJson(partialJson);
-              block.arguments = JSON.parse(partialJson);
-            } catch {
-              // incomplete json
-            }
           }
           stream.push({ type: "toolcall_delta", contentIndex: toolIndex, delta: event.argsDelta, partial: output });
         } else if (event.kind === "finish") {
@@ -512,11 +505,14 @@ export function streamDevin(
           output.stopReason =
             event.reason === "tool_calls" ? "toolUse" : event.reason === "length" ? "length" : "stop";
         } else if (event.kind === "usage") {
-          output.usage.input = event.promptTokens ?? 0;
+          // Empty trailers must not erase usage already received. Metrics are
-          output.usage.output = event.completionTokens ?? 0;
+          // snapshots, and a frame may contain only some of the components.
-          output.usage.cacheRead = event.cachedInputTokens ?? 0;
+          if (!event.totalTokens) continue;
-          output.usage.cacheWrite = event.cacheCreationInputTokens ?? 0;
+          output.usage.input = event.promptTokens ?? output.usage.input;
-          output.usage.totalTokens = event.totalTokens ?? output.usage.input + output.usage.output;
+          output.usage.output = event.completionTokens ?? output.usage.output;
+          output.usage.cacheRead = event.cachedInputTokens ?? output.usage.cacheRead;
+          output.usage.cacheWrite = event.cacheCreationInputTokens ?? output.usage.cacheWrite;
+          output.usage.totalTokens = output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite;
           calculateCost(model, output.usage);
         }
       }

test/stream-deltas.test.mjs

@@ -0,0 +1,86 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { setImmediate } from "node:timers/promises";
+import { streamDevin } from "../src/stream.ts";
+import { clearCachedUserJwt } from "../src/jwt.ts";
+import { encodeMessage, encodeString, encodeTag, encodeVarintField, frameConnectStream } from "../src/wire.ts";
+
+const model = {
+  id: "test-model", name: "Test", api: "devin-local", provider: "devin",
+  baseUrl: "https://devin.invalid", reasoning: false, input: ["text"],
+  cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 1.25 },
+  contextWindow: 1_000_000, maxTokens: 1000,
+};
+
+async function run(t, frames) {
+  clearCachedUserJwt();
+  t.after(clearCachedUserJwt);
+  t.mock.method(globalThis, "fetch", async (url) => {
+    if (url === "https://devin.invalid/exa.auth_pb.AuthService/GetUserJwt") {
+      return new Response(encodeString(1, "eyJtest.jwt"));
+    }
+    assert.equal(url, "https://devin.invalid/exa.api_server_pb.ApiServerService/GetChatMessage");
+    const chunks = [...frames.map((frame) => frameConnectStream(frame, false)), Buffer.from([2, 0, 0, 0, 2, 123, 125])];
+    return new Response(new ReadableStream({
+      async pull(controller) {
+        await setImmediate();
+        if (chunks.length) controller.enqueue(chunks.shift());
+        else controller.close();
+      },
+    }));
+  });
+  const stream = streamDevin(model, { messages: [] }, {
+    apiKey: "synthetic-test-key", env: { DEVIN_API_SERVER_URL: "https://devin.invalid" },
+  });
+  const events = [];
+  for await (const event of stream) events.push(structuredClone(event));
+  const result = await stream.result();
+  assert.notEqual(result.stopReason, "error", result.errorMessage);
+  return { events, result };
+}
+
+function usage(metrics) {
+  return encodeMessage(28, Buffer.concat(Object.entries(metrics).map(([name, value]) => {
+    const float = Buffer.alloc(4);
+    float.writeFloatLE(value);
+    return encodeMessage(2, Buffer.concat([
+      encodeString(5, name), encodeMessage(4, Buffer.concat([encodeTag(2, 5), float])),
+    ]));
+  })));
+}
+
+const metrics = { input_tokens: 461, output_tokens: 262, cached_input_tokens: 228_162, cache_creation_input_tokens: 17 };
+
+test("counts all four usage components and preserves them after an empty frame (#3, #4)", async (t) => {
+  const { result } = await run(t, [usage(metrics), usage(Object.fromEntries(Object.keys(metrics).map((key) => [key, 0])))]);
+  assert.deepEqual({ ...result.usage, cost: undefined }, {
+    input: 461, output: 262, cacheRead: 228_162, cacheWrite: 17, totalTokens: 228_902, cost: undefined,
+  });
+  assert.equal(result.usage.cost.input, 461 / 1_000_000);
+  assert.equal(result.usage.cost.cacheRead, 228_162 * 0.1 / 1_000_000);
+});
+
+test("accepts cache-only usage and combines separate component snapshots without double counting", async (t) => {
+  const { result } = await run(t, [
+    usage({ cached_input_tokens: 58_905 }), usage({ input_tokens: 280, output_tokens: 118 }),
+    usage({ input_tokens: 280, output_tokens: 118 }), usage({ cache_creation_input_tokens: 17 }),
+  ]);
+  assert.equal(result.usage.totalTokens, 59_320);
+  assert.equal(result.usage.cacheRead, 58_905);
+  assert.equal(result.usage.cacheWrite, 17);
+});
+
+test("keeps growing tool arguments across split keys, strings, and Unicode escapes (#2)", async (t) => {
+  const fragments = ['{"path":"file.txt","cont', 'ent":"hello ', '\\u26', '3a"}'];
+  const frames = fragments.map((part, index) => encodeMessage(6, Buffer.concat([
+    ...(index === 0 ? [encodeString(1, "call-1"), encodeString(2, "write")] : []), encodeString(3, part),
+  ])));
+  frames.push(encodeVarintField(5, 10));
+  const { events, result } = await run(t, frames);
+  const deltas = events.filter((event) => event.type === "toolcall_delta");
+  assert.equal(deltas.length, 4);
+  for (const event of deltas) assert.equal(event.partial.content[0].arguments.path, "file.txt");
+  assert.equal(deltas[1].partial.content[0].arguments.content, "hello");
+  assert.deepEqual(result.content[0].arguments, { path: "file.txt", content: "hello ☺" });
+  assert.equal(result.stopReason, "toolUse");
+});