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<T> 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
This commit is contained in:
Mario Zechner 2026-10-24 08:21:41 +00:00
parent f715a15c10
commit da412b0392
12 changed files with 719 additions and 21 deletions

View file

@ -14,6 +14,7 @@
- `Watchable<T>`: a value that changes and can be watched (`value`, `watch(context)` returning a `WatchHandle<T>`), the common interface of documents, derives and other observable values.
- `NativeWatchHandle<T>` and its implementation `NativeWatch<T>`: 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<T>` names the keys it accepts ([#9265](https://github.com/earendil-works/pi/issues/9265))
### Fixed

View file

@ -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) {

View file

@ -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 ─────────────────────────────────────────────────────────────

View file

@ -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> = 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<string, JsonValue>;
type Stored = JsonValue;
@ -77,6 +82,11 @@ type OverlayNode = {
deleteKey?: string;
deletes?: Set<string>;
readded?: Set<string>;
/**
* 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<string, string>;
array?: ArrayOverlay;
dirty?: boolean;
subtreeDirty?: boolean;
@ -330,6 +340,33 @@ export function track<T extends object>(initial: T): Tracker<T> {
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<T extends object>(draft: Draft<T>, key: StringKeys<T>, text: string): void {
// Trust boundary: JavaScript callers can pass any object and key; the types cover TypeScript callers.
const node = (draft as Record<PropertyKey, unknown>)[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<T extends object>(context: OverlayContext): PreparedImpl<T> {
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;

View file

@ -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<T extends object>(tracker: ReturnType<typeof track<T>>, mutate: (draft: Draft<T>) => 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<Block>;
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.
});
});

View file

@ -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

View file

@ -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<LiveState>, 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<void> | 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<Record<string, JsonValue>>, "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<Record<string, JsonValue>>, 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);
}

View file

@ -1,28 +1,53 @@
import type { JsonValue } from "@earendil-works/chord";
import { appendText, type Draft } from "@earendil-works/chord/delta";
type JsonContainer = Record<string, JsonValue> | 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<string | number, AssignPlan> };
/**
* 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<string | number, JsonValue>;
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<Record<string, string>>, String(key), plan.text);
return;
}
if (current !== value) slots[key] = value;
}

View file

@ -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<string, JsonValue> | 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<number, { readonly changes: readonly JsonChange[]; readonly snapshot: JsonValue }>;
};
/** 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<number, { readonly field: "text" | "thinking"; text: string }>();
readonly #toolCalls = new Map<number, ToolCallState>();
/** 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<string | number, AssignPlan>();
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<number, { readonly changes: readonly JsonChange[]; readonly snapshot: JsonValue }>();
for (const [index, call] of this.#toolCalls) {
const block = message.content[index] as Record<string, JsonValue>;
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<Record<string, JsonValue>>, 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<string, JsonValue>[] }).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<Record<string, JsonValue>>[] }).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<Record<string, JsonValue>>, 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<JsonContainer>;
for (const segment of change.path.slice(0, -1)) {
parent = (parent as Record<string | number, Draft<JsonContainer>>)[segment]!;
}
const slots = parent as Record<string | number, JsonValue>;
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<Record<string, string>>, String(key), change.text);
}
/** The value a `set` change reports: JSON data, which pi-ai types with read-only arrays. */
function value(change: Extract<JsonChange, { kind: "set" }>): JsonValue {
return change.value as JsonValue;
}
function isRecord(value: JsonValue | undefined): boolean {
return typeof value === "object" && value !== null && !Array.isArray(value);
}

View file

@ -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<void>): Promise<number> {
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<number> {
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 };

View file

@ -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<string, JsonValue> };
/**
* 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<Live>({ 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<string, unknown> }).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" } });
});
});

View file

@ -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",