diff --git a/extensions/workflows/artifacts.ts b/extensions/workflows/artifacts.ts index a19fb6cd..0fa06a95 100644 --- a/extensions/workflows/artifacts.ts +++ b/extensions/workflows/artifacts.ts @@ -24,6 +24,7 @@ import { import { encodeCompleteJson, safeStringify, + toSerializable, truncateUtf8, writeFileAtomic, } from "./serialization.ts"; @@ -379,7 +380,7 @@ function boundEntry(entry: TranscriptEntry, maxBytes: number) { /** Keep the initial prompt plus the newest useful context within the artifact cap. */ export function boundedArtifactTranscript( - transcript: TranscriptEntry[], + transcript: readonly TranscriptEntry[], options: { maxBytes?: number; entryMaxBytes?: number } = {}, ) { if (transcript.length === 0) return []; @@ -440,12 +441,149 @@ export interface BoundedTranscriptsArtifact { omitted: { agents: number; entries: number }; } +/** + * One agent's bounded transcript, encoded once for `transcripts.json`. Live + * progress replaces an agent's transcript array instead of mutating it, so + * unchanged agents keep the same array and reuse this between checkpoints. + */ +interface EncodedAgentTranscript { + /** The entries after `toSerializable`, nested one level deep. */ + readonly serialized: string; + /** `toSerializable` nodes; undefined when the shape is not plain data. */ + readonly nodes: number | undefined; +} + +// Omitted agents need only sizing metadata. Never retain their measured JSON: +// live workflows can keep thousands of transcript arrays reachable. +const transcriptMeasurements = new WeakMap< + readonly TranscriptEntry[], + { + readonly length: number; + readonly last: TranscriptEntry | undefined; + readonly entryCount: number; + readonly measuredBytes: number; + } +>(); + +// A process-wide LRU bounds both encoded text and bookkeeping, even across +// concurrent workflows or caller-supplied artifact budgets. Charge UTF-16 +// storage conservatively using twice the UTF-8 byte length. +const ENCODED_TRANSCRIPTS_MAX_BYTES = 4 * 1024 * 1024; +const ENCODED_TRANSCRIPTS_MAX_ENTRIES = 128; +const encodedTranscripts = new WeakMap< + readonly TranscriptEntry[], + WeakRef +>(); +const retainedTranscripts = new Map(); +let encodedTranscriptBytes = 0; + +function forgetEncodedTranscript(transcript: readonly TranscriptEntry[]) { + const cached = encodedTranscripts.get(transcript)?.deref(); + if (cached) { + encodedTranscriptBytes -= retainedTranscripts.get(cached) ?? 0; + retainedTranscripts.delete(cached); + } + encodedTranscripts.delete(transcript); +} + +const nestOneLevel = (json: string) => json.replaceAll("\n", "\n "); + +/** Nodes `toSerializable` would visit, or undefined for non-plain values. */ +function plainNodes(value: unknown, depth = 0): number | undefined { + if (value === null || typeof value !== "object") return 1; + if (depth > 8) return undefined; + const prototype = Object.getPrototypeOf(value); + if ( + !Array.isArray(value) && + prototype !== Object.prototype && + prototype !== null + ) + return undefined; + let nodes = 1; + for (const child of Array.isArray(value) ? value : Object.values(value)) { + const childNodes = plainNodes(child, depth + 1); + if (childNodes === undefined) return undefined; + nodes += childNodes; + } + return nodes; +} + +function measureAgentTranscript(transcript: readonly TranscriptEntry[]) { + const last = transcript.at(-1); + const cached = transcriptMeasurements.get(transcript); + // Guard the identity cache against an in-place append or replacement. + if (cached?.length === transcript.length && cached.last === last) + return cached; + forgetEncodedTranscript(transcript); + const entries = boundedArtifactTranscript(transcript); + const measurement = { + length: transcript.length, + last, + entryCount: entries.length, + measuredBytes: textBytes(nestOneLevel(JSON.stringify(entries, null, 2))), + }; + transcriptMeasurements.set(transcript, measurement); + return measurement; +} + +function encodeAgentTranscript(transcript: readonly TranscriptEntry[]) { + const cached = encodedTranscripts.get(transcript)?.deref(); + const cachedBytes = cached && retainedTranscripts.get(cached); + if (cached && cachedBytes !== undefined) { + retainedTranscripts.delete(cached); + retainedTranscripts.set(cached, cachedBytes); + return cached; + } + const entries = boundedArtifactTranscript(transcript); + const encoded: EncodedAgentTranscript = { + serialized: nestOneLevel( + JSON.stringify( + toSerializable(entries, { maxNodes: Number.POSITIVE_INFINITY }), + null, + 2, + ), + ), + nodes: plainNodes(entries), + }; + const retainedBytes = 2 * textBytes(encoded.serialized); + if (retainedBytes <= ENCODED_TRANSCRIPTS_MAX_BYTES) { + while ( + encodedTranscriptBytes + retainedBytes > ENCODED_TRANSCRIPTS_MAX_BYTES || + retainedTranscripts.size >= ENCODED_TRANSCRIPTS_MAX_ENTRIES + ) { + const oldest = retainedTranscripts.keys().next().value!; + encodedTranscriptBytes -= retainedTranscripts.get(oldest)!; + retainedTranscripts.delete(oldest); + } + encodedTranscripts.set(transcript, new WeakRef(encoded)); + retainedTranscripts.set(encoded, retainedBytes); + encodedTranscriptBytes += retainedBytes; + } + return encoded; +} + +/** `JSON.stringify({ [index]: entries }, null, 2)` built from encoded agents. */ +function assembleTranscripts( + agents: readonly { readonly index: number; readonly json: string }[], +) { + if (agents.length === 0) return "{}"; + // Integer-like keys serialize in ascending numeric order. + const ordered = [...agents].sort((a, b) => a.index - b.index); + return `{\n${ordered.map((agent) => ` "${agent.index}": ${agent.json}`).join(",\n")}\n}`; +} + /** * Assemble `transcripts.json`, keeping whole agents in index order until the * next one would exceed the byte budget. The dropped tail is counted and * reported (`transcriptsOmitted`) rather than silently discarded by the * serializer's node budget, which previously stopped at 20_000 nodes — roughly * 11 full agents — while the file sat at a quarter of its byte cap (issue #558). + * + * Agent measurements are weakly cached; only the kept prefix is encoded, with + * a bounded cache for serialized fragments. A checkpoint + * where one child changed does not re-serialize every finished child. The + * result is byte-identical to serializing the whole object; any shape the + * fragment assembly cannot reproduce exactly takes the whole-object path. */ export function boundedTranscriptsArtifact( agents: readonly { index: number; transcript: TranscriptEntry[] }[], @@ -456,6 +594,65 @@ export function boundedTranscriptsArtifact( options.maxBytes ?? WORKFLOW_TRANSCRIPTS_MAX_BYTES, ); const maxNodes = options.maxNodes ?? WORKFLOW_TRANSCRIPTS_MAX_NODES; + const indexes = new Set(agents.map((agent) => agent.index)); + const integerKeys = + indexes.size === agents.length && + agents.every( + (agent) => + Number.isInteger(agent.index) && + agent.index >= 0 && + agent.index < 2 ** 32 - 1, + ); + if (!integerKeys) return wholeTranscriptsArtifact(agents, maxBytes, maxNodes); + + const measured = agents.map((agent) => ({ + index: agent.index, + transcript: agent.transcript, + ...measureAgentTranscript(agent.transcript), + })); + // Byte size of the pretty object holding the first `count` agents. + let kept = 0; + let bytes = 2; + for (const agent of measured) { + // "{}" grows to "{\n\n}", then by ",\n": 2 bytes either way. + const next = + bytes + 2 + textBytes(` "${agent.index}": `) + agent.measuredBytes; + if (next > maxBytes) break; + bytes = next; + kept++; + } + let omittedEntries = 0; + for (let i = kept; i < measured.length; i++) + omittedEntries += measured[i]!.entryCount; + + const keptAgents = measured.slice(0, kept).map((agent) => ({ + index: agent.index, + ...encodeAgentTranscript(agent.transcript), + })); + const nodes = keptAgents.reduce( + (total, agent) => + total === undefined || agent.nodes === undefined + ? undefined + : total + agent.nodes, + 1, + ); + const content = assembleTranscripts( + keptAgents.map((agent) => ({ index: agent.index, json: agent.serialized })), + ); + if (nodes === undefined || nodes > maxNodes || textBytes(content) > maxBytes) + return wholeTranscriptsArtifact(agents, maxBytes, maxNodes); + return { + content, + omitted: { agents: measured.length - kept, entries: omittedEntries }, + }; +} + +/** The original whole-object assembly, kept for shapes fragments cannot match. */ +function wholeTranscriptsArtifact( + agents: readonly { index: number; transcript: TranscriptEntry[] }[], + maxBytes: number, + maxNodes: number, +): BoundedTranscriptsArtifact { // Bound each agent first; a single agent is capped well under maxBytes by // boundedArtifactTranscript, so at least one agent always fits. const bounded = agents.map((agent) => ({ diff --git a/tests/extensions/workflows/artifact-node-budget.test.ts b/tests/extensions/workflows/artifact-node-budget.test.ts index efb80134..ac7a6a2f 100644 --- a/tests/extensions/workflows/artifact-node-budget.test.ts +++ b/tests/extensions/workflows/artifact-node-budget.test.ts @@ -1,4 +1,5 @@ import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; import { mkdtempSync, readFileSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -291,3 +292,104 @@ test("transcriptsOmitted surfaces through the completion display projection", () ); assert.match(entry.expanded, /omitted from transcripts\.json/); }); + +test("FAIL: rebuilding transcripts.json does not re-encode agents whose transcript is unchanged", () => { + let reads = 0; + const counted = (entries: TranscriptEntry[]) => + entries.map( + (entry) => + new Proxy(entry, { + get(target, property, receiver) { + if (property === "text") reads++; + return Reflect.get(target, property, receiver); + }, + }), + ); + const agents = Array.from({ length: 40 }, (_, index) => + agent(index, counted(transcript(8))), + ); + const first = boundedTranscriptsArtifact(agents); + assert.ok(reads > 0); + + // Live progress replaces one child's array; the others keep theirs. + reads = 0; + agents[0]!.transcript = [ + ...transcript(8), + { role: "assistant", text: "next step" }, + ]; + const second = boundedTranscriptsArtifact(agents); + assert.equal(reads, 0, "finished agents were not re-read"); + assert.notEqual(second.content, first.content); + + // The cached result is byte-identical to encoding fresh copies. + const fresh = boundedTranscriptsArtifact( + agents.map((item) => ({ + index: item.index, + transcript: item.transcript.map((entry) => ({ ...entry })), + })), + ); + assert.equal(second.content, fresh.content); + assert.deepEqual(second.omitted, fresh.omitted); +}); + +test("transcript encoding retains bounded memory while workflow agents remain live", () => { + // Isolate GC and heap accounting from the test runner. Keep source arrays live + // so a WeakMap alone cannot hide retained encodings. The old cache retained + // ~73 MiB for this fixture; allow ample headroom above a few MiB of output. + const result = spawnSync( + process.execPath, + [ + "--experimental-strip-types", + "--expose-gc", + "--input-type=module", + "--eval", + ` + import { boundedTranscriptsArtifact } from ${JSON.stringify( + new URL("../../../extensions/workflows/artifacts.ts", import.meta.url) + .href, + )}; + const agents = Array.from({ length: 256 }, (_, index) => ({ + index, + transcript: Array.from({ length: 202 }, (_, i) => ({ + role: i % 2 ? "toolResult" : "tool", + text: "Read src/module-" + i + ".ts", + name: "Read", + toolCallId: "toolu_" + i.toString(36).padStart(8, "0"), + timestamp: 1700000000000 + i, + startedAt: 1700000000000 + i, + finishedAt: 1700000000120 + i, + durationMs: 120, + })), + })); + global.gc(); + const before = process.memoryUsage().heapUsed; + const first = boundedTranscriptsArtifact(agents); + // Separate workflows/prefixes must share a bound, rather than each + // keeping a full cache forever. Exercise warm reads and eviction too. + for (let start = 0; start < agents.length; start += 32) { + const batch = agents.slice(start, start + 32); + const cold = boundedTranscriptsArtifact(batch); + const warm = boundedTranscriptsArtifact(batch); + if (cold.content !== warm.content) throw new Error("cache changed bytes"); + } + const again = boundedTranscriptsArtifact(agents); + if (again.content !== first.content) throw new Error("eviction changed bytes"); + global.gc(); + console.log(JSON.stringify({ + retained: process.memoryUsage().heapUsed - before, + omitted: again.omitted.agents, + liveAgents: agents.length, + })); + `, + ], + { encoding: "utf8", timeout: 30_000 }, + ); + assert.equal(result.status, 0, result.stderr || String(result.error)); + const memory = JSON.parse(result.stdout); + assert.equal(memory.liveAgents, 256); + assert.ok(memory.omitted > 0, "fixture must exceed the artifact byte budget"); + assert.ok( + memory.retained < 24 * 1024 * 1024, + `encoding retained ${(memory.retained / 1024 / 1024).toFixed(1)} MiB`, + ); +});