assistant-stream 0.3.25 → 0.3.26

package/src/core/object/ObjectStreamResponse.test.ts

@@ -0,0 +1,32 @@
+import { describe, expect, it } from "vitest";
+import { fromObjectStreamResponse } from "./ObjectStreamResponse";
+
+const createResponse = (contentType?: string) => {
+  const headers = new Headers({
+    "Assistant-Stream-Format": "object-stream/v0",
+  });
+  if (contentType) headers.set("Content-Type", contentType);
+
+  return new Response(new Uint8Array(), { headers });
+};
+
+describe("fromObjectStreamResponse", () => {
+  it.each([
+    "text/event-stream",
+    "text/event-stream; charset=utf-8",
+    "Text/Event-Stream; Charset=UTF-8",
+  ])("accepts the event-stream content type %s", (contentType) => {
+    expect(() =>
+      fromObjectStreamResponse(createResponse(contentType)),
+    ).not.toThrow();
+  });
+
+  it.each([undefined, "application/json", "text/event-streaming"])(
+    "rejects the content type %s",
+    (contentType) => {
+      expect(() =>
+        fromObjectStreamResponse(createResponse(contentType)),
+      ).toThrow("Response is not an event stream");
+    },
+  );
+});

package/src/core/object/ObjectStreamResponse.ts

