assistant-stream 0.3.27 → 0.3.28

package/src/core/serialization/data-stream/serialization.test.ts

@@ -0,0 +1,133 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { LineDecoderStream } from "../../utils/stream/LineDecoderStream";
+import {
+  DataStreamChunkDecoder,
+  DataStreamChunkEncoder,
+} from "./serialization";
+import { type DataStreamChunk, DataStreamStreamChunkType } from "./chunk-types";
+
+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 createLineStream(lines: string[]): ReadableStream<string> {
+  return new ReadableStream<string>({
+    start(controller) {
+      for (const line of lines) controller.enqueue(line);
+      controller.close();
+    },
+  });
+}
+
+describe("DataStreamChunkDecoder", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("decodes text-delta chunks", async () => {
+    const chunks = await collectChunks(
+      createLineStream(['0:"hello"', '0:" world"']).pipeThrough(
+        new DataStreamChunkDecoder(),
+      ),
+    );
+    expect(chunks).toEqual([
+      { type: DataStreamStreamChunkType.TextDelta, value: "hello" },
+      { type: DataStreamStreamChunkType.TextDelta, value: " world" },
+    ]);
+  });
+
+  it("round-trips through DataStreamChunkEncoder", async () => {
+    const values: DataStreamChunk[] = [
+      { type: DataStreamStreamChunkType.TextDelta, value: "hello" },
+      { type: DataStreamStreamChunkType.Error, value: "boom" },
+    ];
+    const source = new ReadableStream<DataStreamChunk>({
+      start(controller) {
+        for (const v of values) controller.enqueue(v);
+        controller.close();
+      },
+    });
+    const decoded = await collectChunks(
+      source
+        .pipeThrough(new DataStreamChunkEncoder())
+        .pipeThrough(new LineDecoderStream())
+        .pipeThrough(new DataStreamChunkDecoder()),
+    );
+    expect(decoded).toEqual(values);
+  });
+
+  it("drops chunks carrying prototype-pollution keys", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const chunks = await collectChunks(
+      createLineStream([
+        '0:{"__proto__":{"polluted":true}}',
+        '0:{"constructor":{"prototype":{"polluted":true}}}',
+        '0:"ok"',
+      ]).pipeThrough(new DataStreamChunkDecoder()),
+    );
+    expect(warn).toHaveBeenCalledTimes(2);
+    expect(chunks).toEqual([
+      { type: DataStreamStreamChunkType.TextDelta, value: "ok" },
+    ]);
+  });
+
+  it("drops chunks that are not valid JSON", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const invalid = ["0:{not json", "0:not json"];
+    const chunks = await collectChunks(
+      createLineStream([...invalid, '0:"ok"']).pipeThrough(
+        new DataStreamChunkDecoder(),
+      ),
+    );
+    expect(warn.mock.calls.map((c) => c[0])).toEqual(
+      invalid.map((l) => `Dropped invalid data-stream chunk: ${l}`),
+    );
+    expect(chunks).toEqual([
+      { type: DataStreamStreamChunkType.TextDelta, value: "ok" },
+    ]);
+  });
+
+  it("skips blank and whitespace-only lines silently", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const chunks = await collectChunks(
+      createLineStream([
+        "",
+        '0:"hello"',
+        "   ",
+        "\t",
+        '0:" world"',
+      ]).pipeThrough(new DataStreamChunkDecoder()),
+    );
+    expect(warn).not.toHaveBeenCalled();
+    expect(chunks).toEqual([
+      { type: DataStreamStreamChunkType.TextDelta, value: "hello" },
+      { type: DataStreamStreamChunkType.TextDelta, value: " world" },
+    ]);
+  });
+
+  it("drops a chunk without a colon separator with a warning", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const chunks = await collectChunks(
+      createLineStream(["garbage", '0:"ok"']).pipeThrough(
+        new DataStreamChunkDecoder(),
+      ),
+    );
+    expect(warn).toHaveBeenCalledWith(
+      "Dropped invalid data-stream chunk: garbage",
+    );
+    expect(chunks).toEqual([
+      { type: DataStreamStreamChunkType.TextDelta, value: "ok" },
+    ]);
+  });
+});

