assistant-stream 0.3.25 → 0.3.27

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

@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import {
   AssistantTransportEncoder,
   AssistantTransportDecoder,
@@ -47,6 +47,18 @@ async function encodeAndDecode(
   return reconstructedStream.pipeThrough(new AssistantTransportDecoder());
 }
 
+function sseStream(frames: string[]): ReadableStream<AssistantStreamChunk> {
+  const sseText =
+    frames.map((frame) => `data: ${frame}\n\n`).join("") + "data: [DONE]\n\n";
+
+  return new ReadableStream<Uint8Array>({
+    start(controller) {
+      controller.enqueue(new TextEncoder().encode(sseText));
+      controller.close();
+    },
+  }).pipeThrough(new AssistantTransportDecoder());
+}
+
 describe("AssistantTransportEncoder", () => {
   it("should encode text-delta chunks to SSE format", async () => {
     const chunks: AssistantStreamChunk[] = [
@@ -186,6 +198,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 +234,191 @@ describe("AssistantTransportDecoder", () => {
       "Stream ended abruptly without receiving [DONE] marker",
     );
   });
+
+  it("drops frames that are not objects", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        "null",
+        "5",
+        '"text"',
+        "[1,2]",
+        '{"type":"text-delta","textDelta":"Hello","path":[]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "text-delta", textDelta: "Hello", path: [] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(1);
+    expect(warn).toHaveBeenCalledWith(
+      expect.stringContaining("(not-an-object)"),
+    );
+    warn.mockRestore();
+  });
+
+  it("drops frames with a missing or unknown type", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"path":[]}',
+        '{"type":5,"path":[]}',
+        '{"type":"bogus","path":[]}',
+        '{"type":"toString","path":[]}',
+        '{"type":"part-finish","path":[]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([{ type: "part-finish", path: [] }]);
+    expect(warn).toHaveBeenCalledTimes(1);
+    warn.mockRestore();
+  });
+
+  it("drops frames with an invalid path", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"text-delta","textDelta":"b","path":"0"}',
+        '{"type":"text-delta","textDelta":"c","path":[-1]}',
+        '{"type":"text-delta","textDelta":"d","path":[1.5]}',
+        '{"type":"text-delta","textDelta":"e","path":["0"]}',
+        '{"type":"text-delta","textDelta":"f","path":[0]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "text-delta", textDelta: "f", path: [0] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(1);
+    warn.mockRestore();
+  });
+
+  it("normalizes a missing path to an empty path on message-level chunks", async () => {
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"update-state","operations":[{"type":"set","path":["messages"],"value":[]}]}',
+        '{"type":"error","error":"boom"}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      {
+        type: "update-state",
+        operations: [{ type: "set", path: ["messages"], value: [] }],
+        path: [],
+      },
+      { type: "error", error: "boom", path: [] },
+    ]);
+  });
+
+  it("drops part-addressed chunks with a missing path", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"text-delta","textDelta":"a"}',
+        '{"type":"part-finish"}',
+        '{"type":"result","result":{"ok":true}}',
+        '{"type":"text-delta","textDelta":"ok","path":[0]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "text-delta", textDelta: "ok", path: [0] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(3);
+    warn.mockRestore();
+  });
+
+  it("drops frames that are not valid JSON", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"text-delta"',
+        "not json",
+        '{"type":"text-delta","textDelta":"ok","path":[0]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "text-delta", textDelta: "ok", path: [0] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(1);
+    warn.mockRestore();
+  });
+
+  it("drops frames missing per-type required fields", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"part-start","path":[]}',
+        '{"type":"part-start","part":[],"path":[]}',
+        '{"type":"text-delta","textDelta":5,"path":[0]}',
+        '{"type":"annotations","path":[]}',
+        '{"type":"data","path":[]}',
+        '{"type":"update-state","path":[]}',
+        '{"type":"message-finish","path":[]}',
+        '{"type":"step-finish","path":[]}',
+        '{"type":"error","path":[]}',
+        '{"type":"error","error":42,"path":[]}',
+        '{"type":"message-finish","finishReason":"stop","path":[]}',
+        '{"type":"error","error":"boom","path":[]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "error", path: [] },
+      { type: "error", error: 42, path: [] },
+      { type: "message-finish", finishReason: "stop", path: [] },
+      { type: "error", error: "boom", path: [] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(7);
+    expect(warn).toHaveBeenCalledWith(
+      expect.stringContaining("(invalid-fields:part-start)"),
+    );
+    warn.mockRestore();
+  });
+
+  it("tolerates a result without isError or result but rejects a non-boolean isError", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const decodedChunks = await collectChunks(
+      sseStream([
+        '{"type":"result","result":null,"path":[0]}',
+        '{"type":"result","path":[0]}',
+        '{"type":"result","result":1,"isError":"no","path":[0]}',
+      ]),
+    );
+
+    expect(decodedChunks).toEqual([
+      { type: "result", result: null, path: [0] },
+      { type: "result", path: [0] },
+    ]);
+    expect(warn).toHaveBeenCalledTimes(1);
+    warn.mockRestore();
+  });
+
+  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,8 +1,76 @@
+import sjson from "secure-json-parse";
 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";
 
