diff --git a/packages/coding-agent/src/collab/replication-shrink.ts b/packages/coding-agent/src/collab/replication-shrink.ts index 311aa5899..a401701aa 100644 --- a/packages/coding-agent/src/collab/replication-shrink.ts +++ b/packages/coding-agent/src/collab/replication-shrink.ts @@ -4,19 +4,22 @@ * The host wraps every {@link CollabFrame} in an AES-GCM envelope and ships it * through the relay's WebSocket. WebSocket servers enforce a per-frame * `maxPayloadLength` (Bun's default is 16 MB; many proxies cap lower). A - * single oversized SessionEntry — typically a `read`/`bash`/`search` tool - * result that captured a multi-megabyte blob — would otherwise ship as its own - * oversized chunk and trip that limit, killing the host's WebSocket with - * `1006 Received too big message`. `CollabSocket` treats 1006 as transient and - * reconnects, the next guest hello triggers the same oversized send, and the - * loop never breaks (issue #3739). + * single oversized payload — typically a `read`/`bash`/`search` tool result + * captured as one multi-megabyte string, or a tool result whose `content` + * array holds thousands of small blocks — would otherwise ship as its own + * oversized frame and trip that limit, killing the host's WebSocket with + * `1006 Received too big message`. `CollabSocket` treats 1006 as transient + * and reconnects, the next guest hello triggers the same oversized send, and + * the loop never breaks (issue #3739). * * This helper bounds any JSON-serializable payload below * {@link MAX_REPLICATED_PAYLOAD_BYTES}. Already-small payloads pass through * untouched; oversized ones are returned as a deep-cloned shadow where long - * strings are head-truncated with an `[…N bytes elided for collab session]` - * marker. Guests still see the structural mirror; tool outputs degrade - * gracefully instead of looping the whole session. + * strings are head-truncated AND long arrays are head-clipped, with + * `[…N chars elided for collab session]` / `[…N items elided for collab + * session]` markers. Both axes are needed: string truncation alone leaves + * the cap unenforced for a payload built of many short strings, where no + * field exceeds the per-string floor. */ /** @@ -28,36 +31,62 @@ export const MAX_REPLICATED_PAYLOAD_BYTES = 1 * 1024 * 1024; /** - * Starting per-string head-truncation cap. The shrinker halves this until the - * shrunk payload fits {@link MAX_REPLICATED_PAYLOAD_BYTES}; the floor - * (`MIN_STRING_CAP_BYTES`) bounds the worst case so a payload with very many - * long strings still converges in a handful of passes. + * Progressive shrink passes. Each pass tightens both the per-string cap and + * the per-array head limit; the loop stops at the first pass whose output + * fits {@link MAX_REPLICATED_PAYLOAD_BYTES}. The schedule is concrete (not + * recomputed) so the failure modes the helper guards against are visible: + * + * - One giant string → the first pass already truncates it under 64 KB. + * - Array of many small blocks (e.g. a tool result with thousands of + * `{type:"text", text:"..."}` content items) → later passes head-clip the + * array to a small sample with a `[…N items elided]` summary element. + * + * The final pass clamps every string to 64 B and every array to one element, + * so even pathological mixes converge. */ -const INITIAL_STRING_CAP_BYTES = 64 * 1024; +interface ShrinkPass { + stringCap: number; + arrayLimit: number; +} -const MIN_STRING_CAP_BYTES = 256; +const SHRINK_PASSES: readonly ShrinkPass[] = [ + { stringCap: 64 * 1024, arrayLimit: 256 }, + { stringCap: 16 * 1024, arrayLimit: 128 }, + { stringCap: 4 * 1024, arrayLimit: 64 }, + { stringCap: 1 * 1024, arrayLimit: 32 }, + { stringCap: 256, arrayLimit: 16 }, + { stringCap: 256, arrayLimit: 4 }, + { stringCap: 64, arrayLimit: 1 }, +]; + +const STRING_ELISION_RESERVE = 80; /** - * Recursively walk `value`, head-truncating any string longer than `cap`. - * Returns a deep-cloned copy when truncation occurs and the original - * reference for already-small subtrees, so unchanged subgraphs avoid the - * structural clone cost. + * Recursively walk `value`, head-truncating any string longer than + * `stringCap` and head-clipping any array longer than `arrayLimit`. Returns + * a freshly built deep clone — every object/array is rebuilt so the + * recursive output can be safely serialized in isolation. Cheap to call on + * small values: short strings, numbers, and booleans pass through without + * allocation. */ -function truncateStrings(value: unknown, cap: number): unknown { +function shrinkWalk(value: unknown, stringCap: number, arrayLimit: number): unknown { if (typeof value === "string") { - if (value.length <= cap) return value; - const headLen = Math.max(0, cap - 80); + if (value.length <= stringCap) return value; + const headLen = Math.max(0, stringCap - STRING_ELISION_RESERVE); return `${value.slice(0, headLen)}\n…[${value.length - headLen} chars elided for collab session]`; } if (Array.isArray(value)) { - const out: unknown[] = new Array(value.length); - for (let i = 0; i < value.length; i++) out[i] = truncateStrings(value[i], cap); + const keep = Math.min(value.length, arrayLimit); + const elided = value.length - keep; + const out: unknown[] = new Array(elided > 0 ? keep + 1 : keep); + for (let i = 0; i < keep; i++) out[i] = shrinkWalk(value[i], stringCap, arrayLimit); + if (elided > 0) out[keep] = `…[${elided} items elided for collab session]`; return out; } if (value && typeof value === "object") { const src = value as Record; const out: Record = {}; - for (const k in src) out[k] = truncateStrings(src[k], cap); + for (const k in src) out[k] = shrinkWalk(src[k], stringCap, arrayLimit); return out; } return value; @@ -65,19 +94,18 @@ function truncateStrings(value: unknown, cap: number): unknown { /** * Return `value` unchanged when its JSON serialization already fits - * {@link MAX_REPLICATED_PAYLOAD_BYTES}; otherwise return a deep-cloned shadow - * with long strings head-truncated until the payload fits. The cap is halved - * across passes so a value with one giant string and one with many medium - * strings both converge. + * {@link MAX_REPLICATED_PAYLOAD_BYTES}; otherwise return a deep-cloned + * shadow shrunk along both string and array axes until the payload fits. + * The function is generic over `T` because the wire shape is preserved: + * only string leaves and array tails change; discriminator fields, ids, and + * other small metadata pass through untouched. */ export function shrinkForReplication(value: T): T { if (JSON.stringify(value).length <= MAX_REPLICATED_PAYLOAD_BYTES) return value; - let cap = INITIAL_STRING_CAP_BYTES; - let shrunk = value as unknown; - while (cap >= MIN_STRING_CAP_BYTES) { - shrunk = truncateStrings(value, cap); + let shrunk: unknown = value; + for (const pass of SHRINK_PASSES) { + shrunk = shrinkWalk(value, pass.stringCap, pass.arrayLimit); if (JSON.stringify(shrunk).length <= MAX_REPLICATED_PAYLOAD_BYTES) return shrunk as T; - cap = Math.floor(cap / 2); } return shrunk as T; } diff --git a/packages/coding-agent/test/collab/replication-shrink.test.ts b/packages/coding-agent/test/collab/replication-shrink.test.ts index 101ecc1ba..9087cbfa7 100644 --- a/packages/coding-agent/test/collab/replication-shrink.test.ts +++ b/packages/coding-agent/test/collab/replication-shrink.test.ts @@ -30,6 +30,10 @@ import { unpackEnvelope, } from "@oh-my-pi/pi-coding-agent/collab/protocol"; import { CollabSocket } from "@oh-my-pi/pi-coding-agent/collab/relay-client"; +import { + MAX_REPLICATED_PAYLOAD_BYTES, + shrinkForReplication, +} from "@oh-my-pi/pi-coding-agent/collab/replication-shrink"; import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; import type { SessionEntry } from "@oh-my-pi/pi-coding-agent/session/session-entries"; @@ -271,3 +275,60 @@ describe("collab replication shrinking (#3739)", () => { expect(shrunkContent).toContain("chars elided for collab session"); }); }); + +describe("shrinkForReplication (#3740 review)", () => { + it("passes already-small values through by reference", () => { + const small = { type: "message", id: "x", text: "hi" }; + expect(shrinkForReplication(small)).toBe(small); + }); + + it("clamps a single giant string under the cap with an elision marker", () => { + const giant = { content: "x".repeat(5 * 1024 * 1024) }; + const shrunk = shrinkForReplication(giant); + const size = JSON.stringify(shrunk).length; + expect(size).toBeLessThanOrEqual(MAX_REPLICATED_PAYLOAD_BYTES); + expect(shrunk).not.toBe(giant); + expect(shrunk.content).toContain("chars elided for collab session"); + }); + + it("clamps a payload built of many short strings (no individual oversized) under the cap", () => { + // Realistic shape: a tool result content array with thousands of small + // text blocks. ~3 MB total; no individual string crosses the 64 B + // floor of the final shrink pass, so the helper MUST clip the array, + // not just the strings. + const content = Array.from({ length: 100_000 }, (_, i) => ({ + type: "text", + text: `block-${i}`, + })); + const payload = { role: "toolResult", content }; + const original = JSON.stringify(payload).length; + expect(original).toBeGreaterThan(MAX_REPLICATED_PAYLOAD_BYTES); + const shrunk = shrinkForReplication(payload); + const shrunkSize = JSON.stringify(shrunk).length; + expect(shrunkSize).toBeLessThanOrEqual(MAX_REPLICATED_PAYLOAD_BYTES); + // Array was clipped with a summary marker reporting the dropped count. + const shrunkContent = shrunk.content; + expect(Array.isArray(shrunkContent)).toBe(true); + const marker = shrunkContent.find( + (item: unknown) => typeof item === "string" && item.includes("items elided for collab session"), + ); + expect(marker).toBeDefined(); + }); + + it("preserves the wire discriminator on a fully-shrunk payload", () => { + // Even the worst-case final pass keeps the discriminator key/value + // pairs intact (only string leaves and array tails are touched), so + // guests still see the entry/event as the correct kind. + const evt = { + type: "tool_execution_end", + toolCallId: "call-1", + toolName: "read", + result: { text: "x".repeat(8 * 1024 * 1024) }, + }; + const shrunk = shrinkForReplication(evt); + expect(shrunk.type).toBe("tool_execution_end"); + expect(shrunk.toolCallId).toBe("call-1"); + expect(shrunk.toolName).toBe("read"); + expect(JSON.stringify(shrunk).length).toBeLessThanOrEqual(MAX_REPLICATED_PAYLOAD_BYTES); + }); +});