assistant-stream 0.3.25 → 0.3.27

package/src/core/gorp/changeTree.test.ts

@@ -0,0 +1,168 @@
+import { describe, expect, it } from "vitest";
+import {
+  diffKeys,
+  lookupChange,
+  markChanged,
+  mergeChanged,
+  type ChangeNode,
+} from "./changeTree";
+
+describe("markChanged", () => {
+  it("collapses to true when path is empty", () => {
+    expect(markChanged({}, [])).toBe(true);
+  });
+
+  it("keeps `true` after attempting to descend into it", () => {
+    const node: ChangeNode = { a: true };
+    const result = markChanged(node, ["a", "b"]);
+    expect(result).toBe(node);
+    expect(node.a).toBe(true);
+  });
+
+  it("overwrites a nested subtree with `true` when the path collapses to the parent", () => {
+    const node: ChangeNode = { a: { b: true } };
+    markChanged(node, ["a"]);
+    expect(node.a).toBe(true);
+  });
+
+  it("creates intermediate nodes lazily without disturbing siblings", () => {
+    const node: ChangeNode = { a: { existing: true } };
+    markChanged(node, ["a", "b", "c"]);
+    expect(node).toEqual({
+      a: { existing: true, b: { c: true } },
+    });
+  });
+
+  it("is idempotent when the same leaf is marked twice", () => {
+    const node: ChangeNode = {};
+    markChanged(node, ["a", "b"]);
+    markChanged(node, ["a", "b"]);
+    expect(node).toEqual({ a: { b: true } });
+  });
+
+  it("rejects unsafe path segments", () => {
+    const node: ChangeNode = {};
+    expect(() => markChanged(node, ["__proto__", "polluted"])).toThrow(
+      /Unsafe gorp path segment/,
+    );
+    expect(({} as Record<string, unknown>).polluted).toBeUndefined();
+  });
+});
+
+describe("mergeChanged", () => {
+  it("returns true when either side is true", () => {
+    expect(mergeChanged(true, {})).toBe(true);
+    expect(mergeChanged({}, true)).toBe(true);
+  });
+
+  it("merges disjoint subtrees in place", () => {
+    const target: ChangeNode = { a: true };
+    const source: ChangeNode = { b: true };
+    const result = mergeChanged(target, source);
+    expect(result).toBe(target);
+    expect(target).toEqual({ a: true, b: true });
+  });
+
+  it("clones source subtrees before assigning them", () => {
+    const target: ChangeNode = {};
+    const source: ChangeNode = { a: { b: true } };
+    mergeChanged(target, source);
+    (source as { a: { c?: true } }).a.c = true;
+    expect(target).toEqual({ a: { b: true } });
+  });
+
+  it("deep-merges overlapping subtrees", () => {
+    const target: ChangeNode = { a: { x: true } };
+    const source: ChangeNode = { a: { y: true } };
+    mergeChanged(target, source);
+    expect(target).toEqual({ a: { x: true, y: true } });
+  });
+
+  it("propagates `true` upward when one side has it at a shared key", () => {
+    const target: ChangeNode = { a: { x: true } };
+    const source: ChangeNode = { a: true };
+    mergeChanged(target, source);
+    expect(target).toEqual({ a: true });
+  });
+
+  it("rejects unsafe source keys", () => {
+    const target: ChangeNode = {};
+    const source = JSON.parse('{"__proto__":true}') as ChangeNode;
+    expect(() => mergeChanged(target, source)).toThrow(
+      /Unsafe gorp path segment/,
+    );
+    expect(({} as Record<string, unknown>).polluted).toBeUndefined();
+  });
+
+  it("rejects nested unsafe source keys before assigning a subtree", () => {
+    const target: ChangeNode = {};
+    const source = JSON.parse('{"a":{"constructor":true}}') as ChangeNode;
+    expect(() => mergeChanged(target, source)).toThrow(
+      /Unsafe gorp path segment/,
+    );
+  });
+});
+
+describe("lookupChange", () => {
+  it("returns false for paths that don't exist", () => {
+    expect(lookupChange({ a: true }, ["b"])).toBe(false);
+  });
+
+  it("returns true when traversal hits a `true` sentinel", () => {
+    expect(lookupChange({ a: true }, ["a", "b", "c"])).toBe(true);
+  });
+
+  it("returns the subtree when the path lands on an interior node", () => {
+    const node: ChangeNode = { a: { b: true } };
+    const result = lookupChange(node, ["a"]);
+    expect(result).toEqual({ b: true });
+  });
+
+  it("returns the root node for the empty path", () => {
+    const node: ChangeNode = { a: true };
+    expect(lookupChange(node, [])).toBe(node);
+  });
+
+  it("returns false for inherited object members", () => {
+    expect(lookupChange({}, ["toString"])).toBe(false);
+    expect(lookupChange({ a: {} }, ["a", "hasOwnProperty"])).toBe(false);
+    expect(lookupChange({ a: true }, ["constructor"])).toBe(false);
+  });
+});
+
+describe("inherited object members", () => {
+  it("markChanged does not descend into inherited members", () => {
+    const node: ChangeNode = {};
+    markChanged(node, ["toString", "x"]);
+    expect(node).toEqual({ toString: { x: true } });
+    expect(lookupChange(node, ["toString", "x"])).toBe(true);
+    expect(lookupChange(node, ["valueOf"])).toBe(false);
+  });
+
+  it("mergeChanged ignores inherited target members", () => {
+    const target: ChangeNode = {};
+    const source: ChangeNode = { toString: { x: true } };
+    mergeChanged(target, source);
+    expect(target).toEqual({ toString: { x: true } });
+  });
+});
+
+describe("diffKeys", () => {
+  it("returns new keys in order, then deleted keys in reverse", () => {
+    expect(diffKeys({ b: 1, c: 2 }, { a: 1, b: 2, d: 3 })).toEqual([
+      "b",
+      "c",
+      "d",
+      "a",
+    ]);
+  });
+
+  it("returns new keys when the old value is not an object", () => {
+    expect(diffKeys({ a: 1 }, null)).toEqual(["a"]);
+    expect(diffKeys({ a: 1 }, "str")).toEqual(["a"]);
+  });
+
+  it("returns deleted keys when the new value is not an object", () => {
+    expect(diffKeys(null, { a: 1, b: 2 })).toEqual(["b", "a"]);
+  });
+});