+type ChunkFields = Record<string, unknown>;
+
+type ChunkRule = {
+  kind: "message" | "part-addressed";
+  valid: (chunk: ChunkFields) => boolean;
+};
+
+const noFields = () => true;
+const requiredObject = (key: string) => (c: ChunkFields) => {
+  const value = c[key];
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+};
+const requiredString = (key: string) => (c: ChunkFields) =>
+  typeof c[key] === "string";
+const requiredArray = (key: string) => (c: ChunkFields) =>
+  Array.isArray(c[key]);
+const optionalBoolean = (key: string) => (c: ChunkFields) =>
+  c[key] === undefined || typeof c[key] === "boolean";
+
+const KNOWN_CHUNK_TYPES: Record<AssistantStreamChunk["type"], ChunkRule> = {
+  "part-start": { kind: "message", valid: requiredObject("part") },
+  "part-finish": { kind: "part-addressed", valid: noFields },
+  "tool-call-args-text-finish": { kind: "part-addressed", valid: noFields },
+  "text-delta": { kind: "part-addressed", valid: requiredString("textDelta") },
+  annotations: { kind: "message", valid: requiredArray("annotations") },
+  data: { kind: "message", valid: requiredArray("data") },
+  "step-start": { kind: "message", valid: noFields },
+  "step-finish": { kind: "message", valid: requiredString("finishReason") },
+  "message-finish": { kind: "message", valid: requiredString("finishReason") },
+  result: { kind: "part-addressed", valid: optionalBoolean("isError") },
+  error: { kind: "message", valid: noFields },
+  "update-state": { kind: "message", valid: requiredArray("operations") },
+};
+
+const parseChunk = (data: string): AssistantStreamChunk | string => {
+  let value: unknown;
+  try {
+    value = sjson.parse(data);
+  } catch {
+    return "unparseable";
+  }
+  if (typeof value !== "object" || value === null || Array.isArray(value))
+    return "not-an-object";
+  const { type, path } = value as { type?: unknown; path?: unknown };
+  if (
+    typeof type !== "string" ||
+    !Object.prototype.hasOwnProperty.call(KNOWN_CHUNK_TYPES, type)
+  )
+    return "unknown-type";
+  const rule = KNOWN_CHUNK_TYPES[type as AssistantStreamChunk["type"]];
+  if (!rule.valid(value as Record<string, unknown>))
+    return `invalid-fields:${type}`;
+  if (path === undefined) {
+    if (rule.kind !== "message") return `missing-path:${type}`;
+    return { ...value, path: [] } as unknown as AssistantStreamChunk;
+  }
+  if (
+    !Array.isArray(path) ||
+    !path.every((entry) => Number.isInteger(entry) && entry >= 0)
+  )
+    return "invalid-path";
+  return value as AssistantStreamChunk;
+};
+
 /**
  * AssistantTransportEncoder encodes AssistantStreamChunks into SSE format
  * and emits [DONE] when the stream completes.
@@ -35,72 +103,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].
@@ -112,13 +114,13 @@ export class AssistantTransportDecoder extends PipeableTransformStream<
   constructor() {
     super((readable) => {
       let receivedDone = false;
+      const warnedReasons = new Set<string>();
 
       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":
@@ -128,7 +130,17 @@ export class AssistantTransportDecoder extends PipeableTransformStream<
                     // Stop processing when we encounter [DONE]
                     controller.terminate();
                   } else {
-                    controller.enqueue(JSON.parse(event.data));
+                    const chunk = parseChunk(event.data);
+                    if (typeof chunk === "string") {
+                      if (!warnedReasons.has(chunk)) {
+                        warnedReasons.add(chunk);
+                        console.warn(
+                          `Dropped invalid assistant-transport chunk (${chunk}): ${event.data.slice(0, 200)}`,
+                        );
+                      }
+                    } else {
+                      controller.enqueue(chunk);
+                    }
                   }
                   break;
                 default:

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

@@ -2,7 +2,7 @@ import type {
   ReadonlyJSONObject,
   ReadonlyJSONValue,
 } from "../../../utils/json/json-value";
-import type { ObjectStreamOperation } from "../../object/types";
+import type { GorpStreamOperation } from "../../gorp/types";
 
 export type DataStreamChunk = {
   [K in DataStreamStreamChunkType]: {
@@ -100,7 +100,7 @@ type DataStreamStreamChunkValue = {
   [DataStreamStreamChunkType.File]: { data: string; mimeType: string };
 
   // aui-extensions
-  [DataStreamStreamChunkType.AuiUpdateStateOperations]: ObjectStreamOperation[];
+  [DataStreamStreamChunkType.AuiUpdateStateOperations]: GorpStreamOperation[];
   [DataStreamStreamChunkType.AuiTextDelta]: {
     textDelta: string;
     parentId: string;