package/src/core/serialization/data-stream/serialization.ts

@@ -1,3 +1,4 @@
+import sjson from "secure-json-parse";
 import type { DataStreamChunk, DataStreamStreamChunkType } from "./chunk-types";
 
 export class DataStreamChunkEncoder extends TransformStream<
@@ -21,10 +22,27 @@ export class DataStreamChunkDecoder extends TransformStream<
     super({
       transform: (chunk, controller) => {
         const index = chunk.indexOf(":");
-        if (index === -1) throw new Error("Invalid stream part");
+        if (index === -1) {
+          // Blank lines are benign data-stream framing (keepalive newlines,
+          // replay-buffer separators emitted on resume/reconnect).
+          if (chunk.trim().length === 0) return;
+          console.warn(
+            `Dropped invalid data-stream chunk: ${chunk.slice(0, 200)}`,
+          );
+          return;
+        }
+        let value;
+        try {
+          value = sjson.parse(chunk.slice(index + 1));
+        } catch {
+          console.warn(
+            `Dropped invalid data-stream chunk: ${chunk.slice(0, 200)}`,
+          );
+          return;
+        }
         controller.enqueue({
           type: chunk.slice(0, index) as DataStreamStreamChunkType,
-          value: JSON.parse(chunk.slice(index + 1)),
+          value,
         });
       },
     });

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

@@ -604,6 +604,30 @@ describe("UIMessageStreamDecoder", () => {
     warn.mockRestore();
   });
 