package/src/core/gorp/changeTree.ts

@@ -0,0 +1,96 @@
+export type ChangeNode = true | { [key: string]: ChangeNode };
+
+export function createChangeNode(): { [key: string]: ChangeNode } {
+  return Object.create(null);
+}
+
+const UNSAFE_PATH_SEGMENTS = new Set(["__proto__", "constructor", "prototype"]);
+
+export function assertSafePathSegment(segment: string): void {
+  if (UNSAFE_PATH_SEGMENTS.has(segment)) {
+    throw new Error(`Unsafe gorp path segment: ${segment}`);
+  }
+}
+
+export function markChanged(
+  node: ChangeNode,
+  path: readonly string[],
+): ChangeNode {
+  if (node === true || path.length === 0) return true;
+  let cursor = node;
+  for (let i = 0; i < path.length - 1; i++) {
+    const key = path[i]!;
+    assertSafePathSegment(key);
+    const next = Object.hasOwn(cursor, key) ? cursor[key] : undefined;
+    if (next === true) return node;
+    if (next === undefined) cursor[key] = createChangeNode();
+    cursor = cursor[key] as { [k: string]: ChangeNode };
+  }
+  const leaf = path[path.length - 1]!;
+  assertSafePathSegment(leaf);
+  cursor[leaf] = true;
+  return node;
+}
+
+export function mergeChanged(
+  target: ChangeNode,
+  source: ChangeNode,
+): ChangeNode {
+  if (target === true || source === true) return true;
+  for (const key of Object.keys(source)) {
+    assertSafePathSegment(key);
+    const sChild = source[key]!;
+    const tChild = Object.hasOwn(target, key) ? target[key] : undefined;
+    target[key] =
+      tChild === undefined
+        ? cloneChangeNode(sChild)
+        : mergeChanged(tChild, sChild);
+  }
+  return target;
+}
+
+function cloneChangeNode(source: ChangeNode): ChangeNode {
+  if (source === true) return true;
+  const clone = createChangeNode();
+  for (const key of Object.keys(source)) {
+    assertSafePathSegment(key);
+    clone[key] = cloneChangeNode(source[key]!);
+  }
+  return clone;
+}
+
+export function lookupChange(
+  node: ChangeNode,
+  path: readonly string[],
+): ChangeNode | false {
+  for (const key of path) {
+    if (node === true) return true;
+    const next = Object.hasOwn(node, key) ? node[key] : undefined;
+    if (next === undefined) return false;
+    node = next;
+  }
+  return node;
+}
+
+export function lookupValue(state: unknown, path: readonly string[]): unknown {
+  let node = state;
+  for (const key of path) {
+    assertSafePathSegment(key);
+    if (typeof node !== "object" || node === null) return undefined;
+    node = Object.hasOwn(node, key)
+      ? (node as Record<string, unknown>)[key]
+      : undefined;
+  }
+  return node;
+}
+
+export function diffKeys(newNode: unknown, oldNode: unknown): string[] {
+  const newKeys =
+    typeof newNode === "object" && newNode !== null ? Object.keys(newNode) : [];
+  const oldKeys =
+    typeof oldNode === "object" && oldNode !== null ? Object.keys(oldNode) : [];
+  if (oldKeys.length === 0) return newKeys;
+  const newSet = new Set(newKeys);
+  const deletedKeys = oldKeys.filter((k) => !newSet.has(k)).reverse();
+  return [...newKeys, ...deletedKeys];
+}

