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>
This commit is contained in:
Kashyab Ambarani
2026-09-25 03:38:47 -07:00
co-authored by Max19970 SannaMarcoDev Jeremias Volker
parent b9e5ae098c
commit 5d6090ae46
2 changed files with 99 additions and 17 deletions
+86
View File
@@ -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");
});