assistant-stream 0.2.45 → 0.2.47

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

@@ -0,0 +1,370 @@
+import { describe, expect, it, vi } from "vitest";
+import { UIMessageStreamDecoder } from "./UIMessageStream";
+import type { AssistantStreamChunk } from "../../AssistantStreamChunk";
+
+// Helper function to collect all chunks from a stream
+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;
+}
+
+// Helper function to create a UI Message Stream from events
+function createUIMessageStream(events: string[]): ReadableStream<Uint8Array> {
+  const encoder = new TextEncoder();
+  const sseText = events.map((e) => `data: ${e}\n\n`).join("");
+
+  return new ReadableStream({
+    start(controller) {
+      controller.enqueue(encoder.encode(sseText));
+      controller.close();
+    },
+  });
+}
+
+describe("UIMessageStreamDecoder", () => {
+  it("should decode text deltas", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "text-start", id: "text_1" }),
+      JSON.stringify({ type: "text-delta", textDelta: "Hello" }),
+      JSON.stringify({ type: "text-delta", textDelta: " world" }),
+      JSON.stringify({ type: "text-end" }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { promptTokens: 10, completionTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    // Find text-delta chunks
+    const textDeltas = chunks.filter(
+      (c): c is AssistantStreamChunk & { type: "text-delta" } =>
+        c.type === "text-delta",
+    );
+    expect(textDeltas).toHaveLength(2);
+    expect(textDeltas[0]?.textDelta).toBe("Hello");
+    expect(textDeltas[1]?.textDelta).toBe(" world");
+  });
+
+  it("should decode reasoning parts", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "reasoning-start", id: "reasoning_1" }),
+      JSON.stringify({ type: "reasoning-delta", delta: "Let me think..." }),
+      JSON.stringify({ type: "reasoning-end" }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { promptTokens: 10, completionTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    // Find part-start for reasoning
+    const partStarts = chunks.filter(
+      (c): c is AssistantStreamChunk & { type: "part-start" } =>
+        c.type === "part-start",
+    );
+    expect(partStarts.some((p) => p.part.type === "reasoning")).toBe(true);
+  });
+
+  it("should decode tool calls", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({
+        type: "tool-call-start",
+        id: "tc_1",
+        toolCallId: "call_abc",
+        toolName: "weather",
+      }),
+      JSON.stringify({ type: "tool-call-delta", argsText: '{"city":' }),
+      JSON.stringify({ type: "tool-call-delta", argsText: '"NYC"}' }),
+      JSON.stringify({ type: "tool-call-end" }),
+      JSON.stringify({
+        type: "tool-result",
+        toolCallId: "call_abc",
+        result: { temp: 72 },
+      }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { promptTokens: 10, completionTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    // Find tool-call part-start
+    const toolCallStart = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "part-start" } =>
+        c.type === "part-start" && c.part.type === "tool-call",
+    );
+    expect(toolCallStart).toBeDefined();
+    if (toolCallStart?.part.type === "tool-call") {
+      expect(toolCallStart.part.toolName).toBe("weather");
+      expect(toolCallStart.part.toolCallId).toBe("call_abc");
+    }
+
+    // Find result
+    const result = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "result" } =>
+        c.type === "result",
+    );
+    expect(result).toBeDefined();
+    expect(result?.result).toEqual({ temp: 72 });
+  });
+
+  it("should decode source parts", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({
+        type: "source",
+        source: {
+          sourceType: "url",
+          id: "src_1",
+          url: "https://example.com",
+          title: "Example",
+        },
+      }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { promptTokens: 10, completionTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const sourceStart = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "part-start" } =>
+        c.type === "part-start" && c.part.type === "source",
+    );
+    expect(sourceStart).toBeDefined();
+    if (sourceStart?.part.type === "source") {
+      expect(sourceStart.part.url).toBe("https://example.com");
+      expect(sourceStart.part.title).toBe("Example");
+    }
+  });
+
+  it("should decode file parts", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({
+        type: "file",
+        file: {
+          mimeType: "image/png",
+          data: "base64data...",
+        },
+      }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { promptTokens: 10, completionTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const fileStart = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "part-start" } =>
+        c.type === "part-start" && c.part.type === "file",
+    );
+    expect(fileStart).toBeDefined();
+    if (fileStart?.part.type === "file") {
+      expect(fileStart.part.mimeType).toBe("image/png");
+      expect(fileStart.part.data).toBe("base64data...");
+    }
+  });
+
+  it("should handle data-* chunks", async () => {
+    const onData = vi.fn();
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({
+        type: "data-weather",
+        data: { temp: 72, city: "NYC" },
+      }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { promptTokens: 10, completionTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(
+      new UIMessageStreamDecoder({ onData }),
+    );
+    await collectChunks(decodedStream);
+
+    expect(onData).toHaveBeenCalledWith({
+      type: "data-weather",
+      name: "weather",
+      data: { temp: 72, city: "NYC" },
+      transient: undefined,
+    });
+  });
+
+  it("should handle transient data-* chunks", async () => {
+    const onData = vi.fn();
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({
+        type: "data-progress",
+        transient: true,
+        data: { percent: 50 },
+      }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { promptTokens: 10, completionTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(
+      new UIMessageStreamDecoder({ onData }),
+    );
+    const chunks = await collectChunks(decodedStream);
+
+    // Transient data should call onData
+    expect(onData).toHaveBeenCalledWith({
+      type: "data-progress",
+      name: "progress",
+      data: { percent: 50 },
+      transient: true,
+    });
+
+    // Transient data should NOT emit a data chunk
+    const dataChunks = chunks.filter((c) => c.type === "data");
+    expect(dataChunks).toHaveLength(0);
+  });
+
+  it("should handle step lifecycle", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "start-step", messageId: "step_1" }),
+      JSON.stringify({ type: "text-start", id: "text_1" }),
+      JSON.stringify({ type: "text-delta", textDelta: "Hello" }),
+      JSON.stringify({ type: "text-end" }),
+      JSON.stringify({
+        type: "finish-step",
+        finishReason: "stop",
+        usage: { promptTokens: 10, completionTokens: 5 },
+        isContinued: false,
+      }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { promptTokens: 10, completionTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    // Find step-start chunks
+    const stepStarts = chunks.filter((c) => c.type === "step-start");
+    expect(stepStarts.length).toBeGreaterThanOrEqual(1);
+
+    // Find step-finish
+    const stepFinish = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "step-finish" } =>
+        c.type === "step-finish",
+    );
+    expect(stepFinish).toBeDefined();
+    expect(stepFinish?.finishReason).toBe("stop");
+  });
+
+  it("should handle errors", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "error", errorText: "Something went wrong" }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const errorChunk = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "error" } => c.type === "error",
+    );
+    expect(errorChunk).toBeDefined();
+    expect(errorChunk?.error).toBe("Something went wrong");
+  });
+
+  it("should throw when stream ends without [DONE]", async () => {
+    const encoder = new TextEncoder();
+    const sseText =
+      'data: {"type":"text-delta","textDelta":"Hello"}\n\n' +
+      'data: {"type":"text-delta","textDelta":" world"}\n\n';
+
+    const stream = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(encoder.encode(sseText));
+        controller.close();
+      },
+    });
+
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+
+    await expect(collectChunks(decodedStream)).rejects.toThrow(
+      "Stream ended abruptly without receiving [DONE] marker",
+    );
+  });
+
+  it("should ignore unknown chunk types for forward compatibility", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "unknown-future-type", data: {} }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { promptTokens: 10, completionTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    // Should not throw, should complete successfully
+    expect(chunks.some((c) => c.type === "message-finish")).toBe(true);
+  });
+});

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

