diff --git a/packages/adapter-utils/src/usage-checkpoint.test.ts b/packages/adapter-utils/src/usage-checkpoint.test.ts index 2b622ef60f0..a158f587c15 100644 --- a/packages/adapter-utils/src/usage-checkpoint.test.ts +++ b/packages/adapter-utils/src/usage-checkpoint.test.ts @@ -29,6 +29,29 @@ describe("usage checkpoint stream", () => { expect(parser.mock.calls[0][0]).not.toContain("private content"); }); + it("bounds large content strings before buffering and preserves usage on either side", async () => { + const saved = vi.fn(); const parser = vi.fn(createParser()); const output = vi.fn(); + const log = createUsageCheckpointLog(output, saved, parser); + const line = JSON.stringify({ type: "result", usage: { inputTokens: 2, costUsd: 0.25 }, messages: [{ content: '"\\\n'.repeat(2 * 1024 * 1024) }], model: "after-large-content" }); + for (let offset = 0; offset < line.length; offset += 4093) await log("stdout", line.slice(offset, offset + 4093)); + await log.flush(); + expect(saved).toHaveBeenLastCalledWith(expect.objectContaining({ costUsd: 0.25, complete: true })); + expect(parser.mock.calls[0][0]).toContain("after-large-content"); + expect(parser.mock.calls[0][0].length).toBeLessThan(1000); + expect(output.mock.calls.map(call => call[1]).join("")).toBe(line); + }); + + it("preserves Unicode and escaped quotes across the string limit and chunk boundaries", async () => { + for (const suffix of ['\\u1234', '\\"', '\\\\']) { + const saved = vi.fn(); const parser = vi.fn(createParser()); + const log = createUsageCheckpointLog(vi.fn(), saved, parser); + const line = '{"type":"result","text":"' + 'a'.repeat(249) + suffix + 'discarded","usage":{"costUsd":0.25}}\n'; + for (const character of line) await log("stdout", character); + await log.flush(); + expect(saved).toHaveBeenCalledWith(expect.objectContaining({ costUsd: 0.25 })); + } + }); + it("surfaces a persistence failure outside the local process's best-effort log path", async () => { const failure = new Error("Receipt storage unavailable"); const output = vi.fn(); const saved = vi.fn().mockRejectedValue(failure); diff --git a/packages/adapter-utils/src/usage-checkpoint.ts b/packages/adapter-utils/src/usage-checkpoint.ts index 1e0c44bd86e..123ae66615e 100644 --- a/packages/adapter-utils/src/usage-checkpoint.ts +++ b/packages/adapter-utils/src/usage-checkpoint.ts @@ -24,6 +24,41 @@ function project(value: unknown, depth = 0, field = ""): unknown { : project(entry, depth + 1, key)])); } +/** Bound strings before buffering whole JSON lines. Provider records can carry + * many MiB of conversation text alongside a small usage object. Preserve JSON + * escapes across chunks so discarded text cannot corrupt the accounting fields. */ +function createStringCompactor() { + let inString = false, length = 0, escapeRemaining = 0; + let escapeCode = false, keepEscape = false; + return (chunk: string) => { + let output = ""; + for (const character of chunk) { + if (character === "\n") { + inString = false; escapeRemaining = 0; escapeCode = false; + output += character; + } else if (!inString) { + output += character; + if (character === '"') { inString = true; length = 0; } + } else if (escapeRemaining > 0) { + if (keepEscape) output += character; + escapeRemaining = escapeCode && character === "u" ? 4 : escapeRemaining - 1; + escapeCode = false; + length++; + } else if (character === '"') { + output += character; inString = false; + } else if (character === "\\") { + keepEscape = length < 250; + if (keepEscape) output += character; + escapeRemaining = 1; escapeCode = true; length++; + } else { + if (length < 250) output += character; + length++; + } + } + return output; + }; +} + /** The consumer incrementally parses only new records (never the full history). * A fresh instance belongs to one CLI attempt. Call flush() after its process * exits, before retrying or final result handling. Checkpoint failures are @@ -34,6 +69,7 @@ export function createUsageCheckpointLog( consume: (records: string) => AdapterUsageCheckpoint | null, ) { const attemptId = randomUUID(); + const compactStrings = createStringCompactor(); let remainder = "", accounting = "", previous = ""; let snapshot: AdapterUsageCheckpoint | null = null; let failure: unknown; @@ -63,7 +99,7 @@ export function createUsageCheckpointLog( const log = async (stream: "stdout" | "stderr", chunk: string) => { if (onUsage && !failed && stream === "stdout") { try { - remainder += chunk; + remainder += compactStrings(chunk); const lines = remainder.split("\n"); remainder = lines.pop() ?? ""; if (remainder.length > 8 * 1024 * 1024) throw new Error("Accounting protocol line exceeds 8 MiB"); for (const line of lines) retain(line); diff --git a/server/src/__tests__/adapter-cost-failure-paths.test.ts b/server/src/__tests__/adapter-cost-failure-paths.test.ts index bde2346f785..b831a1e0476 100644 --- a/server/src/__tests__/adapter-cost-failure-paths.test.ts +++ b/server/src/__tests__/adapter-cost-failure-paths.test.ts @@ -98,10 +98,12 @@ describe("CLI adapter accounting on timeout", () => { } const end = fixture.name === "pi" && !timedOut ? '\n{"type":"agent_end","messages":[]}' : ""; const first = JSON.stringify(early) + "\n"; - const noise = JSON.stringify({ type: "text", part: { text: "x".repeat(5 * 1024 * 1024) } }) + "\n"; + const noise = JSON.stringify({ type: "text", part: { text: "x".repeat(9 * 1024 * 1024) } }) + "\n"; const last = JSON.stringify(fixture.event) + end; processResult.mockImplementation(async (_run, _target, _command, _args, options) => { - for (const chunk of [first, noise, last]) await options.onLog("stdout", chunk); + for (const record of [first, noise, last]) { + for (let offset = 0; offset < record.length; offset += 65521) await options.onLog("stdout", record.slice(offset, offset + 65521)); + } return { exitCode: timedOut ? null : 0, signal: timedOut ? "SIGTERM" : null, timedOut, stdout: (first + noise + last).slice(-4 * 1024 * 1024), stderr: "", pid: 123, startedAt: new Date().toISOString() }; }); const onUsage = checkpoint ? vi.fn() : undefined;