package/src/core/{object/createObjectStream.ts → gorp/createGorpStream.ts}

@@ -1,32 +1,35 @@
 import type { ReadonlyJSONValue } from "../../utils";
 import { withPromiseOrValue } from "../utils/withPromiseOrValue";
-import { ObjectStreamAccumulator } from "./ObjectStreamAccumulator";
-import type { ObjectStreamOperation, ObjectStreamChunk } from "./types";
+import { GorpStreamAccumulator } from "./GorpStreamAccumulator";
+import type { GorpStreamOperation, GorpStreamChunk } from "./types";
 
-type ObjectStreamController = {
+type GorpStreamController = {
   readonly abortSignal: AbortSignal;
 
-  enqueue(operations: readonly ObjectStreamOperation[]): void;
+  enqueue(operations: readonly GorpStreamOperation[]): void;
 };
 
-class ObjectStreamControllerImpl implements ObjectStreamController {
-  private _controller: ReadableStreamDefaultController<ObjectStreamChunk>;
+class GorpStreamControllerImpl implements GorpStreamController {
+  private _controller: ReadableStreamDefaultController<GorpStreamChunk>;
   private _abortController = new AbortController();
-  private _accumulator: ObjectStreamAccumulator;
+  private _accumulator: GorpStreamAccumulator;
+  private _cancelled = false;
 
   get abortSignal() {
     return this._abortController.signal;
   }
 
   constructor(
-    controller: ReadableStreamDefaultController<ObjectStreamChunk>,
+    controller: ReadableStreamDefaultController<GorpStreamChunk>,
     defaultValue: ReadonlyJSONValue,
   ) {
     this._controller = controller;
-    this._accumulator = new ObjectStreamAccumulator(defaultValue);
+    this._accumulator = new GorpStreamAccumulator(defaultValue);
   }
 
-  enqueue(operations: readonly ObjectStreamOperation[]) {
+  enqueue(operations: readonly GorpStreamOperation[]) {
+    if (this._cancelled) return;
+
     this._accumulator.append(operations);
 
     this._controller.enqueue({
@@ -36,23 +39,26 @@ class ObjectStreamControllerImpl implements ObjectStreamController {
   }
 
   __internalError(error: unknown) {
+    if (this._cancelled) return;
     this._controller.error(error);
   }
 
   __internalClose() {
+    if (this._cancelled) return;
     this._controller.close();
   }
 
   __internalCancel(reason?: unknown) {
+    this._cancelled = true;
     this._abortController.abort(reason);
   }
 }
 
 const getStreamControllerPair = (defaultValue: ReadonlyJSONValue) => {
-  let controller!: ObjectStreamControllerImpl;
-  const stream = new ReadableStream<ObjectStreamChunk>({
+  let controller!: GorpStreamControllerImpl;
+  const stream = new ReadableStream<GorpStreamChunk>({
     start(c) {
-      controller = new ObjectStreamControllerImpl(c, defaultValue);
+      controller = new GorpStreamControllerImpl(c, defaultValue);
     },
     cancel(reason: unknown) {
       controller.__internalCancel(reason);
@@ -62,15 +68,15 @@ const getStreamControllerPair = (defaultValue: ReadonlyJSONValue) => {
   return [stream, controller] as const;
 };
 
-type CreateObjectStreamOptions = {
-  execute: (controller: ObjectStreamController) => void | PromiseLike<void>;
+type CreateGorpStreamOptions = {
+  execute: (controller: GorpStreamController) => void | PromiseLike<void>;
   defaultValue?: ReadonlyJSONValue;
 };
 
-export const createObjectStream = ({
+export const createGorpStream = ({
   execute,
   defaultValue = {},
-}: CreateObjectStreamOptions) => {
+}: CreateGorpStreamOptions) => {
   const [stream, controller] = getStreamControllerPair(defaultValue);
 
   withPromiseOrValue(

package/src/core/{object → gorp}/types.ts

@@ -1,6 +1,6 @@
 import type { ReadonlyJSONValue } from "../../utils";
 
-export type ObjectStreamOperation =
+export type AssistantTransportStateOperation =
   | {
       readonly type: "set";
       readonly path: readonly string[];
@@ -12,7 +12,9 @@ export type ObjectStreamOperation =
       readonly value: string;
     };
 
-export type ObjectStreamChunk = {
+export type GorpStreamOperation = AssistantTransportStateOperation;
+
+export type GorpStreamChunk = {
   readonly snapshot: ReadonlyJSONValue;
-  readonly operations: readonly ObjectStreamOperation[];
+  readonly operations: readonly GorpStreamOperation[];
 };

package/src/core/modules/assistant-stream.test.ts

@@ -1,6 +1,10 @@
-import { describe, it, expect } from "vitest";
-import { createAssistantStreamResponse } from "./assistant-stream";
+import { describe, it, expect, vi } from "vitest";
+import {
+  createAssistantStream,
+  createAssistantStreamResponse,
+} from "./assistant-stream";
 import { AssistantStream } from "../AssistantStream";
+import type { AssistantStreamChunk } from "../AssistantStreamChunk";
 import { DataStreamDecoder } from "../serialization/data-stream/DataStream";
 import { AssistantMessageAccumulator } from "../accumulators/assistant-message-accumulator";
 import type { AssistantMessage } from "../utils/types";
@@ -21,6 +25,154 @@ const accumulate = async (response: Response): Promise<AssistantMessage> => {
   return last!;
 };
 
+const collectChunks = async (
+  stream: AssistantStream,
+): Promise<AssistantStreamChunk[]> => {
+  const chunks: AssistantStreamChunk[] = [];
+  await stream.pipeTo(
+    new WritableStream({
+      write(chunk) {
+        chunks.push(chunk);
+      },
+    }),
+  );
+  return chunks;
+};
+
+const captureUnhandledRejections = async (
+  callback: () => Promise<void>,
+): Promise<unknown[]> => {
+  const reasons: unknown[] = [];
+  const listener = (reason: unknown) => reasons.push(reason);
+  process.on("unhandledRejection", listener);
+  try {
+    await callback();
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    return reasons;
+  } finally {
+    process.off("unhandledRejection", listener);
+  }
+};
+
+describe("createAssistantStream task settlement", () => {
+  it("emits callback failures without leaking an unhandled rejection", async () => {
+    let chunks: AssistantStreamChunk[] = [];
+    const unhandledRejections = await captureUnhandledRejections(async () => {
+      chunks = await collectChunks(
+        createAssistantStream(async () => {
+          throw new Error("provider failed");
+        }),
+      );
+    });
+
+    expect(chunks).toEqual([
+      {
+        type: "error",
+        path: [],
+        error: "Error: provider failed",
+      },
+    ]);
+    expect(unhandledRejections).toEqual([]);
+  });
+
+  it("does not settle the stream again after cancellation", async () => {
+    let finishCallback!: () => void;
+    const callbackPending = new Promise<void>((resolve) => {
+      finishCallback = resolve;
+    });
+
+    const unhandledRejections = await captureUnhandledRejections(async () => {
+      const reader = createAssistantStream(() => callbackPending).getReader();
+      await reader.cancel("consumer stopped");
+      finishCallback();
+      await callbackPending;
+    });
+
+    expect(unhandledRejections).toEqual([]);
+  });
+
+  it("reports callback failures after the controller is explicitly closed", async () => {
+    const error = new Error("cleanup failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+
+    try {
+      const chunks = await collectChunks(
+        createAssistantStream((controller) => {
+          controller.close();
+          throw error;
+        }),
+      );
+
+      expect(chunks).toEqual([]);
+      expect(consoleError).toHaveBeenCalledOnce();
+      expect(consoleError).toHaveBeenCalledWith(error);
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
+  it("does not report callback failures caused after cancellation", async () => {
+    let failCallback!: (error: Error) => void;
+    const callbackPending = new Promise<void>((_, reject) => {
+      failCallback = reject;
+    });
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+
+    try {
+      const unhandledRejections = await captureUnhandledRejections(async () => {
+        const reader = createAssistantStream(() => callbackPending).getReader();
+        await reader.cancel("consumer stopped");
+        failCallback(new Error("provider stopped"));
+        await callbackPending.catch(() => undefined);
+      });
+
+      expect(unhandledRejections).toEqual([]);
+      expect(consoleError).not.toHaveBeenCalled();
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
+  it("does not settle the outer stream again after a merged stream errors", async () => {
+    let finishCallback!: () => void;
+    const callbackPending = new Promise<void>((resolve) => {
+      finishCallback = resolve;
+    });
+    const streamError = new Error("merged stream failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+
+    try {
+      const unhandledRejections = await captureUnhandledRejections(async () => {
+        const stream = createAssistantStream(async (controller) => {
+          controller.merge(
+            new ReadableStream({
+              start(streamController) {
+                streamController.error(streamError);
+              },
+            }),
+          );
+          await callbackPending;
+        });
+
+        await expect(collectChunks(stream)).rejects.toBe(streamError);
+        finishCallback();
+        await callbackPending;
+      });
+
+      expect(unhandledRejections).toEqual([]);
+      expect(consoleError).toHaveBeenCalledWith(streamError);
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+});
+
 describe("AssistantStreamController withParentId", () => {
   it("attaches parentId to text parts across a data-stream round trip", async () => {
     const response = createAssistantStreamResponse((controller) => {

package/src/core/modules/assistant-stream.ts

@@ -115,7 +115,15 @@ class AssistantStreamControllerImpl implements AssistantStreamController {
   }
 
   get __internal_isClosed() {
-    return this._state.merger.isSealed();
+    return (
+      this._state.merger.isSealed() ||
+      this._state.merger.isCancelled() ||
+      this._state.merger.isErrored()
+    );
+  }
+
+  get __internal_isCancelled() {
+    return this._state.merger.isCancelled();
   }
 
   __internal_getReadable() {
@@ -284,10 +292,11 @@ class AssistantStreamControllerImpl implements AssistantStreamController {
  * Creates an {@link AssistantStream} and writes to it with an
  * {@link AssistantStreamController}.
  *
- * The callback may write synchronously or asynchronously. If it throws, an
- * `error` chunk is emitted before the error is rethrown; when the callback
- * settles, the stream is closed automatically unless the controller was
- * already closed.
+ * The callback may write synchronously or asynchronously. If it throws while
+ * the stream is open, an `error` chunk is emitted. Failures after an explicit
+ * close are logged, while failures after consumer cancellation are discarded.
+ * When the callback settles, the stream is closed automatically unless the
+ * controller was already closed.
  */
 export function createAssistantStream(
   callback: (controller: AssistantStreamController) => PromiseLike<void> | void,
@@ -304,8 +313,9 @@ export function createAssistantStream(
           path: [],
           error: String(e),
         });
+      } else if (!controller.__internal_isCancelled) {
+        console.error(e);
       }
-      throw e;
     } finally {
       if (!controller.__internal_isClosed) {
         controller.close();

package/src/core/modules/tool-call.test.ts

@@ -0,0 +1,67 @@
+import { describe, expect, it, vi } from "vitest";
+import { createAssistantStreamController } from "./assistant-stream";
+import { ToolResponse } from "../tool/ToolResponse";
+import { toolResultStream } from "../tool/toolResultStream";
+import type { ToolCallReader } from "../tool/tool-types";
+import type { AssistantStreamChunk } from "../AssistantStreamChunk";
+
+type Reader = ToolCallReader<Record<string, unknown>, unknown>;
+
+describe("ToolCallStreamController", () => {
+  it("delivers a backend response before an args parse failure", async () => {
+    const [stream, controller] = createAssistantStreamController();
+    let resolveToolReader!: (reader: Reader) => void;
+    const toolReaderPromise = new Promise<Reader>((resolve) => {
+      resolveToolReader = resolve;
+    });
+    const streamCall = vi.fn((reader: Reader) => {
+      resolveToolReader(reader);
+    });
+    const execute = vi.fn();
+    const output = stream.pipeThrough(
+      toolResultStream(
+        {
+          weatherSearch: {
+            parameters: { type: "object", properties: {} },
+            execute,
+            streamCall,
+          },
+        },
+        new AbortController().signal,
+        async () => undefined,
+      ),
+    );
+    const chunks: AssistantStreamChunk[] = [];
+    const drain = output.pipeTo(
+      new WritableStream({
+        write(chunk) {
+          chunks.push(chunk);
+        },
+      }),
+    );
+
+    const toolCall = controller.addToolCallPart({
+      toolCallId: "tool-1",
+      toolName: "weatherSearch",
+    });
+    toolCall.argsText.append('{"query":"London","longitude":0');
+    const reader = await toolReaderPromise;
+    expect(await reader.args.get("query")).toBe("London");
+
+    toolCall.setResponse(new ToolResponse({ result: { source: "backend" } }));
+    toolCall.close();
+    controller.close();
+
+    const response = await reader.response.get();
+    expect(response.result).toEqual({ source: "backend" });
+    expect(response.isError).toBe(false);
+    expect(execute).not.toHaveBeenCalled();
+    await drain;
+    const results = chunks.filter((chunk) => chunk.type === "result");
+    expect(results).toHaveLength(1);
+    expect(results[0]).toMatchObject({
+      result: { source: "backend" },
+      isError: false,
+    });
+  });
+});

package/src/core/modules/tool-call.ts

@@ -65,9 +65,6 @@ class ToolCallStreamControllerImpl implements ToolCallStreamController {
   private _argsTextController!: TextStreamController;
 
   async setResponse(response: ToolResponseLike<ReadonlyJSONValue>) {
-    this._argsTextController.close();
-    await Promise.resolve(); // flush microtask queue
-    // TODO switch argsTextController to be something that doesn'#t require this
     this._controller.enqueue({
       type: "result",
       path: [],
@@ -83,6 +80,9 @@ class ToolCallStreamControllerImpl implements ToolCallStreamController {
         ? { messages: response.messages }
         : {}),
     });
+    this._argsTextController.close();
+    await Promise.resolve(); // flush microtask queue
+    // TODO switch argsTextController to be something that doesn'#t require this
   }
 
   async close() {