assistant-stream 0.3.27 → 0.3.29

package/src/core/serialization/PlainText.ts

@@ -1,5 +1,9 @@
 import type { AssistantStreamEncoder } from "../AssistantStream";
 import type { AssistantStreamChunk } from "../AssistantStreamChunk";
+import {
+  AssistantMetaTransformStream,
+  type AssistantMetaStreamChunk,
+} from "../utils/stream/AssistantMetaTransformStream";
 import { AssistantTransformStream } from "../utils/stream/AssistantTransformStream";
 import { PipeableTransformStream } from "../utils/stream/PipeableTransformStream";
 
@@ -9,43 +13,20 @@ export class PlainTextEncoder
 {
   headers = new Headers({
     "Content-Type": "text/plain; charset=utf-8",
-    "x-vercel-ai-data-stream": "v1",
   });
 
   constructor() {
     super((readable) => {
-      const transform = new TransformStream<AssistantStreamChunk, string>({
+      const transform = new TransformStream<AssistantMetaStreamChunk, string>({
         transform(chunk, controller) {
-          const type = chunk.type;
-          switch (type) {
-            case "text-delta":
-              controller.enqueue(chunk.textDelta);
-              break;
-
-            case "part-start":
-            case "part-finish":
-            case "step-start":
-            case "step-finish":
-            case "message-finish":
-            case "error":
-              break;
-
-            default: {
-              const unsupportedType:
-                | "tool-call-args-text-finish"
-                | "data"
-                | "annotations"
-                | "tool-call-begin"
-                | "tool-call-delta"
-                | "result"
-                | "update-state" = type;
-              throw new Error(`unsupported chunk type: ${unsupportedType}`);
-            }
+          if (chunk.type === "text-delta" && chunk.meta.type === "text") {
+            controller.enqueue(chunk.textDelta);
           }
         },
       });
 
       return readable
+        .pipeThrough(new AssistantMetaTransformStream())
         .pipeThrough(transform)
         .pipeThrough(new TextEncoderStream());
     });

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

@@ -208,7 +208,7 @@ export class DataStreamEncoder
   }
 }
 
-const TOOL_CALL_ARGS_CLOSING_CHUNKS = [
+const TOOL_CALL_ARGS_CLOSING_CHUNKS: DataStreamStreamChunkType[] = [
   DataStreamStreamChunkType.StartToolCall,
   DataStreamStreamChunkType.ToolCall,
   DataStreamStreamChunkType.TextDelta,

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

@@ -25,29 +25,31 @@ type LanguageModelV1Usage = {
   outputTokens: number;
 };
 
-export enum DataStreamStreamChunkType {
-  TextDelta = "0",
-  Data = "2",
-  Error = "3",
-  Annotation = "8",
-  ToolCall = "9",
-  ToolCallResult = "a",
-  StartToolCall = "b",
-  ToolCallArgsTextDelta = "c",
-  FinishMessage = "d",
-  FinishStep = "e",
-  StartStep = "f",
-  ReasoningDelta = "g",
-  Source = "h",
-  RedactedReasoning = "i",
-  ReasoningSignature = "j",
-  File = "k",
+export const DataStreamStreamChunkType = {
+  TextDelta: "0",
+  Data: "2",
+  Error: "3",
+  Annotation: "8",
+  ToolCall: "9",
+  ToolCallResult: "a",
+  StartToolCall: "b",
+  ToolCallArgsTextDelta: "c",
+  FinishMessage: "d",
+  FinishStep: "e",
+  StartStep: "f",
+  ReasoningDelta: "g",
+  Source: "h",
+  RedactedReasoning: "i",
+  ReasoningSignature: "j",
+  File: "k",
 
-  AuiUpdateStateOperations = "aui-state",
-  AuiTextDelta = "aui-text-delta",
-  AuiReasoningDelta = "aui-reasoning-delta",
-  AuiDataPart = "aui-data",
-}
+  AuiUpdateStateOperations: "aui-state",
+  AuiTextDelta: "aui-text-delta",
+  AuiReasoningDelta: "aui-reasoning-delta",
+  AuiDataPart: "aui-data",
+} as const;
+export type DataStreamStreamChunkType =
+  (typeof DataStreamStreamChunkType)[keyof typeof DataStreamStreamChunkType];
 type DataStreamStreamChunkValue = {
   [DataStreamStreamChunkType.TextDelta]: string;
   [DataStreamStreamChunkType.Data]: ReadonlyJSONValue[];

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();
-  });
-});