assistant-stream 0.3.26 → 0.3.28

package/src/core/serialization/ui-message-stream/UIMessageStream.ts

@@ -1,3 +1,4 @@
+import sjson from "secure-json-parse";
 import type { AssistantStreamChunk } from "../../AssistantStreamChunk";
 import type { ToolCallStreamController } from "../../modules/tool-call";
 import type { TextStreamController } from "../../modules/text";
@@ -149,6 +150,9 @@ export class UIMessageStreamDecoder extends PipeableTransformStream<
                   `Encountered tool result with unknown id: ${chunk.toolCallId}`,
                 );
               }
+              if (toolCallController.argsText === activeToolCallArgsText) {
+                activeToolCallArgsText = undefined;
+              }
               toolCallController.setResponse({
                 result: chunk.result,
                 isError: chunk.isError ?? false,
@@ -222,7 +226,23 @@ export class UIMessageStreamDecoder extends PipeableTransformStream<
                 return;
               }
 
-              const chunk = JSON.parse(event.data);
+              let chunk;
+              try {
+                chunk = sjson.parse(event.data);
+              } catch {
+                chunk = undefined;
+              }
+              if (
+                typeof chunk !== "object" ||
+                chunk === null ||
+                Array.isArray(chunk) ||
+                typeof chunk.type !== "string"
+              ) {
+                console.warn(
+                  `Dropped invalid UIMessageStream chunk: ${event.data.slice(0, 200)}`,
+                );
+                return;
+              }
               if (
                 chunk.type === "text-delta" &&
                 chunk.textDelta === undefined

package/src/core/tool/ToolCallReader.ts

@@ -441,7 +441,11 @@ export class ToolCallArgsReaderImpl<
 export class ToolCallResponseReaderImpl<
   TResult extends ReadonlyJSONValue,
 > implements ToolCallResponseReader<TResult> {
-  constructor(private readonly promise: Promise<ToolResponse<TResult>>) {}
+  private readonly promise: Promise<ToolResponse<TResult>>;
+
+  constructor(promise: Promise<ToolResponse<TResult>>) {
+    this.promise = promise;
+  }
 
   public get() {
     return this.promise;

package/src/core/tool/ToolExecutionStream.ts

@@ -51,6 +51,7 @@ export class ToolExecutionStream extends PipeableTransformStream<
       string,
       ToolCallReaderImpl<ReadonlyJSONObject, ReadonlyJSONValue>
     >();
+    const toolCallIdsWithBackendResult = new Set<string>();
 
     super((readable) => {
       const transform = new TransformStream<
@@ -109,6 +110,7 @@ export class ToolExecutionStream extends PipeableTransformStream<
                   modelContent: chunk.modelContent,
                 }),
               );
+              toolCallIdsWithBackendResult.add(toolCallId);
               break;
             }
             case "tool-call-args-text-finish": {
@@ -123,6 +125,11 @@ export class ToolExecutionStream extends PipeableTransformStream<
               // resolve. Awaited so the close settles before the writer is reused.
               await streamController.finishArgsText();
 
+              // A backend result is authoritative. Closing the args stream still
+              // emits this finish chunk, but must not parse stale/incomplete args,
+              // execute the frontend tool, or enqueue a second result.
+              if (toolCallIdsWithBackendResult.has(toolCallId)) break;
+
               let isExecuting = false;
               const promise = withPromiseOrValue(
                 () => {
@@ -206,10 +213,13 @@ export class ToolExecutionStream extends PipeableTransformStream<
                 toolCallPromise.then(() => {
                   toolCallPromises.delete(toolCallId);
                   toolCallControllers.delete(toolCallId);
+                  toolCallIdsWithBackendResult.delete(toolCallId);
 
                   controller.enqueue(chunk);
                 });
               } else {
+                toolCallControllers.delete(toolCallId);
+                toolCallIdsWithBackendResult.delete(toolCallId);
                 controller.enqueue(chunk);
               }
             }

package/src/core/tool/tool-types.ts

@@ -226,6 +226,13 @@ type ToolBase<
   unstable_backendDefault?: {
     parameters?: boolean;
   };
+
+  /**
+   * Replaces an already-registered same-priority tool with the same name
+   * instead of throwing. Discouraged escape hatch — overwriting a same-priority
+   * tool is usually a design smell; prefer distinct names or priorities.
+   */
+  overwrite?: boolean;
 };
 
 type BackendTool<

package/src/core/tool/toolResultStream.test.ts

@@ -17,6 +17,56 @@ const createDelayedTool = (delay: number, result?: string): Tool => ({
 });
 
 describe("unstable_runPendingTools", () => {
+  it("removes the abort listener after tool execution settles", async () => {
+    const abortController = new AbortController();
+    const addEventListener = vi.spyOn(
+      abortController.signal,
+      "addEventListener",
+    );
+    const removeEventListener = vi.spyOn(
+      abortController.signal,
+      "removeEventListener",
+    );
+    const message: AssistantMessage = {
+      role: "assistant",
+      status: { type: "requires-action", reason: "tool-calls" },
+      parts: [
+        {
+          type: "tool-call",
+          toolCallId: "1",
+          toolName: "tool",
+          args: {},
+        } as ToolCallPart,
+      ],
+      content: [],
+      metadata: {
+        unstable_state: {},
+        unstable_data: [],
+        unstable_annotations: [],
+        steps: [],
+        custom: {},
+      },
+    };
+
+    await unstable_runPendingTools(
+      message,
+      {
+        tool: {
+          parameters: { type: "object", properties: {} },
+          execute: async () => "done",
+        },
+      },
+      abortController.signal,
+      async () => {},
+    );
+
+    expect(addEventListener).toHaveBeenCalledOnce();
+    expect(removeEventListener).toHaveBeenCalledWith(
+      "abort",
+      addEventListener.mock.calls[0]![1],
+    );
+  });
+
   describe("parallel execution", () => {
     it("should run tool calls in parallel", async () => {
       const tool1 = createDelayedTool(100, "Tool 1");

package/src/core/tool/toolResultStream.ts

@@ -59,9 +59,10 @@ function getToolResponse(
 
     // Create abort promise that resolves after 2 microtasks
     // This gives tools that handle abort a chance to win the race
+    let onAbort!: () => void;
     const abortPromise = new Promise<ToolResponse<ReadonlyJSONValue>>(
       (resolve) => {
-        const onAbort = () => {
+        onAbort = () => {
           queueMicrotask(() => {
             queueMicrotask(() => {
               resolve(
@@ -116,7 +117,11 @@ function getToolResponse(
       return response;
     })();
 
-    return Promise.race([executePromise, abortPromise]);
+    try {
+      return await Promise.race([executePromise, abortPromise]);
+    } finally {
+      abortSignal.removeEventListener("abort", onAbort);
+    }
   };
 
   return getResult(tool.execute);

package/src/core/utils/stream/SSE.test.ts

@@ -0,0 +1,97 @@
+import { describe, expect, it, vi } from "vitest";
+import { SSEDecoder, SSEEncoder } from "./SSE";
+
+async function collectChunks<T>(stream: ReadableStream<T>): Promise<T[]> {
+  const reader = stream.getReader();
+  const chunks: T[] = [];
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      chunks.push(value);
+    }
+  } finally {
+    reader.releaseLock();
+  }
+  return chunks;
+}
+
+function createSSEStream(messages: string[]): ReadableStream<Uint8Array> {
+  const encoder = new TextEncoder();
+  const sseText = messages.map((m) => `data: ${m}\n\n`).join("");
+  return new ReadableStream<Uint8Array>({
+    start(controller) {
+      controller.enqueue(encoder.encode(sseText));
+      controller.close();
+    },
+  });
+}
+
+describe("SSEDecoder", () => {
+  it("decodes message events as JSON", async () => {
+    const chunks = await collectChunks(
+      createSSEStream([
+        JSON.stringify({ hello: "world" }),
+        JSON.stringify([1, 2, 3]),
+      ]).pipeThrough(new SSEDecoder()),
+    );
+    expect(chunks).toEqual([{ hello: "world" }, [1, 2, 3]]);
+  });
+
+  it("round-trips values through SSEEncoder and SSEDecoder", async () => {
+    type V = { hello: string } | number[];
+    const values: V[] = [{ hello: "world" }, [1, 2, 3]];
+    const source = new ReadableStream<V>({
+      start(controller) {
+        for (const v of values) controller.enqueue(v);
+        controller.close();
+      },
+    });
+    const decoded = await collectChunks(
+      source.pipeThrough(new SSEEncoder<V>()).pipeThrough(new SSEDecoder<V>()),
+    );
+    expect(decoded).toEqual(values);
+  });
+
+  it("drops frames carrying prototype-pollution keys", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const chunks = await collectChunks(
+      createSSEStream([
+        '{"__proto__":{"polluted":true},"ok":1}',
+        '{"constructor":{"prototype":{"polluted":true}},"ok":2}',
+        JSON.stringify({ ok: 3 }),
+      ]).pipeThrough(new SSEDecoder()),
+    );
+    expect(warn).toHaveBeenCalledTimes(2);
+    expect(chunks).toEqual([{ ok: 3 }]);
+    warn.mockRestore();
+  });
+
+  it("drops frames that are not valid JSON", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const invalid = ["{not json", "not json"];
+    const chunks = await collectChunks(
+      createSSEStream([...invalid, JSON.stringify({ ok: 1 })]).pipeThrough(
+        new SSEDecoder(),
+      ),
+    );
+    expect(warn.mock.calls.map((c) => c[0])).toEqual(
+      invalid.map((m) => `Dropped invalid SSE message: ${m}`),
+    );
+    expect(chunks).toEqual([{ ok: 1 }]);
+    warn.mockRestore();
+  });
+
+  it("errors the stream on an unknown event type", async () => {
+    const encoder = new TextEncoder();
+    const stream = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(encoder.encode("event: custom\ndata: x\n\n"));
+        controller.close();
+      },
+    });
+    await expect(
+      collectChunks(stream.pipeThrough(new SSEDecoder())),
+    ).rejects.toThrow("Unknown SSE event type: custom");
+  });
+});

package/src/core/utils/stream/SSE.ts

@@ -1,3 +1,4 @@
+import sjson from "secure-json-parse";
 import { PipeableTransformStream } from "./PipeableTransformStream";
 import {
   SSEEventDecoderStream,
@@ -44,9 +45,19 @@ export class SSEDecoder<T> extends PipeableTransformStream<
           new TransformStream<PipelineSSEEvent, T>({
             transform(event, controller) {
               switch (event.event) {
-                case "message":
-                  controller.enqueue(JSON.parse(event.data));
+                case "message": {
+                  let value;
+                  try {
+                    value = sjson.parse(event.data);
+                  } catch {
+                    console.warn(
+                      `Dropped invalid SSE message: ${event.data.slice(0, 200)}`,
+                    );
+                    break;
+                  }
+                  controller.enqueue(value);
                   break;
+                }
                 default:
                   throw new Error(`Unknown SSE event type: ${event.event}`);
               }

package/src/core/utils/stream/merge.ts

@@ -9,9 +9,18 @@ type MergeStreamItem = {
 export const createMergeStream = () => {
   const list: MergeStreamItem[] = [];
   let sealed = false;
+  let cancelled = false;
+  let errored = false;
   let controller: ReadableStreamDefaultController<AssistantStreamChunk>;
   let currentPull: ReturnType<typeof promiseWithResolvers<void>> | undefined;
 
+  const cancelAllReaders = () => {
+    list.forEach((item) => {
+      void item.reader.cancel().catch(() => undefined);
+    });
+    list.length = 0;
+  };
+
   const handlePull = (item: MergeStreamItem) => {
     if (!item.promise) {
       // TODO for most streams, we can directly pipeTo to avoid the microTask queue
@@ -23,6 +32,8 @@ export const createMergeStream = () => {
         .read()
         .then(({ done, value }) => {
           item.promise = undefined;
+          if (cancelled || errored) return;
+
           if (done) {
             list.splice(list.indexOf(item), 1);
             if (sealed && list.length === 0) {
@@ -36,12 +47,11 @@ export const createMergeStream = () => {
           currentPull = undefined;
         })
         .catch((e) => {
-          console.error(e);
+          if (cancelled || errored) return;
 
-          list.forEach((item) => {
-            item.reader.cancel();
-          });
-          list.length = 0;
+          errored = true;
+          console.error(e);
+          cancelAllReaders();
 
           controller.error(e);
 
@@ -64,10 +74,10 @@ export const createMergeStream = () => {
       return currentPull.promise;
     },
     cancel() {
-      list.forEach((item) => {
-        item.reader.cancel();
-      });
-      list.length = 0;
+      cancelled = true;
+      cancelAllReaders();
+      currentPull?.resolve();
+      currentPull = undefined;
     },
   });
 
@@ -76,11 +86,23 @@ export const createMergeStream = () => {
     isSealed() {
       return sealed;
     },
+    isCancelled() {
+      return cancelled;
+    },
+    isErrored() {
+      return errored;
+    },
     seal() {
+      if (cancelled || errored) return;
       sealed = true;
       if (list.length === 0) controller.close();
     },
     addStream(stream: ReadableStream<AssistantStreamChunk>) {
+      if (cancelled || errored) {
+        void stream.cancel().catch(() => undefined);
+        return;
+      }
+
       if (sealed)
         throw new Error(
           "Cannot add streams after the run callback has settled.",

package/src/gorp.test.ts

@@ -0,0 +1,32 @@
+import { describe, expect, it } from "vitest";
+import * as entry from "./index";
+
+describe("assistant-transport state exports", () => {
+  it("streams operations end to end through the SSE wire", async () => {
+    const stream = entry.createObjectStream({
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "set", path: ["message"], value: "Hello" },
+        ]);
+        controller.enqueue([
+          { type: "append-text", path: ["message"], value: " World" },
+        ]);
+      },
+    });
+
+    const response = new entry.ObjectStreamResponse(stream);
+    const decoded = entry.fromObjectStreamResponse(response);
+
+    const tracker = new entry.AssistantTransportDeltaTracker();
+    const reader = decoded.getReader();
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      tracker.append(value.operations);
+      expect(tracker.state).toEqual(value.snapshot);
+    }
+
+    expect(tracker.state).toEqual({ message: "Hello World" });
+    expect(tracker.isChangedAt(["message"])).toBe(true);
+  });
+});

package/src/index.ts

@@ -62,12 +62,27 @@ export {
 export type { TextStreamController } from "./core/modules/text";
 export type { ToolCallStreamController } from "./core/modules/tool-call";
 
-export { createObjectStream } from "./core/object/createObjectStream";
-export {
-  ObjectStreamResponse,
-  fromObjectStreamResponse,
-} from "./core/object/ObjectStreamResponse";
-export type { ObjectStreamChunk } from "./core/object/types";
+import { createGorpStream } from "./core/gorp/createGorpStream";
+import {
+  GorpStreamResponse,
+  fromGorpStreamResponse,
+} from "./core/gorp/GorpStreamResponse";
+import type { GorpStreamChunk } from "./core/gorp/types";
+import { GorpStreamDeltaTracker } from "./core/gorp/GorpStreamDeltaTracker";
+
+export { GorpStreamDeltaTracker as AssistantTransportDeltaTracker };
+export type { AssistantTransportStateOperation } from "./core/gorp/types";
+
+/** @deprecated Use the assistant-transport surface instead. */
+export const createObjectStream = createGorpStream;
+/** @deprecated Use the assistant-transport surface instead. */
+export const ObjectStreamResponse = GorpStreamResponse;
+/** @deprecated Use the assistant-transport surface instead. */
+export type ObjectStreamResponse = GorpStreamResponse;
+/** @deprecated Use the assistant-transport surface instead. */
+export const fromObjectStreamResponse = fromGorpStreamResponse;
+/** @deprecated Use the assistant-transport surface instead. */
+export type ObjectStreamChunk = GorpStreamChunk;
 
 export {
   toGenericMessages,

package/src/resumable/ResumableStreamContext.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import { createResumableStreamContext } from "./ResumableStreamContext";
 import { ResumableStreamError } from "./errors";
 import { createInMemoryResumableStreamStore } from "./stores/InMemoryResumableStreamStore";
@@ -8,6 +8,10 @@ const dec = new TextDecoder();
 
 const bytes = (s: string): Uint8Array => enc.encode(s);
 
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
 async function collect(stream: ReadableStream<Uint8Array>): Promise<string> {
   const reader = stream.getReader();
   let out = "";
@@ -271,4 +275,115 @@ describe("createResumableStreamContext", () => {
     expect(errors[0]!.id).toBe("a");
     expect((errors[0]!.error as Error).message).toBe("boom");
   });
+
+  it("continues acquisition when onAcquire throws", async () => {
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    const hookError = new Error("acquire observer failed");
+    const store = createInMemoryResumableStreamStore();
+    const ctx = createResumableStreamContext({
+      store,
+      onAcquire: () => {
+        throw hookError;
+      },
+    });
+
+    const stream = await ctx.run("a", () => makeStringStream(["hello"]));
+
+    expect(await collect(stream)).toBe("hello");
+    expect(await ctx.status("a")).toBe("done");
+    expect(consoleError).toHaveBeenCalledWith(
+      "resumable stream onAcquire hook failed:",
+      hookError,
+    );
+  });
+
+  it("continues streaming when onAppend throws", async () => {
+    vi.spyOn(console, "error").mockImplementation(() => {});
+    const store = createInMemoryResumableStreamStore();
+    const ctx = createResumableStreamContext({
+      store,
+      onAppend: () => {
+        throw new Error("append observer failed");
+      },
+    });
+
+    const stream = await ctx.run("a", () =>
+      makeStringStream(["hello ", "world"]),
+    );
+
+    expect(await collect(stream)).toBe("hello world");
+    expect(await ctx.status("a")).toBe("done");
+  });
+
+  it("continues streaming when async onAppend rejects", async () => {
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    const hookError = new Error("async append observer failed");
+    const store = createInMemoryResumableStreamStore();
+    const ctx = createResumableStreamContext({
+      store,
+      onAppend: async () => {
+        throw hookError;
+      },
+    });
+
+    const stream = await ctx.run("a", () => makeStringStream(["hello"]));
+
+    expect(await collect(stream)).toBe("hello");
+    expect(await ctx.status("a")).toBe("done");
+    await vi.waitFor(() => {
+      expect(consoleError).toHaveBeenCalledWith(
+        "resumable stream onAppend hook failed:",
+        hookError,
+      );
+    });
+  });
+
+  it("keeps successful completion when onFinalize throws", async () => {
+    vi.spyOn(console, "error").mockImplementation(() => {});
+    const onError = vi.fn();
+    const onFinalize = vi.fn(() => {
+      throw new Error("finalize observer failed");
+    });
+    const store = createInMemoryResumableStreamStore();
+    const ctx = createResumableStreamContext({ store, onError, onFinalize });
+
+    const stream = await ctx.run("a", () => makeStringStream(["hello"]));
+
+    expect(await collect(stream)).toBe("hello");
+    expect(await ctx.status("a")).toBe("done");
+    expect(onFinalize).toHaveBeenCalledOnce();
+    expect(onError).not.toHaveBeenCalled();
+  });
+
+  it("finalizes producer errors when onError throws", async () => {
+    vi.spyOn(console, "error").mockImplementation(() => {});
+    const tasks: Promise<unknown>[] = [];
+    const store = createInMemoryResumableStreamStore();
+    const ctx = createResumableStreamContext({
+      store,
+      waitUntil: (task) => tasks.push(task),
+      onError: () => {
+        throw new Error("error observer failed");
+      },
+    });
+    const producerError = new Error("producer failed");
+
+    const stream = await ctx.run(
+      "a",
+      () =>
+        new ReadableStream<Uint8Array>({
+          start(controller) {
+            controller.error(producerError);
+          },
+        }),
+    );
+    await Promise.allSettled(tasks);
+
+    expect(await ctx.status("a")).toBe("error");
+    await expect(collect(stream)).rejects.toThrow("producer failed");
+  });
 });

package/src/resumable/ResumableStreamContext.ts

@@ -1,4 +1,5 @@
 import { ResumableStreamError } from "./errors";
+import { withPromiseOrValue } from "../core/utils/withPromiseOrValue";
 import type {
   ResumableStreamRole,
   ResumableStreamStatus,
@@ -42,6 +43,21 @@ export interface ResumableStreamContext {
   delete(streamId: string): Promise<void>;
 }
 
+const invokeObservabilityHook = <TArgs extends unknown[]>(
+  name: string,
+  hook: ((...args: TArgs) => void) | undefined,
+  ...args: TArgs
+): void => {
+  if (!hook) return;
+  void withPromiseOrValue(
+    () => hook(...args),
+    () => {},
+    (error) => {
+      console.error(`resumable stream ${name} hook failed:`, error);
+    },
+  );
+};
+
 export function createResumableStreamContext(
   options: ResumableStreamContextOptions,
 ): ResumableStreamContext {
@@ -53,7 +69,7 @@ export function createResumableStreamContext(
   return {
     async run(streamId, makeStream) {
       const role = await store.acquire(streamId, acquireOptions);
-      onAcquire?.(streamId, role);
+      invokeObservabilityHook("onAcquire", onAcquire, streamId, role);
       if (role === "producer") {
         startProducerTask(store, streamId, makeStream, {
           waitUntil,
@@ -119,13 +135,18 @@ function startProducerTask(
         const { done, value } = await reader.read();
         if (done) break;
         await store.append(streamId, value);
-        onAppend?.(streamId, value.byteLength);
+        invokeObservabilityHook(
+          "onAppend",
+          onAppend,
+          streamId,
+          value.byteLength,
+        );
       }
       await store.finalize(streamId, "done");
-      onFinalize?.(streamId, "done");
+      invokeObservabilityHook("onFinalize", onFinalize, streamId, "done");
     } catch (err) {
       cancelled = true;
-      onError?.(streamId, err);
+      invokeObservabilityHook("onError", onError, streamId, err);
       const message = err instanceof Error ? err.message : String(err);
       try {
         await reader?.cancel(err);
@@ -134,7 +155,13 @@ function startProducerTask(
       }
       try {
         await store.finalize(streamId, "error", message);
-        onFinalize?.(streamId, "error", message);
+        invokeObservabilityHook(
+          "onFinalize",
+          onFinalize,
+          streamId,
+          "error",
+          message,
+        );
       } catch (finalizeErr) {
         console.error("resumable stream finalize failed:", finalizeErr);
       }