+  it("drops frames carrying prototype-pollution keys", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const events = [
+      '{"type":"text-delta","delta":"x","__proto__":{"polluted":true}}',
+      '{"type":"text-delta","delta":"y","constructor":{"prototype":{"polluted":true}}}',
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "text-delta", id: "t", delta: "ok" }),
+      "[DONE]",
+    ];
+
+    const chunks = await collectChunks(
+      createUIMessageStream(events).pipeThrough(new UIMessageStreamDecoder()),
+    );
+
+    expect(warn).toHaveBeenCalledTimes(2);
+    const deltas = chunks.filter(
+      (c): c is AssistantStreamChunk & { type: "text-delta" } =>
+        c.type === "text-delta",
+    );
+    expect(deltas).toHaveLength(1);
+    expect(deltas[0]!.textDelta).toBe("ok");
+    warn.mockRestore();
+  });
+
   it("drops frames that are not valid JSON", async () => {
     const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
     const invalidFrames = ['{"type":"te', "not json"];

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";
@@ -227,7 +228,7 @@ export class UIMessageStreamDecoder extends PipeableTransformStream<
 
               let chunk;
               try {
-                chunk = JSON.parse(event.data);
+                chunk = sjson.parse(event.data);
               } catch {
                 chunk = 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/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/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/gorp.test.ts

@@ -2,16 +2,6 @@ import { describe, expect, it } from "vitest";
 import * as entry from "./index";
 
 describe("assistant-transport state exports", () => {
-  it("exposes the assistant-transport surface and no gorp-shaped names", () => {
-    expect(typeof entry.AssistantTransportDeltaTracker).toBe("function");
-    expect(typeof entry.createObjectStream).toBe("function");
-    expect(typeof entry.ObjectStreamResponse).toBe("function");
-    expect(typeof entry.fromObjectStreamResponse).toBe("function");
-    for (const name of Object.keys(entry)) {
-      expect(name.toLowerCase()).not.toContain("gorp");
-    }
-  });
-
   it("streams operations end to end through the SSE wire", async () => {
     const stream = entry.createObjectStream({
       execute: (controller) => {

package/src/core/accumulators/TimingTracker.test.ts

@@ -1,100 +0,0 @@
-import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
-import { TimingTracker } from "./TimingTracker";
-
-describe("TimingTracker", () => {
-  beforeEach(() => {
-    vi.useFakeTimers();
-    vi.setSystemTime(1000);
-  });
-
-  afterEach(() => {
-    vi.useRealTimers();
-  });
-
-  it("should record stream start time", () => {
-    const tracker = new TimingTracker();
-    const timing = tracker.getTiming();
-    expect(timing.streamStartTime).toBe(1000);
-  });
-
-  it("should count chunks", () => {
-    const tracker = new TimingTracker();
-    tracker.recordChunk();
-    tracker.recordChunk();
-    tracker.recordChunk();
-    const timing = tracker.getTiming();
-    expect(timing.totalChunks).toBe(3);
-  });
-
-  it("should record first token time", () => {
-    const tracker = new TimingTracker();
-    vi.advanceTimersByTime(50);
-    tracker.recordFirstToken();
-    vi.advanceTimersByTime(100);
-    tracker.recordFirstToken(); // should not overwrite
-    const timing = tracker.getTiming();
-    expect(timing.firstTokenTime).toBe(50);
-  });
-
-  it("should not set firstTokenTime if no text tokens", () => {
-    const tracker = new TimingTracker();
-    const timing = tracker.getTiming();
-    expect(timing.firstTokenTime).toBeUndefined();
-  });
-
-  it("should track unique tool calls", () => {
-    const tracker = new TimingTracker();
-    tracker.recordToolCallStart("tc-1");
-    tracker.recordToolCallStart("tc-2");
-    tracker.recordToolCallStart("tc-1"); // duplicate
-    const timing = tracker.getTiming();
-    expect(timing.toolCallCount).toBe(2);
-  });
-
-  it("should compute totalStreamTime", () => {
-    const tracker = new TimingTracker();
-    vi.advanceTimersByTime(200);
-    const timing = tracker.getTiming();
-    expect(timing.totalStreamTime).toBe(200);
-  });
-
-  it("should use outputTokens when available", () => {
-    const tracker = new TimingTracker();
-    vi.advanceTimersByTime(1000);
-    const timing = tracker.getTiming(42, "some text");
-    expect(timing.tokenCount).toBe(42);
-    expect(timing.tokensPerSecond).toBe(42);
-  });
-
-  it("should estimate tokens from text when no outputTokens", () => {
-    const tracker = new TimingTracker();
-    vi.advanceTimersByTime(1000);
-    // 20 chars / 4 = 5 tokens
-    const timing = tracker.getTiming(undefined, "12345678901234567890");
-    expect(timing.tokenCount).toBe(5);
-    expect(timing.tokensPerSecond).toBe(5);
-  });
-
-  it("should handle zero outputTokens by falling back to text", () => {
-    const tracker = new TimingTracker();
-    vi.advanceTimersByTime(1000);
-    const timing = tracker.getTiming(0, "1234567890"); // 10 chars -> 3 tokens
-    expect(timing.tokenCount).toBe(3);
-  });
-
-  it("should handle no tokens and no text", () => {
-    const tracker = new TimingTracker();
-    vi.advanceTimersByTime(100);
-    const timing = tracker.getTiming();
-    expect(timing.tokenCount).toBeUndefined();
-    expect(timing.tokensPerSecond).toBeUndefined();
-  });
-
-  it("should handle zero stream time (no division by zero)", () => {
-    const tracker = new TimingTracker();
-    // Don't advance time — totalStreamTime = 0
-    const timing = tracker.getTiming(100, "text");
-    expect(timing.totalStreamTime).toBe(0);
-    expect(timing.tokensPerSecond).toBeUndefined();
-  });
-});

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

@@ -1,168 +0,0 @@
-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"]);
-  });
-});