assistant-stream 0.3.26 → 0.3.27

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

package/src/core/serialization/assistant-transport/AssistantTransport.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import {
   AssistantTransportEncoder,
   AssistantTransportDecoder,
@@ -47,6 +47,18 @@ async function encodeAndDecode(
   return reconstructedStream.pipeThrough(new AssistantTransportDecoder());
 }
 
+function sseStream(frames: string[]): ReadableStream<AssistantStreamChunk> {
+  const sseText =
+    frames.map((frame) => `data: ${frame}\n\n`).join("") + "data: [DONE]\n\n";
+
+  return new ReadableStream<Uint8Array>({
+    start(controller) {
+      controller.enqueue(new TextEncoder().encode(sseText));
+      controller.close();
+    },
+  }).pipeThrough(new AssistantTransportDecoder());
+}
+
 describe("AssistantTransportEncoder", () => {
   it("should encode text-delta chunks to SSE format", async () => {
     const chunks: AssistantStreamChunk[] = [
@@ -223,6 +235,167 @@ describe("AssistantTransportDecoder", () => {
     );
   });
 
+  it("drops frames that are not objects", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        "null",
+        "5",
+        '"text"',
+        "[1,2]",
+        '{"type":"text-delta","textDelta":"Hello","path":[]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "text-delta", textDelta: "Hello", path: [] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(1);
+    expect(warn).toHaveBeenCalledWith(
+      expect.stringContaining("(not-an-object)"),
+    );
+    warn.mockRestore();
+  });
+
+  it("drops frames with a missing or unknown type", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"path":[]}',
+        '{"type":5,"path":[]}',
+        '{"type":"bogus","path":[]}',
+        '{"type":"toString","path":[]}',
+        '{"type":"part-finish","path":[]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([{ type: "part-finish", path: [] }]);
+    expect(warn).toHaveBeenCalledTimes(1);
+    warn.mockRestore();
+  });
+
+  it("drops frames with an invalid path", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"text-delta","textDelta":"b","path":"0"}',
+        '{"type":"text-delta","textDelta":"c","path":[-1]}',
+        '{"type":"text-delta","textDelta":"d","path":[1.5]}',
+        '{"type":"text-delta","textDelta":"e","path":["0"]}',
+        '{"type":"text-delta","textDelta":"f","path":[0]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "text-delta", textDelta: "f", path: [0] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(1);
+    warn.mockRestore();
+  });
+
+  it("normalizes a missing path to an empty path on message-level chunks", async () => {
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"update-state","operations":[{"type":"set","path":["messages"],"value":[]}]}',
+        '{"type":"error","error":"boom"}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      {
+        type: "update-state",
+        operations: [{ type: "set", path: ["messages"], value: [] }],
+        path: [],
+      },
+      { type: "error", error: "boom", path: [] },
+    ]);
+  });
+
+  it("drops part-addressed chunks with a missing path", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"text-delta","textDelta":"a"}',
+        '{"type":"part-finish"}',
+        '{"type":"result","result":{"ok":true}}',
+        '{"type":"text-delta","textDelta":"ok","path":[0]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "text-delta", textDelta: "ok", path: [0] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(3);
+    warn.mockRestore();
+  });
+
+  it("drops frames that are not valid JSON", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"text-delta"',
+        "not json",
+        '{"type":"text-delta","textDelta":"ok","path":[0]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "text-delta", textDelta: "ok", path: [0] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(1);
+    warn.mockRestore();
+  });
+
+  it("drops frames missing per-type required fields", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"part-start","path":[]}',
+        '{"type":"part-start","part":[],"path":[]}',
+        '{"type":"text-delta","textDelta":5,"path":[0]}',
+        '{"type":"annotations","path":[]}',
+        '{"type":"data","path":[]}',
+        '{"type":"update-state","path":[]}',
+        '{"type":"message-finish","path":[]}',
+        '{"type":"step-finish","path":[]}',
+        '{"type":"error","path":[]}',
+        '{"type":"error","error":42,"path":[]}',
+        '{"type":"message-finish","finishReason":"stop","path":[]}',
+        '{"type":"error","error":"boom","path":[]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "error", path: [] },
+      { type: "error", error: 42, path: [] },
+      { type: "message-finish", finishReason: "stop", path: [] },
+      { type: "error", error: "boom", path: [] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(7);
+    expect(warn).toHaveBeenCalledWith(
+      expect.stringContaining("(invalid-fields:part-start)"),
+    );
+    warn.mockRestore();
+  });
+
+  it("tolerates a result without isError or result but rejects a non-boolean isError", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"result","result":null,"path":[0]}',
+        '{"type":"result","path":[0]}',
+        '{"type":"result","result":1,"isError":"no","path":[0]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "result", result: null, path: [0] },
+      { type: "result", path: [0] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(1);
+    warn.mockRestore();
+  });
+
   it("round-trips error chunks with code and severity", async () => {
     const originalChunks: AssistantStreamChunk[] = [
       {