From da412b0392e30def79a31768621605cc11c1d33f Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sat, 24 Oct 2026 08:21:41 +0000 Subject: [PATCH] feat(chord,durable): appendText for cheap string appends; Durable partial commits append instead of reassigning Follow-up to #9265. - chord: Delta appendText(draft, key, text) / StringKeys record an append op without the O(length) prefix compare that a plain `draft[key] += text` costs in prepare() - durable: partial commits append text and thinking deltas with appendText (assignJson plan, checked by length), and apply tool-call argument changes from their own StreamingJsonParser of the raw toolcall_delta text instead of reassigning the arguments snapshot - durable "." entry-graph budget 73 -> 76 for the parser it now imports --- packages/chord/CHANGELOG.md | 1 + packages/chord/src/delta/diff.ts | 3 +- packages/chord/src/delta/index.ts | 2 +- packages/chord/src/delta/tracker.ts | 48 +++- packages/chord/test/delta-append-text.test.ts | 190 ++++++++++++++++ packages/durable/CHANGELOG.md | 1 + packages/durable/src/harness/generation.ts | 29 ++- packages/durable/src/harness/json.ts | 33 ++- .../durable/src/harness/partial-changes.ts | 166 ++++++++++++++ packages/durable/test/harness-events.test.ts | 48 +++- .../test/harness-partial-changes.test.ts | 213 ++++++++++++++++++ scripts/check-entry-graphs.mjs | 6 +- 12 files changed, 719 insertions(+), 21 deletions(-) create mode 100644 packages/chord/test/delta-append-text.test.ts create mode 100644 packages/durable/src/harness/partial-changes.ts create mode 100644 packages/durable/test/harness-partial-changes.test.ts diff --git a/packages/chord/CHANGELOG.md b/packages/chord/CHANGELOG.md index aa6634e23..bce91731e 100644 --- a/packages/chord/CHANGELOG.md +++ b/packages/chord/CHANGELOG.md @@ -14,6 +14,7 @@ - `Watchable`: a value that changes and can be watched (`value`, `watch(context)` returning a `WatchHandle`), the common interface of documents, derives and other observable values. - `NativeWatchHandle` and its implementation `NativeWatch`: the same lifecycle for in-process values, with listener `(value, context)` and no ops; an overflow collapses to the newest value. - `ownedSlice(text, start, end?)` in `@earendil-works/chord/delta`: a slice of a string that does not keep the whole string alive ([#9265](https://github.com/earendil-works/pi/issues/9265)) +- `appendText(draft, key, text)` in `@earendil-works/chord/delta`: append to a string in a draft at O(text) cost, recorded as an append op without comparing the old and new strings (a plain `draft[key] += text` makes `prepare()` read the whole string, O(length) per change); `StringKeys` names the keys it accepts ([#9265](https://github.com/earendil-works/pi/issues/9265)) ### Fixed diff --git a/packages/chord/src/delta/diff.ts b/packages/chord/src/delta/diff.ts index ff0bf2121..5fbebe05b 100644 --- a/packages/chord/src/delta/diff.ts +++ b/packages/chord/src/delta/diff.ts @@ -60,7 +60,8 @@ const permutation = (before: readonly JsonValue[], after: readonly JsonValue[]): return result; }; -// O(length): comparing two revisions reads both strings, and slicing a rope makes V8 flatten all of it. +// O(length): comparing two revisions reads both strings, and slicing a rope makes V8 flatten all of it. A writer that +// knows what it appended records it with appendText() instead of diffing. const emitString = (before: string, after: string, path: NonEmptyPath, operations: Op[]): void => { if (before === after) return; if (after.length > before.length && after.slice(0, before.length) === before) { diff --git a/packages/chord/src/delta/index.ts b/packages/chord/src/delta/index.ts index 11d94b47c..c961cf5c9 100644 --- a/packages/chord/src/delta/index.ts +++ b/packages/chord/src/delta/index.ts @@ -123,7 +123,7 @@ export function ownedSlice(text: string, start: number, end = text.length): stri export { diffRevisions } from "./diff.ts"; export type { Draft } from "./draft.ts"; -export { type Change, type Prepared, type Tracker, track } from "./tracker.ts"; +export { appendText, type Change, type Prepared, type StringKeys, type Tracker, track } from "./tracker.ts"; // ─── Path safety ───────────────────────────────────────────────────────────── diff --git a/packages/chord/src/delta/tracker.ts b/packages/chord/src/delta/tracker.ts index bb267bd78..db5eb6b66 100644 --- a/packages/chord/src/delta/tracker.ts +++ b/packages/chord/src/delta/tracker.ts @@ -14,6 +14,11 @@ import { export type { Draft, JsonValue, Op }; +/** Keys of `T` whose value type is exactly `string`: required, and not a string literal type. */ +export type StringKeys = T extends readonly unknown[] + ? never + : { [Key in keyof T]: string extends T[Key] ? (T[Key] extends string ? Key : never) : never }[keyof T]; + type Primitive = null | boolean | number | string; type Container = JsonValue[] | Record; type Stored = JsonValue; @@ -77,6 +82,11 @@ type OverlayNode = { deleteKey?: string; deletes?: Set; readded?: Set; + /** + * Keys whose only changes in this transaction are `appendText()` calls, with the text appended so far. `prepare()` + * records each as one append op without comparing the old and new strings. + */ + appends?: Map; array?: ArrayOverlay; dirty?: boolean; subtreeDirty?: boolean; @@ -330,6 +340,33 @@ export function track(initial: T): Tracker { return new TrackerImpl(initial); } +/** + * Append `text` to the string at `draft[key]`, costing O(text.length). The change is recorded as one append op + * (several appends to the same key in one change join into one op) without comparing the old and new strings. A plain + * assignment such as `draft[key] += text` gives the same op, but `prepare()` then has to check that the new string + * extends the old one, which reads the whole string: O(length) per change. A plain assignment or deletion of the key in + * the same change falls back to that check. + */ +export function appendText(draft: Draft, key: StringKeys, text: string): void { + // Trust boundary: JavaScript callers can pass any object and key; the types cover TypeScript callers. + const node = (draft as Record)[NODE] as OverlayNode | undefined; + if (node === undefined) throw new TypeError("appendText() needs a draft"); + assertWritable(node.context); + const name = String(key); + const current = Array.isArray(nodeBase(node)) || !objectHas(node, name) ? undefined : objectValue(node, name); + if (typeof current !== "string") throw new TypeError(`appendText() needs a string at ${JSON.stringify(name)}`); + if (text.length === 0) return; + // The op and the new value keep the text; a copy, so they never keep a larger string the caller cut it from alive. + const owned = ownedSlice(text, 0); + const appended = node.appends?.get(name) ?? (hasObjectWrite(node, name) ? undefined : ""); + setObjectWrite(node, name, current + owned); + markDirty(node); + if (appended !== undefined) { + node.appends ??= new Map(); + node.appends.set(name, appended + owned); + } +} + function materializePrepared(context: OverlayContext): PreparedImpl { const operations = ensureOperations(context); const base = context.baseValue as T; @@ -588,6 +625,7 @@ function setProperty(node: OverlayNode, property: PropertyKey, supplied: unknown const current = objectValue(node, key); const wasDeleted = isObjectDeleted(node, key); if (!wasDeleted && !isContainer(stored) && current === stored) return true; + node.appends?.delete(key); setObjectWrite(node, key, stored); if (wasDeleted && Object.hasOwn(nodeBase(node), key)) { if (node.readded === undefined) node.readded = new Set(); @@ -604,6 +642,7 @@ function deleteProperty(node: OverlayNode, property: PropertyKey): boolean { if (Array.isArray(nodeBase(node))) throw new TypeError("Overlay arrays cannot contain holes"); const key = String(property); if (!objectHas(node, key)) return true; + node.appends?.delete(key); deleteObjectWrite(node, key); node.readded?.delete(key); setObjectDeletion(node, key); @@ -1846,6 +1885,11 @@ function hasPlacementAncestor(node: OverlayNode): boolean { function emitObjectWrite(node: OverlayNode, path: Path, operations: Op[], key: string, value: Stored): boolean { const nextPath = [...path, key] as unknown as NonEmptyPath; + const appended = node.appends?.get(key); + if (appended !== undefined) { + operations.push(["a", nextPath, appended]); + return false; + } const before = node.readded?.has(key) || !Object.hasOwn(nodeBase(node), key) ? undefined @@ -2054,8 +2098,8 @@ function emitChangedValue( if (isContainer(before) && isContainer(after) && equalTrustedJson(before, after)) return false; if (typeof before === "string" && typeof after === "string") { // O(length): any slice of a rope makes V8 flatten all of it, so this prefix check copies the whole string once per - // change, even when the writer only appended. The op text gets a copy of its own, so it does not keep the whole - // of `after` alive. + // change, even when the writer only appended. Streaming writers use appendText(), which skips it. The op text gets + // a copy of its own, so it does not keep the whole of `after` alive. if (after.length > before.length && after.slice(0, before.length) === before) { operations.push(["a", path, ownedSlice(after, before.length)]); return true; diff --git a/packages/chord/test/delta-append-text.test.ts b/packages/chord/test/delta-append-text.test.ts new file mode 100644 index 000000000..51822c611 --- /dev/null +++ b/packages/chord/test/delta-append-text.test.ts @@ -0,0 +1,190 @@ +import { GCProfiler, getHeapStatistics, setFlagsFromString } from "node:v8"; +import { runInNewContext } from "node:vm"; +import { describe, expect, it } from "vitest"; +import { appendText, applyImmutable, type Draft, type Op, type StringKeys, track } from "../src/delta/index.ts"; + +// #9265: streaming writers appended to growing strings by assignment, and the tracker's prefix check copied the whole +// string on every change. + +type Block = { type: "text"; text: string; signature?: string; tags: string[] }; +type Doc = { title: string; content: Block[]; meta: { note: string } }; + +function doc(): Doc { + return { title: "t", content: [{ type: "text", text: "hello", tags: ["a"] }], meta: { note: "n" } }; +} + +/** Run `mutate` in one change, adopt it, and return its ops; checks that replaying the ops gives the new value. */ +function commit(tracker: ReturnType>, mutate: (draft: Draft) => void): Op[] { + const base = tracker.value; + const change = tracker.beginChange(); + mutate(change.state); + const prepared = change.prepare(); + expect(applyImmutable(base, prepared.ops)).toEqual(prepared.value); + tracker.adopt(prepared); + return [...prepared.ops]; +} + +describe("appendText", () => { + it("records an append op and the appended value", () => { + const tracker = track(doc()); + const before = tracker.value; + expect(commit(tracker, (draft) => appendText(draft.content[0]!, "text", " world"))).toEqual([ + ["a", ["content", 0, "text"], " world"], + ]); + expect(tracker.value.content[0]!.text).toBe("hello world"); + expect(before.content[0]!.text).toBe("hello"); + expect(commit(tracker, (draft) => appendText(draft, "title", "!"))).toEqual([["a", ["title"], "!"]]); + expect(commit(tracker, (draft) => appendText(draft.meta, "note", "ote"))).toEqual([ + ["a", ["meta", "note"], "ote"], + ]); + expect(tracker.value).toEqual({ + title: "t!", + content: [{ type: "text", text: "hello world", tags: ["a"] }], + meta: { note: "note" }, + }); + }); + + it("joins several appends to one key in one change into one op, next to other changes", () => { + const tracker = track(doc()); + const ops = commit(tracker, (draft) => { + appendText(draft.content[0]!, "text", " a"); + draft.title = "new"; + appendText(draft.content[0]!, "text", " b"); + appendText(draft.content[0]!, "text", ""); + appendText(draft.content[0]!, "text", " c"); + }); + expect(ops).toContainEqual(["a", ["content", 0, "text"], " a b c"]); + expect(ops).toContainEqual(["s", ["title"], "new"]); + expect(ops).toHaveLength(2); + expect(tracker.value.content[0]!.text).toBe("hello a b c"); + }); + + it("falls back to the string comparison when the same key is also assigned or deleted", () => { + const tracker = track(doc()); + expect( + commit(tracker, (draft) => { + appendText(draft.content[0]!, "text", " x"); + draft.content[0]!.text += " y"; + }), + ).toEqual([["a", ["content", 0, "text"], " x y"]]); + expect( + commit(tracker, (draft) => { + appendText(draft.content[0]!, "text", " z"); + draft.content[0]!.text = "replaced"; + }), + ).toEqual([["s", ["content", 0, "text"], "replaced"]]); + expect( + commit(tracker, (draft) => { + draft.title = "assigned"; + appendText(draft, "title", " then appended"); + }), + ).toEqual([["s", ["title"], "assigned then appended"]]); + expect( + commit(tracker, (draft) => { + appendText(draft.meta, "note", "x"); + delete (draft.meta as { note?: string }).note; + }), + ).toEqual([["d", ["meta", "note"]]]); + }); + + it("appends inside a block the same change inserted or replaced", () => { + const tracker = track(doc()); + commit(tracker, (draft) => { + draft.content.push({ type: "text", text: "new", tags: [] }); + appendText(draft.content[1]!, "text", " block"); + draft.content[0] = { type: "text", text: "other", tags: [] }; + appendText(draft.content[0]!, "text", " one"); + }); + expect(tracker.value.content.map((block) => block.text)).toEqual(["other one", "new block"]); + }); + + it("rejects what a JavaScript caller can get wrong", () => { + const tracker = track(doc()); + const append = appendText as (draft: unknown, key: string, text: string) => void; + expect(() => append({ text: "plain object" }, "text", "x")).toThrow("appendText() needs a draft"); + const change = tracker.beginChange(); + expect(() => append(change.state, "content", "x")).toThrow('appendText() needs a string at "content"'); + expect(() => append(change.state, "missing", "x")).toThrow('appendText() needs a string at "missing"'); + expect(() => append(change.state.content, "0", "x")).toThrow('appendText() needs a string at "0"'); + const block = change.state.content[0]!; + change.prepare(); + expect(() => appendText(block, "text", "x")).toThrow("Cannot use a settled overlay"); + const aborted = tracker.beginChange(); + const state = aborted.state; + aborted.abort(); + expect(() => appendText(state, "title", "x")).toThrow("Cannot use a settled overlay"); + }); + + it("accepts only keys whose type is exactly string", () => { + type Keys = StringKeys; + const keys: Keys[] = ["text"]; + expect(keys).toEqual(["text"]); + const tracker = track(doc()); + const change = tracker.beginChange(); + const block = change.state.content[0]!; + // @ts-expect-error a string literal type + expect(() => appendText(block, "type", "x")).not.toThrow(); + // @ts-expect-error an optional string + expect(() => appendText(block, "signature", "x")).toThrow(); + // @ts-expect-error an array + expect(() => appendText(block, "tags", "x")).toThrow(); + // @ts-expect-error array keys + expect(() => appendText(change.state.content, 0, "x")).toThrow(); + change.abort(); + }); + + it("keeps an op text that does not hold a larger string it was cut from alive", () => { + setFlagsFromString("--expose-gc"); + const gc = runInNewContext("gc") as () => void; + gc(); + const before = getHeapStatistics().used_heap_size; + const tracker = track(doc()); + const ops: Op[] = []; + for (let index = 0; index < 50; index++) { + const source = `${index}`.repeat(100_000); + ops.push(...commit(tracker, (draft) => appendText(draft.content[0]!, "text", source.slice(0, 40)))); + } + gc(); + expect(getHeapStatistics().used_heap_size - before).toBeLessThan(1_000_000); + expect(ops).toHaveLength(50); + }); +}); + +describe("appendText cost", () => { + /** Bytes `run` allocates, counting what collections during it freed. */ + function allocatedBytes(run: () => void): number { + const profiler = new GCProfiler(); + profiler.start(); + const before = getHeapStatistics().used_heap_size; + run(); + const after = getHeapStatistics().used_heap_size; + let freed = 0; + for (const collection of profiler.stop().statistics) { + freed += collection.beforeGC.heapStatistics.usedHeapSize - collection.afterGC.heapStatistics.usedHeapSize; + } + return after - before + freed; + } + + /** One change per 40-character chunk, appended to one streamed text until it has `length` characters. */ + function stream(length: number): number { + const tracker = track(doc()); + const chunk = "0123456789abcdefghij0123456789abcdefghij"; + return allocatedBytes(() => { + for (let streamed = 0; streamed < length; streamed += chunk.length) { + const change = tracker.beginChange(); + appendText(change.state.content[0]!, "text", chunk); + tracker.adopt(change.prepare()); + } + }); + } + + it("costs O(chunk) per change: a 500 KB stream allocates linearly", () => { + const small = stream(50_000); + const large = stream(500_000); + // Ten times the text: linear growth is about 10x. Measured about 10 KB per change, mostly the change itself. + expect(large / small).toBeLessThan(20); + expect(large / 12_500).toBeLessThan(40_000); + // The same stream through `draft.text += chunk` makes prepare() compare the whole string on every change: on + // Node 24, 6 GB allocated and about 2 s for 500 KB, against 120 MB and under 0.1 s here. Not asserted. + }); +}); diff --git a/packages/durable/CHANGELOG.md b/packages/durable/CHANGELOG.md index 41cf5091a..ea1481081 100644 --- a/packages/durable/CHANGELOG.md +++ b/packages/durable/CHANGELOG.md @@ -22,6 +22,7 @@ - Fixed a conversation watcher (`watchEvents()`, conversation views) holding a copy of a streamed text or tool call per partial commit, through Chord append ops that kept the whole string alive; a 140 KB answer followed by a 140 KB tool call held well over 100 MB ([#9265](https://github.com/earendil-works/pi/issues/9265)) - Fixed a nested call's error summary keeping the call's whole output alive through a slice of it ([#9265](https://github.com/earendil-works/pi/issues/9265)) - Fixed task span error messages and codemode call previews keeping the whole text they were cut from alive ([#9265](https://github.com/earendil-works/pi/issues/9265)) +- Fixed partial commits of a streamed answer copying every growing text, thinking and tool-call argument string in full on each commit (about 20 s of CPU and a 100 MB heap peak for a 500 KB tool call): they now append new text with Chord's `appendText()`, and apply the argument changes of their own parse of the streamed argument text ([#9265](https://github.com/earendil-works/pi/issues/9265)) ## [1.4.0] - 2026-10-21 diff --git a/packages/durable/src/harness/generation.ts b/packages/durable/src/harness/generation.ts index c8d26791c..d617713bd 100644 --- a/packages/durable/src/harness/generation.ts +++ b/packages/durable/src/harness/generation.ts @@ -31,8 +31,8 @@ import type { import { addTools, sectionExtension } from "./agent.ts"; import { compactionPolicy, createCompaction, estimateContext, selectCut } from "./compaction.ts"; import { applyBoundary, type Boundary, prepareBoundary, withdrawQueuedInputs } from "./inbox.ts"; -import { assignJson } from "./json.ts"; import { endRun, LiveDoc, type LiveState, type ToolSlot } from "./live.ts"; +import { PartialChanges } from "./partial-changes.ts"; import { planSystemEntries, renderSections, replaySections } from "./prompt.ts"; import { ensureProviderSessionId } from "./provider.ts"; import { appendToolResult, harnessError, ToolTask, type ToolTaskResult } from "./tool.ts"; @@ -441,7 +441,8 @@ export async function convertPartial(tx: Tx, live: Draft, conversatio /** * Stream one request and return the terminal message. Partials commit as trailing writes at most every * `progress.partialIntervalMs` (default 100 ms) with one commit in flight; `finally` stops the throttle and awaits that - * commit, so no stale partial lands after the outcome. + * commit, so no stale partial lands after the outcome. A commit appends what text and thinking grew by and applies + * the argument changes since the last one (`PartialChanges`), so its cost follows the new text, not the strings' length. */ async function streamResponse( runtime: Runtime, @@ -456,20 +457,27 @@ async function streamResponse( let timer: (() => void) | undefined; let inFlight: Promise | undefined; let stopped = false; + const changes = new PartialChanges(); const flush = (): void => { timer = undefined; const partial = pending; pending = undefined; if (partial === undefined || stopped) return; inFlight = (async () => { - // Copy synchronously: the provider keeps mutating its partial. - const message = copyJson(partial, { omitUndefinedProperties: true }); - await runtime.commit(async (tx) => { - const live = await tx.doc(LiveDoc, runtime.conversationId); - live.generation ??= { attempt }; - assignJson(live.generation as Draft>, "message", message); - return undefined; - }, context); + // Copy synchronously: the provider keeps mutating its partial. The changes since the last commit belong to it. + const message = copyJson(partial, { omitUndefinedProperties: true }) as { content: JsonValue[] }; + const taken = changes.take(message); + try { + await runtime.commit(async (tx) => { + const live = await tx.doc(LiveDoc, runtime.conversationId); + live.generation ??= { attempt }; + changes.commit(live.generation as Draft>, message, taken); + return undefined; + }, context); + } catch (error) { + changes.restore(taken); + throw error; + } })() .catch((error: unknown) => { // Rejections after an abort mark or close are expected; the committed state stays consistent. @@ -486,6 +494,7 @@ async function streamResponse( // A partial without content, such as pi-ai's opening `start` event, shows nothing; a deferred response // never gets past it, so it never leaves a partial. if (event.type === "done" || event.type === "error" || event.partial.content.length === 0) continue; + changes.observe(event); pending = event.partial; if (timer === undefined && inFlight === undefined) timer = runtime.setTimer(flush, interval); } diff --git a/packages/durable/src/harness/json.ts b/packages/durable/src/harness/json.ts index 66cfcd891..c96119fe8 100644 --- a/packages/durable/src/harness/json.ts +++ b/packages/durable/src/harness/json.ts @@ -1,28 +1,53 @@ import type { JsonValue } from "@earendil-works/chord"; +import { appendText, type Draft } from "@earendil-works/chord/delta"; type JsonContainer = Record | JsonValue[]; +/** How `assignJson()` writes one value instead of assigning it leaf by leaf. */ +export type AssignPlan = + /** The string grew by `text` since it was last assigned: append it, when the lengths agree. */ + | { readonly kind: "append"; readonly text: string } + /** Leave the value as it is; the caller writes it. */ + | { readonly kind: "keep" } + /** Plans for some members or items of a container. */ + | { readonly kind: "fields"; readonly fields: ReadonlyMap }; + /** * Assign `value` at `target[key]` leaf by leaf. Chord records a container assignment as one full set and only emits an * append when a string leaf is reassigned with a longer string, so writing the partial whole would store and publish the - * complete message on every flush. + * complete message on every flush. That append costs a compare of the whole string, though; where `plan` says what a + * string member grew by, and the committed string plus that text has the new string's length, the text is appended + * with `appendText()` instead. */ -export function assignJson(target: JsonContainer, key: string | number, value: JsonValue): void { +export function assignJson(target: JsonContainer, key: string | number, value: JsonValue, plan?: AssignPlan): void { + if (plan?.kind === "keep") return; const slots = target as Record; const current = slots[key]; + const fields = plan?.kind === "fields" ? plan.fields : undefined; if (isRecord(current) && isRecord(value)) { for (const name of Object.keys(current)) if (!Object.hasOwn(value, name)) delete current[name]; - for (const [name, child] of Object.entries(value)) assignJson(current, name, child); + for (const [name, child] of Object.entries(value)) assignJson(current, name, child, fields?.get(name)); return; } if (Array.isArray(current) && Array.isArray(value) && current.length <= value.length) { const items = current as JsonValue[]; for (let index = 0; index < value.length; index++) { - if (index < items.length) assignJson(items, index, value[index]!); + if (index < items.length) assignJson(items, index, value[index]!, fields?.get(index)); else items.push(value[index]!); } return; } + if ( + plan?.kind === "append" && + typeof current === "string" && + typeof value === "string" && + !Array.isArray(target) && + current.length + plan.text.length === value.length + ) { + // The cast names what the checks above established: `target[key]` is a string. + appendText(target as Draft>, String(key), plan.text); + return; + } if (current !== value) slots[key] = value; } diff --git a/packages/durable/src/harness/partial-changes.ts b/packages/durable/src/harness/partial-changes.ts new file mode 100644 index 000000000..355a9c93c --- /dev/null +++ b/packages/durable/src/harness/partial-changes.ts @@ -0,0 +1,166 @@ +// What a streamed answer changed in its partial message between two partial commits, so a commit costs the new text. +import type { JsonValue } from "@earendil-works/chord"; +import { appendText, type Draft } from "@earendil-works/chord/delta"; +import type { AssistantMessageEvent } from "@earendil-works/pi-ai"; +import { type JsonChange, StreamingJsonParser } from "@earendil-works/pi-ai/utils/streaming-json"; +import { type AssignPlan, assignJson } from "./json.ts"; + +type JsonContainer = Record | JsonValue[]; + +/** What one partial commit writes besides the partial message itself; from `PartialChanges.take()`. */ +export type PartialCommit = { + /** The plan `assignJson()` writes the message with: text appends, and the arguments it leaves to `changes`. */ + readonly plan: AssignPlan; + /** Each tracked tool call by content index: its argument changes since the last commit, in order, and its snapshot. */ + readonly toolCalls: ReadonlyMap; +}; + +/** What a tracked tool call changed since the last commit, and whether any commit took its changes yet. */ +type ToolCallState = { readonly parser: StreamingJsonParser; changes: JsonChange[]; committed: boolean }; + +/** + * Collects, from a response's events, what changed in its partial message since the last commit, so the commit appends + * new text instead of reassigning growing strings. Reassigning makes Chord compare the old and new string, which costs + * O(length) per commit. + * + * - Text and thinking: the `text_delta` and `thinking_delta` events. The commit appends them where the committed string + * plus the deltas has the new string's length, and compares otherwise. A block's `*_end` event drops them, because + * providers may replace the text when a block ends. + * - Tool-call arguments: a `StreamingJsonParser` per call, fed with the raw argument text of each `toolcall_delta`, + * reports what each piece changed. The commit applies those changes to the committed arguments in order and leaves + * the arguments out of the message assignment, so the committed arguments are exactly the parsed text so far. + * `toolcall_end` ends the tracking; the call's last commit assigns its final arguments. + */ +export class PartialChanges { + #text = new Map(); + readonly #toolCalls = new Map(); + + /** Record what `event` changed. */ + observe(event: AssistantMessageEvent): void { + switch (event.type) { + case "text_delta": + case "thinking_delta": { + const field = event.type === "text_delta" ? "text" : "thinking"; + const pending = this.#text.get(event.contentIndex); + if (pending?.field === field) pending.text += event.delta; + else this.#text.set(event.contentIndex, { field, text: event.delta }); + return; + } + case "toolcall_start": + this.#toolCalls.set(event.contentIndex, { + parser: new StreamingJsonParser(), + changes: [], + committed: false, + }); + return; + case "toolcall_delta": { + const call = this.#toolCalls.get(event.contentIndex); + if (call === undefined) return; + for (const change of call.parser.append(event.delta)) { + // Arguments are an object: any other root shows as no arguments, so the snapshot is committed instead. + if (change.path.length === 0 && !isRecord(change.kind === "set" ? value(change) : null)) { + this.#toolCalls.delete(event.contentIndex); + return; + } + call.changes.push(change); + } + return; + } + case "text_end": + case "thinking_end": + this.#text.delete(event.contentIndex); + return; + case "toolcall_end": + this.#toolCalls.delete(event.contentIndex); + return; + } + } + + /** + * Take what changed since the last call, for a commit of `message`, a copy of the partial. Clears the arguments of + * tracked tool calls in `message`: the commit writes them from the changes. Call `restore()` if the commit fails. + */ + take(message: { readonly content: readonly JsonValue[] }): PartialCommit { + const blocks = new Map(); + for (const [index, { field, text }] of this.#text) { + blocks.set(index, { kind: "fields", fields: new Map([[field, { kind: "append", text }]]) }); + } + this.#text = new Map(); + const toolCalls = new Map(); + for (const [index, call] of this.#toolCalls) { + const block = message.content[index] as Record; + toolCalls.set(index, { changes: call.changes, snapshot: block.arguments! }); + block.arguments = {}; + blocks.set(index, { kind: "fields", fields: new Map([["arguments", { kind: "keep" }]]) }); + call.changes = []; + } + return { + plan: { kind: "fields", fields: new Map([["content", { kind: "fields", fields: blocks }]]) }, + toolCalls, + }; + } + + /** + * Write `message` with what `commit` took into `generation.message`: the message leaf by leaf with the plan, then the + * argument changes. A tracked call whose block is new here gets its arguments from its changes, which then start at + * its first piece; when the committed message lost a call whose changes were already committed (another writer reset + * it), the call falls back to its snapshot from now on. + */ + commit(generation: Draft>, message: JsonValue, commit: PartialCommit): void { + const committedBlocks = (generation.message as { content?: unknown[] } | undefined)?.content?.length ?? 0; + const apply: [number, readonly JsonChange[]][] = []; + for (const [index, { changes, snapshot }] of commit.toolCalls) { + const call = this.#toolCalls.get(index); + if (call === undefined) continue; + if (index >= committedBlocks && call.committed) { + // The earlier changes are not in the committed message: commit the snapshot and stop tracking the call. + this.#toolCalls.delete(index); + (message as { content: Record[] }).content[index]!.arguments = snapshot; + continue; + } + call.committed = true; + apply.push([index, changes]); + } + assignJson(generation, "message", message, commit.plan); + const content = (generation.message as unknown as { content: Draft>[] }).content; + for (const [index, changes] of apply) { + for (const change of changes) applyChange(content[index]!, change); + } + } + + /** Put back what a failed commit took, so the next commit applies it first. */ + restore(commit: PartialCommit): void { + for (const [index, { changes }] of commit.toolCalls) { + const call = this.#toolCalls.get(index); + if (call !== undefined) call.changes = [...changes, ...call.changes]; + } + } +} + +/** Apply one argument change to the committed arguments of `block`. */ +function applyChange(block: Draft>, change: JsonChange): void { + const key = change.path.at(-1); + if (key === undefined) { + // Only an object root is tracked, and only `set` starts it. + if (change.kind === "set") block.arguments = value(change); + return; + } + let parent = block.arguments as Draft; + for (const segment of change.path.slice(0, -1)) { + parent = (parent as Record>)[segment]!; + } + const slots = parent as Record; + if (change.kind === "set") slots[key] = value(change); + // appendText() takes object members; a string in an array is reassigned, which compares it whole. + else if (Array.isArray(parent)) slots[key] = `${slots[key] as string}${change.text}`; + else appendText(parent as Draft>, String(key), change.text); +} + +/** The value a `set` change reports: JSON data, which pi-ai types with read-only arrays. */ +function value(change: Extract): JsonValue { + return change.value as JsonValue; +} + +function isRecord(value: JsonValue | undefined): boolean { + return typeof value === "object" && value !== null && !Array.isArray(value); +} diff --git a/packages/durable/test/harness-events.test.ts b/packages/durable/test/harness-events.test.ts index ce136a674..e12d7d295 100644 --- a/packages/durable/test/harness-events.test.ts +++ b/packages/durable/test/harness-events.test.ts @@ -1,4 +1,4 @@ -import { getHeapStatistics, setFlagsFromString } from "node:v8"; +import { GCProfiler, getHeapStatistics, setFlagsFromString } from "node:v8"; import { runInNewContext } from "node:vm"; import type { Draft } from "@earendil-works/chord"; import { @@ -628,6 +628,21 @@ const gc = (() => { return runInNewContext("gc") as () => void; })(); +/** Bytes `run` allocates, counting what collections during it freed. */ +async function allocatedBytes(run: () => Promise): Promise { + gc(); + const profiler = new GCProfiler(); + profiler.start(); + const before = getHeapStatistics().used_heap_size; + await run(); + const after = getHeapStatistics().used_heap_size; + let freed = 0; + for (const collection of profiler.stop().statistics) { + freed += collection.beforeGC.heapStatistics.usedHeapSize - collection.afterGC.heapStatistics.usedHeapSize; + } + return after - before + freed; +} + /** Heap bytes in use after a full collection. */ async function heapInUse(): Promise { for (let round = 0; round < 3; round++) { @@ -670,6 +685,37 @@ describe("streamed partials", () => { await harness.close(context); }); + it("commits a growing text and tool call at a cost that follows the new text, not the string lengths", async () => { + /** Bytes allocated while streaming one answer of `size` characters of text and as much tool-call content. */ + const stream = async (size: number) => { + const setup = chatSetup({ tokensPerSecond: 64_000, tokenSize: { min: 64, max: 64 } }); + setup.settings.progress = { partialIntervalMs: 1, outputIntervalMs: 100 }; + const text = 'streamed text with "quotes" and \\ escapes\n'.repeat(size / 40); + setup.faux.setResponses([ + fauxAssistantMessage( + [fauxText(text), fauxToolCall("missing", { path: "a.ts", content: text }, { id: "c1" })], + { + stopReason: "toolUse", + }, + ), + fauxAssistantMessage([fauxText("done")]), + ]); + const { harness, root } = await openChat(new MemoryStorage(), setup); + let commits = 0; + harness.subscribeCommits(() => commits++); + const bytes = await allocatedBytes(async () => { + await (await root.submit({ type: "input", content: "go" }, context)).wait(context); + }); + await harness.close(context); + return { bytes, commits }; + }; + const small = await stream(40_000); + const large = await stream(400_000); + expect(large.commits).toBeGreaterThan(small.commits * 5); + // Per commit, the cost stays flat: comparing the whole strings on every commit made it grow with their length. + expect(large.bytes / large.commits / (small.bytes / small.commits)).toBeLessThan(2); + }, 60_000); + it("keeps a watcher's memory proportional to a long streamed text and tool call, not to their commits", async () => { const setup = chatSetup({ tokensPerSecond: 16_000, tokenSize: { min: 16, max: 16 } }); setup.settings.progress = { partialIntervalMs: 1, outputIntervalMs: 100 }; diff --git a/packages/durable/test/harness-partial-changes.test.ts b/packages/durable/test/harness-partial-changes.test.ts new file mode 100644 index 000000000..3b61133b7 --- /dev/null +++ b/packages/durable/test/harness-partial-changes.test.ts @@ -0,0 +1,213 @@ +import type { JsonValue } from "@earendil-works/chord"; +import { type Op, track } from "@earendil-works/chord/delta"; +import { + type AssistantMessage, + type AssistantMessageEvent, + type ToolCall, + ToolCallArgumentsStream, +} from "@earendil-works/pi-ai"; +import { describe, expect, it } from "vitest"; +import { PartialChanges } from "../src/harness/partial-changes.ts"; + +// #9265: partial commits reassigned the growing strings of a streamed answer, and Chord compared each one whole. + +type Live = { generation: Record }; + +/** + * A response streaming one text block (started empty) and then one tool call, as a provider produces it: the partial + * message it mutates and the events it emits. Partial commits go into a tracked document, as Durable's generation does. + */ +function setup() { + const partial: AssistantMessage = { + role: "assistant", + content: [{ type: "text", text: "" }], + api: "test", + provider: "test", + model: "test", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "pending", + timestamp: 0, + }; + const changes = new PartialChanges(); + const tracker = track({ generation: { attempt: 1 } }); + const args = new ToolCallArgumentsStream(); + const emit = (event: AssistantMessageEvent) => changes.observe(event); + const ops: Op[][] = []; + return { + partial, + tracker, + ops, + text(delta: string) { + (partial.content[0] as { text: string }).text += delta; + emit({ type: "text_delta", contentIndex: 0, delta, partial }); + }, + argument(delta: string) { + if (partial.content[1] === undefined) { + partial.content.push({ type: "toolCall", id: "c1", name: "write", arguments: {} }); + emit({ type: "toolcall_start", contentIndex: 1, partial }); + } + (partial.content[1] as { arguments: Record }).arguments = args.append(delta); + emit({ type: "toolcall_delta", contentIndex: 1, delta, partial }); + }, + /** The provider ends both blocks. */ + endBlocks() { + emit({ type: "text_end", contentIndex: 0, content: (partial.content[0] as { text: string }).text, partial }); + emit({ type: "toolcall_end", contentIndex: 1, toolCall: partial.content[1] as ToolCall, partial }); + }, + /** One partial commit, as the generation's throttle makes it; `fail` aborts it like a rejected commit. */ + commit(fail = false) { + const message = structuredClone(partial) as unknown as { content: JsonValue[] }; + const taken = changes.take(message); + const change = tracker.beginChange(); + changes.commit(change.state.generation, message as unknown as JsonValue, taken); + if (fail) { + change.abort(); + changes.restore(taken); + return; + } + const prepared = change.prepare(); + ops.push([...prepared.ops]); + tracker.adopt(prepared); + }, + committed() { + const message = tracker.value.generation.message as { content: { text?: string; arguments?: unknown }[] }; + return { text: message.content[0]?.text, arguments: message.content[1]?.arguments }; + }, + }; +} + +function pieces(text: string, size: number): string[] { + const result: string[] = []; + for (let index = 0; index < text.length; index += size) result.push(text.slice(index, index + size)); + return result; +} + +const ARGUMENTS = { + path: "src/a.ts", + n: 1234, + ok: true, + content: 'line "one"\n\tline \\ two é 😀\n', + edits: [ + { oldText: "a", newText: "b\nc" }, + { oldText: "", newText: "x".repeat(50) }, + ], + nested: { deep: { list: [1, [2, "three"], { four: null }] } }, +}; + +describe("PartialChanges", () => { + it("commits tool-call arguments equal to the final parse, at every piece size and commit rhythm", () => { + const text = JSON.stringify(ARGUMENTS); + for (const size of [1, 2, 3, 5, 7, 16]) { + for (const every of [1, 3, 10]) { + const stream = setup(); + pieces(text, size).forEach((piece, index) => { + stream.argument(piece); + if (index % every === every - 1) stream.commit(); + }); + stream.commit(); + expect(stream.committed().arguments, `size ${size}, every ${every}`).toEqual(ARGUMENTS); + } + } + }); + + it("handles a number ending and a string starting in one piece, and escapes split across pieces", () => { + const stream = setup(); + for (const piece of [ + '{"n":1', + '2,"s":"a', + "b\\", + 'u00e9\\"', + "q\\u", + "00", + "41", + '","t":tr', + "ue,", + '"u":"\\ud83d', + '\\ude00"}', + ]) { + stream.argument(piece); + stream.commit(); + } + expect(stream.committed().arguments).toEqual({ n: 12, s: 'abé"qA', t: true, u: "😀" }); + }); + + it("appends growing argument and text strings instead of reassigning them", () => { + const stream = setup(); + stream.text("Hello"); + stream.argument('{"path":"a.ts","content":"x'); + stream.commit(); + for (let index = 0; index < 20; index++) { + stream.text(" more"); + stream.argument("yz"); + stream.commit(); + } + for (const ops of stream.ops.slice(1)) { + expect(ops).toEqual([ + ["a", ["generation", "message", "content", 0, "text"], " more"], + ["a", ["generation", "message", "content", 1, "arguments", "content"], "yz"], + ]); + } + expect(stream.committed()).toEqual({ + text: `Hello${" more".repeat(20)}`, + arguments: { path: "a.ts", content: `x${"yz".repeat(20)}` }, + }); + }); + + it("applies the changes of a failed commit in the next one", () => { + const stream = setup(); + stream.argument('{"a":"x'); + stream.commit(); + stream.argument('y","b":[1'); + stream.commit(true); + stream.argument(',2],"c":"z"}'); + stream.commit(); + expect(stream.committed().arguments).toEqual({ a: "xy", b: [1, 2], c: "z" }); + }); + + it("re-applies a failed commit's changes even when it would have added the call", () => { + const stream = setup(); + stream.commit(); + stream.argument('{"a":"x'); + stream.commit(true); + stream.argument('y"'); + stream.commit(); + expect(stream.committed().arguments).toEqual({ a: "xy" }); + }); + + it("falls back to the snapshot when the committed message lost a call whose changes were committed", () => { + const stream = setup(); + stream.argument('{"a":"x'); + stream.commit(); + // Another writer resets the generation's partial. + const change = stream.tracker.beginChange(); + change.state.generation = { attempt: 1 }; + stream.tracker.adopt(change.prepare()); + stream.argument('y"'); + stream.commit(); + expect(stream.committed().arguments).toEqual({ a: "xy" }); + stream.argument(',"b":2}'); + stream.commit(); + expect(stream.committed().arguments).toEqual({ a: "xy", b: 2 }); + }); + + it("commits the final arguments after the call ends, and drops text appends when a block ends", () => { + const stream = setup(); + stream.text("draft"); + stream.argument('{"a":"x'); + stream.commit(); + // Providers may replace a block's text and arguments as it ends; the last commit assigns what they ended with. + stream.text(" more"); + (stream.partial.content[0] as { text: string }).text = "replaced"; + stream.partial.content[1] = { type: "toolCall", id: "c1", name: "write", arguments: { a: "final" } }; + stream.endBlocks(); + stream.commit(); + expect(stream.committed()).toEqual({ text: "replaced", arguments: { a: "final" } }); + }); +}); diff --git a/scripts/check-entry-graphs.mjs b/scripts/check-entry-graphs.mjs index f53303d99..5fbb4658b 100644 --- a/scripts/check-entry-graphs.mjs +++ b/scripts/check-entry-graphs.mjs @@ -64,8 +64,10 @@ const BUDGETS = { // memory and SQLite storages share scan order and cursor handling. Values from Chord's root (copyJson, // replicatedState) reach its host runtime through the root barrel. Chord's context reaches pi-telemetry's root // (index, noop, metrics, memory) for its NOOP telemetry default; harness/telemetry.ts is Durable's schema. - // The Harness's request identity reaches pi-ai's utils/identity.ts and its version constant. - maxFiles: 73, + // The Harness's request identity reaches pi-ai's utils/identity.ts and its version constant. Partial commits + // parse streamed tool-call arguments with pi-ai's utils/streaming-json.ts (and its json-parse.ts) to append what + // argument strings grew by (harness/partial-changes.ts). + maxFiles: 76, forbid: [ "packages/ai/src/index.ts", "packages/ai/src/utils/typebox-helpers.ts",