@@ -0,0 +1,300 @@
+import type { AssistantStreamChunk } from "../../AssistantStreamChunk";
+import type { ToolCallStreamController } from "../../modules/tool-call";
+import type { TextStreamController } from "../../modules/text";
+import { AssistantTransformStream } from "../../utils/stream/AssistantTransformStream";
+import { PipeableTransformStream } from "../../utils/stream/PipeableTransformStream";
+import { LineDecoderStream } from "../../utils/stream/LineDecoderStream";
+import {
+  type UIMessageStreamChunk,
+  type UIMessageStreamDataChunk,
+} from "./chunk-types";
+import { generateId } from "../../utils/generateId";
+
+export type { UIMessageStreamChunk, UIMessageStreamDataChunk };
+
+export type UIMessageStreamDecoderOptions = {
+  onData?: (data: {
+    type: string;
+    name: string;
+    data: unknown;
+    transient?: boolean;
+  }) => void;
+};
+
+type SSEEvent = {
+  event: string;
+  data: string;
+  id?: string | undefined;
+  retry?: number | undefined;
+};
+
+class SSEEventStream extends TransformStream<string, SSEEvent> {
+  constructor() {
+    let eventBuffer: Partial<SSEEvent> = {};
+    let dataLines: string[] = [];
+
+    super({
+      start() {
+        eventBuffer = {};
+        dataLines = [];
+      },
+      transform(line, controller) {
+        if (line.startsWith(":")) return;
+
+        if (line === "") {
+          if (dataLines.length > 0) {
+            controller.enqueue({
+              event: eventBuffer.event || "message",
+              data: dataLines.join("\n"),
+              id: eventBuffer.id,
+              retry: eventBuffer.retry,
+            });
+          }
+          eventBuffer = {};
+          dataLines = [];
+          return;
+        }
+
+        const [field, ...rest] = line.split(":");
+        const value = rest.join(":").trimStart();
+
+        switch (field) {
+          case "event":
+            eventBuffer.event = value;
+            break;
+          case "data":
+            dataLines.push(value);
+            break;
+          case "id":
+            eventBuffer.id = value;
+            break;
+          case "retry":
+            eventBuffer.retry = Number(value);
+            break;
+        }
+      },
+      flush(controller) {
+        if (dataLines.length > 0) {
+          controller.enqueue({
+            event: eventBuffer.event || "message",
+            data: dataLines.join("\n"),
+            id: eventBuffer.id,
+            retry: eventBuffer.retry,
+          });
+        }
+      },
+    });
+  }
+}
+
+const isDataChunk = (
+  chunk: UIMessageStreamChunk,
+): chunk is UIMessageStreamDataChunk => chunk.type.startsWith("data-");
+
+/**
+ * Decodes AI SDK v6 UI Message Stream format into AssistantStreamChunks.
+ */
+export class UIMessageStreamDecoder extends PipeableTransformStream<
+  Uint8Array<ArrayBuffer>,
+  AssistantStreamChunk
+> {
+  constructor(options: UIMessageStreamDecoderOptions = {}) {
+    super((readable) => {
+      const toolCallControllers = new Map<string, ToolCallStreamController>();
+      let activeToolCallArgsText: TextStreamController | undefined;
+      let currentMessageId: string | undefined;
+      let receivedDone = false;
+
+      const transform = new AssistantTransformStream<UIMessageStreamChunk>({
+        transform(chunk, controller) {
+          const type = chunk.type;
+
+          if (isDataChunk(chunk)) {
+            const name = chunk.type.slice(5);
+
+            if (options.onData) {
+              options.onData({
+                type: chunk.type,
+                name,
+                data: chunk.data,
+                ...(chunk.transient !== undefined && {
+                  transient: chunk.transient,
+                }),
+              });
+            }
+
+            if (!chunk.transient) {
+              controller.enqueue({
+                type: "data",
+                path: [],
+                data: [{ name, data: chunk.data }],
+              });
+            }
+            return;
+          }
+
+          switch (type) {
+            case "start":
+              currentMessageId = chunk.messageId;
+              controller.enqueue({
+                type: "step-start",
+                path: [],
+                messageId: chunk.messageId,
+              });
+              break;
+
+            case "text-start":
+            case "text-end":
+            case "reasoning-start":
+            case "reasoning-end":
+              break;
+
+            case "text-delta":
+              controller.appendText(chunk.textDelta);
+              break;
+
+            case "reasoning-delta":
+              controller.appendReasoning(chunk.delta);
+              break;
+
+            case "source":
+              controller.appendSource({
+                type: "source",
+                sourceType: chunk.source.sourceType,
+                id: chunk.source.id,
+                url: chunk.source.url,
+                ...(chunk.source.title && { title: chunk.source.title }),
+              });
+              break;
+
+            case "file":
+              controller.appendFile({
+                type: "file",
+                mimeType: chunk.file.mimeType,
+                data: chunk.file.data,
+              });
+              break;
+
+            case "tool-call-start": {
+              activeToolCallArgsText?.close();
+              activeToolCallArgsText = undefined;
+
+              if (toolCallControllers.has(chunk.toolCallId)) {
+                throw new Error(
+                  `Encountered duplicate tool call id: ${chunk.toolCallId}`,
+                );
+              }
+
+              const toolCallController = controller.addToolCallPart({
+                toolCallId: chunk.toolCallId,
+                toolName: chunk.toolName,
+              });
+              toolCallControllers.set(chunk.toolCallId, toolCallController);
+              activeToolCallArgsText = toolCallController.argsText;
+              break;
+            }
+
+            case "tool-call-delta":
+              activeToolCallArgsText?.append(chunk.argsText);
+              break;
+
+            case "tool-call-end":
+              activeToolCallArgsText?.close();
+              activeToolCallArgsText = undefined;
+              break;
+
+            case "tool-result": {
+              const toolCallController = toolCallControllers.get(
+                chunk.toolCallId,
+              );
+              if (!toolCallController) {
+                throw new Error(
+                  `Encountered tool result with unknown id: ${chunk.toolCallId}`,
+                );
+              }
+              toolCallController.setResponse({
+                result: chunk.result,
+                isError: chunk.isError ?? false,
+              });
+              break;
+            }
+
+            case "start-step":
+              controller.enqueue({
+                type: "step-start",
+                path: [],
+                messageId: chunk.messageId ?? currentMessageId ?? generateId(),
+              });
+              break;
+
+            case "finish-step":
+              controller.enqueue({
+                type: "step-finish",
+                path: [],
+                finishReason: chunk.finishReason,
+                usage: chunk.usage,
+                isContinued: chunk.isContinued,
+              });
+              break;
+
+            case "finish":
+              controller.enqueue({
+                type: "message-finish",
+                path: [],
+                finishReason: chunk.finishReason,
+                usage: chunk.usage,
+              });
+              break;
+
+            case "error":
+              controller.enqueue({
+                type: "error",
+                path: [],
+                error: chunk.errorText,
+              });
+              break;
+
+            default:
+              // ignore unknown types for forward compatibility
+              break;
+          }
+        },
+        flush() {
+          activeToolCallArgsText?.close();
+          toolCallControllers.forEach((ctrl) => ctrl.close());
+          toolCallControllers.clear();
+        },
+      });
+
+      return readable
+        .pipeThrough(new TextDecoderStream())
+        .pipeThrough(new LineDecoderStream())
+        .pipeThrough(new SSEEventStream())
+        .pipeThrough(
+          new TransformStream<SSEEvent, UIMessageStreamChunk>({
+            transform(event, controller) {
+              if (event.event !== "message") {
+                throw new Error(`Unknown SSE event type: ${event.event}`);
+              }
+
+              if (event.data === "[DONE]") {
+                receivedDone = true;
+                controller.terminate();
+                return;
+              }
+
+              controller.enqueue(JSON.parse(event.data));
+            },
+            flush() {
+              if (!receivedDone) {
+                throw new Error(
+                  "Stream ended abruptly without receiving [DONE] marker",
+                );
+              }
+            },
+          }),
+        )
+        .pipeThrough(transform);
+    });
+  }
+}