@@ -99,7 +99,12 @@ export const fromObjectStreamResponse = (
   if (!response.ok)
     throw new Error(`Response failed, status ${response.status}`);
   if (!response.body) throw new Error("Response body is null");
-  if (response.headers.get("Content-Type") !== "text/event-stream") {
+  const mediaType = response.headers
+    .get("Content-Type")
+    ?.split(";", 1)[0]
+    ?.trim()
+    .toLowerCase();
+  if (mediaType !== "text/event-stream") {
     throw new Error("Response is not an event stream");
   }
   if (response.headers.get("Assistant-Stream-Format") !== "object-stream/v0") {

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

@@ -186,6 +186,21 @@ describe("AssistantTransportDecoder", () => {
     });
   });
 
+  it("should discard an unterminated [DONE] event", async () => {
+    const stream = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(new TextEncoder().encode("data: [DONE]\n"));
+        controller.close();
+      },
+    });
+
+    const decodedStream = stream.pipeThrough(new AssistantTransportDecoder());
+
+    await expect(collectChunks(decodedStream)).rejects.toThrow(
+      "Stream ended abruptly without receiving [DONE] marker",
+    );
+  });
+
   it("should throw error when stream ends without [DONE]", async () => {
     // Manually create an SSE stream without [DONE]
     const sseText =
@@ -207,4 +222,30 @@ describe("AssistantTransportDecoder", () => {
       "Stream ended abruptly without receiving [DONE] marker",
     );
   });
+
+  it("round-trips error chunks with code and severity", async () => {
+    const originalChunks: AssistantStreamChunk[] = [
+      {
+        type: "error",
+        path: [],
+        error: "rate limited",
+        code: "provider",
+        severity: "warning",
+      },
+    ];
+
+    const stream = new ReadableStream<AssistantStreamChunk>({
+      start(controller) {
+        for (const chunk of originalChunks) {
+          controller.enqueue(chunk);
+        }
+        controller.close();
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const decodedChunks = await collectChunks(decodedStream);
+
+    expect(decodedChunks).toEqual(originalChunks);
+  });
 });

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

@@ -1,6 +1,9 @@
 import type { AssistantStreamChunk } from "../../AssistantStreamChunk";
 import { PipeableTransformStream } from "../../utils/stream/PipeableTransformStream";
-import { LineDecoderStream } from "../../utils/stream/LineDecoderStream";
+import {
+  SSEEventDecoderStream,
+  type PipelineSSEEvent,
+} from "../../utils/stream/SSEEventDecoderStream";
 import type { AssistantStreamEncoder } from "../../AssistantStream";
 
 /**
@@ -35,72 +38,6 @@ export class AssistantTransportEncoder
   }
 }
 
-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; // Ignore comments
-
-        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,
-          });
-        }
-      },
-    });
-  }
-}
-
 /**
  * AssistantTransportDecoder decodes SSE format into AssistantStreamChunks.
  * It stops decoding when it encounters [DONE].
@@ -115,10 +52,9 @@ export class AssistantTransportDecoder extends PipeableTransformStream<
 
       return readable
         .pipeThrough(new TextDecoderStream())
-        .pipeThrough(new LineDecoderStream())
-        .pipeThrough(new SSEEventStream())
+        .pipeThrough(new SSEEventDecoderStream())
         .pipeThrough(
-          new TransformStream<SSEEvent, AssistantStreamChunk>({
+          new TransformStream<PipelineSSEEvent, AssistantStreamChunk>({
             transform(event, controller) {
               switch (event.event) {
                 case "message":

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

@@ -35,6 +35,34 @@ function createUIMessageStream(events: string[]): ReadableStream<Uint8Array> {
 
 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", id: "text_1", delta: "Hello" }),
+      JSON.stringify({ type: "text-delta", id: "text_1", delta: " world" }),
+      JSON.stringify({ type: "text-end" }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { inputTokens: 10, outputTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    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 legacy textDelta text deltas", async () => {
     const events = [
       JSON.stringify({ type: "start", messageId: "msg_123" }),
       JSON.stringify({ type: "text-start", id: "text_1" }),
@@ -138,7 +166,50 @@ describe("UIMessageStreamDecoder", () => {
     expect(result?.result).toEqual({ temp: 72 });
   });
 
-  it("should decode source parts", async () => {
+  it("should decode source-url parts", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({
+        type: "source-url",
+        sourceId: "src_1",
+        url: "https://example.com",
+        title: "Example",
+      }),
+      JSON.stringify({
+        type: "source-url",
+        sourceId: "src_2",
+        url: "https://example.org",
+      }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { inputTokens: 10, outputTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const sourceStarts = chunks.filter(
+      (c): c is AssistantStreamChunk & { type: "part-start" } =>
+        c.type === "part-start" && c.part.type === "source",
+    );
+    expect(sourceStarts).toHaveLength(2);
+    if (sourceStarts[0]?.part.type === "source") {
+      expect(sourceStarts[0].part.id).toBe("src_1");
+      expect(sourceStarts[0].part.url).toBe("https://example.com");
+      expect(sourceStarts[0].part.title).toBe("Example");
+    }
+    if (sourceStarts[1]?.part.type === "source") {
+      expect(sourceStarts[1].part.id).toBe("src_2");
+      expect(sourceStarts[1].part.url).toBe("https://example.org");
+      expect(sourceStarts[1].part.title).toBeUndefined();
+    }
+  });
+
+  it("should decode legacy source parts", async () => {
     const events = [
       JSON.stringify({ type: "start", messageId: "msg_123" }),
       JSON.stringify({
@@ -168,12 +239,44 @@ describe("UIMessageStreamDecoder", () => {
     );
     expect(sourceStart).toBeDefined();
     if (sourceStart?.part.type === "source") {
+      expect(sourceStart.part.id).toBe("src_1");
       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",
+        url: "https://example.com/image.png",
+        mediaType: "image/png",
+      }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { inputTokens: 10, outputTokens: 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("https://example.com/image.png");
+    }
+  });
+
+  it("should decode legacy file parts", async () => {
     const events = [
       JSON.stringify({ type: "start", messageId: "msg_123" }),
       JSON.stringify({
@@ -277,7 +380,7 @@ describe("UIMessageStreamDecoder", () => {
       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-delta", id: "text_1", delta: "Hello" }),
       JSON.stringify({ type: "text-end" }),
       JSON.stringify({
         type: "finish-step",
@@ -308,6 +411,96 @@ describe("UIMessageStreamDecoder", () => {
     );
     expect(stepFinish).toBeDefined();
     expect(stepFinish?.finishReason).toBe("stop");
+    expect(stepFinish?.usage).toEqual({ inputTokens: 10, outputTokens: 5 });
+    expect(stepFinish?.isContinued).toBe(false);
+  });
+
+  it("should handle the stock v6 lifecycle with bare step chunks", async () => {
+    const events = [
+      JSON.stringify({ type: "start" }),
+      JSON.stringify({ type: "start-step" }),
+      JSON.stringify({ type: "text-start", id: "text_1" }),
+      JSON.stringify({ type: "text-delta", id: "text_1", delta: "Hello" }),
+      JSON.stringify({ type: "text-end" }),
+      JSON.stringify({ type: "finish-step" }),
+      JSON.stringify({ type: "finish", finishReason: "stop" }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const stepStarts = chunks.filter(
+      (c): c is AssistantStreamChunk & { type: "step-start" } =>
+        c.type === "step-start",
+    );
+    expect(stepStarts).toHaveLength(2);
+    expect(stepStarts[0]?.messageId).toBeTruthy();
+    expect(stepStarts[1]?.messageId).toBe(stepStarts[0]?.messageId);
+
+    const stepFinish = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "step-finish" } =>
+        c.type === "step-finish",
+    );
+    expect(stepFinish?.finishReason).toBe("unknown");
+    expect(stepFinish?.usage).toEqual({ inputTokens: 0, outputTokens: 0 });
+    expect(stepFinish?.isContinued).toBe(false);
+
+    const messageFinish = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "message-finish" } =>
+        c.type === "message-finish",
+    );
+    expect(messageFinish?.finishReason).toBe("stop");
+  });
+
+  it("should default finishReason on a bare finish chunk", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "text-start", id: "text_1" }),
+      JSON.stringify({ type: "text-delta", id: "text_1", delta: "Hello" }),
+      JSON.stringify({ type: "text-end" }),
+      JSON.stringify({ type: "finish" }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const messageFinish = chunks.find(
+      (c): c is AssistantStreamChunk & { type: "message-finish" } =>
+        c.type === "message-finish",
+    );
+    expect(messageFinish?.finishReason).toBe("unknown");
+    expect(messageFinish?.usage).toEqual({ inputTokens: 0, outputTokens: 0 });
+  });
+
+  it("should ignore malformed legacy source and file chunks", async () => {
+    const events = [
+      JSON.stringify({ type: "start", messageId: "msg_123" }),
+      JSON.stringify({ type: "source" }),
+      JSON.stringify({ type: "source", source: null }),
+      JSON.stringify({ type: "file" }),
+      JSON.stringify({ type: "file", file: null }),
+      JSON.stringify({
+        type: "finish",
+        finishReason: "stop",
+        usage: { inputTokens: 10, outputTokens: 5 },
+      }),
+      "[DONE]",
+    ];
+
+    const stream = createUIMessageStream(events);
+    const decodedStream = stream.pipeThrough(new UIMessageStreamDecoder());
+    const chunks = await collectChunks(decodedStream);
+
+    const partStarts = chunks.filter(
+      (c) =>
+        c.type === "part-start" &&
+        (c.part.type === "source" || c.part.type === "file"),
+    );
+    expect(partStarts).toHaveLength(0);
   });
 
   it("should handle errors", async () => {
@@ -331,8 +524,8 @@ describe("UIMessageStreamDecoder", () => {
   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';
+      'data: {"type":"text-delta","id":"text_1","delta":"Hello"}\n\n' +
+      'data: {"type":"text-delta","id":"text_1","delta":" world"}\n\n';
 
     const stream = new ReadableStream<Uint8Array>({
       start(controller) {
@@ -348,6 +541,21 @@ describe("UIMessageStreamDecoder", () => {
     );
   });
 
+  it("should discard an unterminated [DONE] event", async () => {
+    const stream = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(new TextEncoder().encode("data: [DONE]\n"));
+        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" }),

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

@@ -3,7 +3,10 @@ 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 {
+  SSEEventDecoderStream,
+  type PipelineSSEEvent,
+} from "../../utils/stream/SSEEventDecoderStream";
 import type {
   UIMessageStreamChunk,
   UIMessageStreamDataChunk,
@@ -21,72 +24,6 @@ export type UIMessageStreamDecoderOptions = {
   }) => 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-");
@@ -271,10 +208,9 @@ export class UIMessageStreamDecoder extends PipeableTransformStream<
 
       return readable
         .pipeThrough(new TextDecoderStream())
-        .pipeThrough(new LineDecoderStream())
-        .pipeThrough(new SSEEventStream())
+        .pipeThrough(new SSEEventDecoderStream())
         .pipeThrough(
-          new TransformStream<SSEEvent, UIMessageStreamChunk>({
+          new TransformStream<PipelineSSEEvent, UIMessageStreamChunk>({
             transform(event, controller) {
               if (event.event !== "message") {
                 throw new Error(`Unknown SSE event type: ${event.event}`);
@@ -286,7 +222,62 @@ export class UIMessageStreamDecoder extends PipeableTransformStream<
                 return;
               }
 
-              controller.enqueue(JSON.parse(event.data));
+              const chunk = JSON.parse(event.data);
+              if (
+                chunk.type === "text-delta" &&
+                chunk.textDelta === undefined
+              ) {
+                const { delta, ...rest } = chunk;
+                controller.enqueue({ ...rest, textDelta: delta ?? "" });
+                return;
+              }
+              if (chunk.type === "start") {
+                controller.enqueue({
+                  ...chunk,
+                  messageId: chunk.messageId ?? generateId(),
+                });
+                return;
+              }
+              if (chunk.type === "source-url") {
+                controller.enqueue({
+                  type: "source",
+                  source: {
+                    sourceType: "url",
+                    id: chunk.sourceId,
+                    url: chunk.url,
+                    ...(chunk.title && { title: chunk.title }),
+                  },
+                });
+                return;
+              }
+              if (chunk.type === "source" && chunk.source == null) return;
+              if (chunk.type === "file" && chunk.file == null) {
+                if (chunk.url === undefined) return;
+                controller.enqueue({
+                  type: "file",
+                  file: { mimeType: chunk.mediaType, data: chunk.url },
+                });
+                return;
+              }
+              if (chunk.type === "finish-step") {
+                controller.enqueue({
+                  ...chunk,
+                  finishReason: chunk.finishReason ?? "unknown",
+                  usage: chunk.usage ?? { inputTokens: 0, outputTokens: 0 },
+                  isContinued: chunk.isContinued ?? false,
+                });
+                return;
+              }
+              if (chunk.type === "finish") {
+                controller.enqueue({
+                  ...chunk,
+                  finishReason: chunk.finishReason ?? "unknown",
+                  usage: chunk.usage ?? { inputTokens: 0, outputTokens: 0 },
+                });
+                return;
+              }
+
+              controller.enqueue(chunk);
             },
             flush() {
               if (!receivedDone) {

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

@@ -40,6 +40,22 @@ describe("LineDecoderStream", () => {
     expect(lines).toEqual(["line1", "line2", "line3"]);
   });
 
+  it("should split lines on CR (\\r)", async () => {
+    const stream = createTextStream(["line1\rline2\rline3\r"]);
+    const lines = await collectLines(
+      stream.pipeThrough(new LineDecoderStream()),
+    );
+    expect(lines).toEqual(["line1", "line2", "line3"]);
+  });
+
+  it("should handle CR line endings split across chunks", async () => {
+    const stream = createTextStream(["line1\r", "line2\r", "line3\r"]);
+    const lines = await collectLines(
+      stream.pipeThrough(new LineDecoderStream()),
+    );
+    expect(lines).toEqual(["line1", "line2", "line3"]);
+  });
+
   it("should handle CRLF split across chunks", async () => {
     // The \r is at the end of one chunk and \n at the start of the next
     const stream = createTextStream(["line1\r", "\nline2\r\n"]);
@@ -65,8 +81,16 @@ describe("LineDecoderStream", () => {
     expect(lines).toEqual(["line1", "", "line2"]);
   });
 
-  it("should handle mixed LF and CRLF", async () => {
-    const stream = createTextStream(["line1\nline2\r\nline3\n"]);
+  it("should handle empty lines with CR", async () => {
+    const stream = createTextStream(["line1\r\rline2\r"]);
+    const lines = await collectLines(
+      stream.pipeThrough(new LineDecoderStream()),
+    );
+    expect(lines).toEqual(["line1", "", "line2"]);
+  });
+
+  it("should handle mixed standard line endings", async () => {
+    const stream = createTextStream(["line1\nline2\r\nline3\r"]);
     const lines = await collectLines(
       stream.pipeThrough(new LineDecoderStream()),
     );