assistant-stream 0.3.25 → 0.3.27

package/src/core/accumulators/assistant-message-accumulator.test.ts

@@ -1,4 +1,4 @@
-import { describe, it, expect } from "vitest";
+import { describe, it, expect, vi } from "vitest";
 import { AssistantMessageAccumulator } from "./assistant-message-accumulator";
 import type { AssistantStreamChunk } from "../AssistantStreamChunk";
 import type { AssistantMessage } from "../utils/types";
@@ -224,4 +224,440 @@ describe("AssistantMessageAccumulator timing", () => {
     expect(last.metadata.timing).toBeDefined();
     expect(last.metadata.timing!.firstTokenTime).toBeTypeOf("number");
   });
+
+  it("does not record firstTokenTime when a text-delta is dropped for an out-of-range path", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      { type: "part-start", path: [0], part: { type: "text" } },
+      { type: "text-delta", path: [5], textDelta: "x" },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.metadata.timing).toBeDefined();
+    expect(last.metadata.timing!.firstTokenTime).toBeUndefined();
+    warn.mockRestore();
+  });
+
+  it("does not record firstTokenTime when a text-delta is dropped for a nested path", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      { type: "part-start", path: [0], part: { type: "text" } },
+      { type: "text-delta", path: [0, 1], textDelta: "x" },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.metadata.timing).toBeDefined();
+    expect(last.metadata.timing!.firstTokenTime).toBeUndefined();
+    warn.mockRestore();
+  });
+
+  it("does not record firstTokenTime when a text-delta is dropped for a wrong part type", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      {
+        type: "part-start",
+        path: [0],
+        part: {
+          type: "source",
+          sourceType: "url",
+          id: "s1",
+          url: "https://example.com",
+        },
+      },
+      { type: "text-delta", path: [0], textDelta: "x" },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.metadata.timing).toBeDefined();
+    expect(last.metadata.timing!.firstTokenTime).toBeUndefined();
+    warn.mockRestore();
+  });
+
+  it("records firstTokenTime on the first applied text-delta, not a preceding dropped one", async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(1000);
+
+    const accumulator = new AssistantMessageAccumulator();
+    const writer = accumulator.writable.getWriter();
+    const messages: AssistantMessage[] = [];
+    const reading = (async () => {
+      const reader = accumulator.readable.getReader();
+      for (;;) {
+        const { done, value } = await reader.read();
+        if (done) break;
+        messages.push(value);
+      }
+    })();
+
+    await writer.write({
+      type: "part-start",
+      path: [0],
+      part: { type: "text" },
+    });
+    vi.setSystemTime(2000);
+    await writer.write({ type: "text-delta", path: [5], textDelta: "x" });
+    vi.setSystemTime(3000);
+    await writer.write({ type: "text-delta", path: [0], textDelta: "hi" });
+    await writer.close();
+    await reading;
+    vi.useRealTimers();
+
+    const last = messages.at(-1)!;
+    expect(last.metadata.timing).toBeDefined();
+    expect(last.metadata.timing!.firstTokenTime).toBe(2000);
+  });
+});
+
+describe("AssistantMessageAccumulator error chunks", () => {
+  it("defaults missing code to unknown", async () => {
+    const messages = await collectStream([
+      { type: "error", path: [], error: "stream failed" },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.status).toEqual({
+      type: "incomplete",
+      reason: "error",
+      error: { code: "unknown", message: "stream failed" },
+    });
+  });
+
+  it("defaults a missing error message", async () => {
+    const messages = await collectStream([
+      { type: "error", path: [] } as unknown as AssistantStreamChunk,
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.status).toEqual({
+      type: "incomplete",
+      reason: "error",
+      error: { code: "unknown", message: "unknown error" },
+    });
+  });
+
+  it("carries code and severity through unchanged", async () => {
+    const messages = await collectStream([
+      {
+        type: "error",
+        path: [],
+        error: "rate limited",
+        code: "provider",
+        severity: "warning",
+      },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.status).toEqual({
+      type: "incomplete",
+      reason: "error",
+      error: {
+        code: "provider",
+        message: "rate limited",
+        severity: "warning",
+      },
+    });
+  });
+
+  it("drops out-of-taxonomy severity at ingestion", async () => {
+    const messages = await collectStream([
+      {
+        type: "error",
+        path: [],
+        error: "boom",
+        code: "unknown",
+        severity: "fatal",
+      } as unknown as AssistantStreamChunk,
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.status).toEqual({
+      type: "incomplete",
+      reason: "error",
+      error: {
+        code: "unknown",
+        message: "boom",
+      },
+    });
+    expect(last.status).toMatchObject({
+      error: expect.not.objectContaining({ severity: expect.anything() }),
+    });
+  });
+});
+
+describe("AssistantMessageAccumulator part path bounds", () => {
+  it("drops a text-delta whose path is out of range with a warning", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      { type: "part-start", path: [0], part: { type: "text" } },
+      { type: "text-delta", path: [5], textDelta: "x" },
+      { type: "text-delta", path: [0], textDelta: "hi" },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.parts).toHaveLength(1);
+    expect(last.parts[0]).toMatchObject({ type: "text", text: "hi" });
+    expect(warn).toHaveBeenCalledWith(
+      "Dropped text-delta chunk: no part at path [5]",
+    );
+    warn.mockRestore();
+  });
+
+  it("drops an out-of-range part-finish instead of fabricating a part", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      { type: "part-start", path: [0], part: { type: "text" } },
+      { type: "part-finish", path: [3] },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.parts).toHaveLength(1);
+    expect(last.parts[0]!.type).toBe("text");
+    warn.mockRestore();
+  });
+
+  it("drops part chunks when no parts exist", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      { type: "text-delta", path: [0], textDelta: "x" },
+      { type: "part-start", path: [0], part: { type: "text" } },
+      { type: "text-delta", path: [0], textDelta: "ok" },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.parts).toHaveLength(1);
+    expect(last.parts[0]).toMatchObject({ type: "text", text: "ok" });
+    warn.mockRestore();
+  });
+
+  it("drops chunks with nested paths", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      { type: "part-start", path: [0], part: { type: "text" } },
+      { type: "text-delta", path: [0, 1], textDelta: "x" },
+      { type: "text-delta", path: [0], textDelta: "ok" },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.parts).toHaveLength(1);
+    expect(last.parts[0]).toMatchObject({ type: "text", text: "ok" });
+    warn.mockRestore();
+  });
+
+  it("drops an out-of-range result without touching existing tool calls", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      {
+        type: "part-start",
+        path: [0],
+        part: { type: "tool-call", toolCallId: "t1", toolName: "f" },
+      },
+      { type: "text-delta", path: [0], textDelta: "{}" },
+      { type: "result", path: [7], result: { ok: true }, isError: false },
+      { type: "result", path: [0], result: { ok: true }, isError: false },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.parts).toHaveLength(1);
+    expect(last.parts[0]).toMatchObject({
+      type: "tool-call",
+      state: "result",
+      result: { ok: true },
+    });
+    warn.mockRestore();
+  });
+});
+
+describe("AssistantMessageAccumulator wrong part type chunks", () => {
+  const sourcePart = {
+    type: "source",
+    sourceType: "url",
+    id: "s1",
+    url: "https://example.com",
+  } as const;
+
+  it("drops a text-delta addressed to a source part", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      { type: "part-start", path: [0], part: sourcePart },
+      { type: "text-delta", path: [0], textDelta: "x" },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.parts[0]).toMatchObject({ type: "source", id: "s1" });
+    expect(warn).toHaveBeenCalledWith(
+      "Dropped text-delta chunk: part is neither text nor tool-call",
+    );
+    warn.mockRestore();
+  });
+
+  it("drops a result addressed to a text part", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      { type: "part-start", path: [0], part: { type: "text" } },
+      { type: "text-delta", path: [0], textDelta: "hi" },
+      { type: "result", path: [0], result: { ok: true }, isError: false },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.parts[0]).toMatchObject({ type: "text", text: "hi" });
+    expect(warn).toHaveBeenCalledWith(
+      "Dropped result chunk: part is not a tool-call",
+    );
+    warn.mockRestore();
+  });
+
+  it("drops an args-text-finish addressed to a text part", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      { type: "part-start", path: [0], part: { type: "text" } },
+      { type: "tool-call-args-text-finish", path: [0] },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.parts[0]).toMatchObject({ type: "text" });
+    expect(warn).toHaveBeenCalledWith(
+      "Dropped tool-call-args-text-finish chunk: part is not a tool-call",
+    );
+    warn.mockRestore();
+  });
+
+  it("inserts a placeholder for an unsupported part-start so later part indices stay aligned", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      {
+        type: "part-start",
+        path: [0],
+        part: { type: "bogus" },
+      } as unknown as AssistantStreamChunk,
+      { type: "part-start", path: [1], part: { type: "text" } },
+      { type: "text-delta", path: [1], textDelta: "ok" },
+      { type: "part-finish", path: [0] },
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.parts).toHaveLength(2);
+    expect(last.parts[0]).toMatchObject({
+      type: "reasoning",
+      text: "",
+      status: { type: "complete", reason: "unknown" },
+    });
+    expect(last.parts[1]).toMatchObject({ type: "text", text: "ok" });
+    expect(warn).toHaveBeenCalledWith(
+      "Unsupported part-start type bogus: inserting an empty reasoning part to preserve part indices",
+    );
+    warn.mockRestore();
+  });
+});
+
+describe("AssistantMessageAccumulator warn dedup", () => {
+  it("warns once per drop class per instance", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    await collectStream([
+      { type: "part-start", path: [0], part: { type: "text" } },
+      { type: "text-delta", path: [5], textDelta: "a" },
+      { type: "text-delta", path: [5], textDelta: "b" },
+      { type: "text-delta", path: [6], textDelta: "c" },
+      { type: "part-finish", path: [9] },
+    ]);
+
+    expect(warn).toHaveBeenCalledTimes(2);
+    expect(warn).toHaveBeenCalledWith(
+      "Dropped text-delta chunk: no part at path [5]",
+    );
+    expect(warn).toHaveBeenCalledWith(
+      "Dropped part-finish chunk: no part at path [9]",
+    );
+    warn.mockRestore();
+  });
+
+  it("resets dedup per accumulator instance", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const stream = [
+      { type: "part-start", path: [0], part: { type: "text" } },
+      { type: "text-delta", path: [5], textDelta: "x" },
+    ] satisfies AssistantStreamChunk[];
+    await collectStream(stream);
+    await collectStream(stream);
+
+    expect(warn).toHaveBeenCalledTimes(2);
+    warn.mockRestore();
+  });
+
+  it("dedupes repeated wrong-part-type drops", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    await collectStream([
+      {
+        type: "part-start",
+        path: [0],
+        part: { type: "source", sourceType: "url", id: "s", url: "https://x" },
+      },
+      { type: "text-delta", path: [0], textDelta: "a" },
+      { type: "text-delta", path: [0], textDelta: "b" },
+    ]);
+
+    expect(warn).toHaveBeenCalledTimes(1);
+    warn.mockRestore();
+  });
+});
+
+describe("AssistantMessageAccumulator warn dedup key independence", () => {
+  it("keeps drop classes as independent keys within one stream", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    await collectStream([
+      {
+        type: "part-start",
+        path: [0],
+        part: { type: "source", sourceType: "url", id: "s", url: "https://x" },
+      },
+      { type: "text-delta", path: [0], textDelta: "wrong-part" },
+      { type: "text-delta", path: [7], textDelta: "no-part" },
+    ]);
+
+    expect(warn).toHaveBeenCalledTimes(2);
+    warn.mockRestore();
+  });
+
+  it("keys unsupported part-starts per type", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const messages = await collectStream([
+      {
+        type: "part-start",
+        path: [0],
+        part: { type: "bogus" },
+      } as unknown as AssistantStreamChunk,
+      {
+        type: "part-start",
+        path: [1],
+        part: { type: "bogus" },
+      } as unknown as AssistantStreamChunk,
+      {
+        type: "part-start",
+        path: [2],
+        part: { type: "video" },
+      } as unknown as AssistantStreamChunk,
+    ]);
+    const last = messages.at(-1)!;
+
+    expect(last.parts).toHaveLength(3);
+    expect(warn).toHaveBeenCalledTimes(2);
+    warn.mockRestore();
+  });
+
+  it("caps the number of warned keys per instance", async () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    await collectStream(
+      Array.from(
+        { length: 20 },
+        (_, i) =>
+          ({
+            type: "part-start",
+            path: [i],
+            part: { type: `bogus-${i}` },
+          }) as unknown as AssistantStreamChunk,
+      ),
+    );
+
+    expect(warn).toHaveBeenCalledTimes(16);
+    warn.mockRestore();
+  });
 });

package/src/core/accumulators/assistant-message-accumulator.ts

@@ -13,7 +13,7 @@ import type {
   FilePart,
   DataPart,
 } from "../utils/types";
-import { ObjectStreamAccumulator } from "../object/ObjectStreamAccumulator";
+import { GorpStreamAccumulator } from "../gorp/GorpStreamAccumulator";
 import type { ReadonlyJSONValue } from "../../utils";
 import { TimingTracker } from "./TimingTracker";
 
@@ -37,20 +37,29 @@ export const createInitialMessage = ({
   },
 });
 
+type WarnOnce = (key: string, message: string) => void;
+
+const MAX_WARNED_KEYS = 16;
+
 const updatePartForPath = (
   message: AssistantMessage,
   chunk: AssistantStreamChunk,
+  warnOnce: WarnOnce,
   updater: (part: AssistantMessagePart) => AssistantMessagePart,
 ): AssistantMessage => {
-  if (message.parts.length === 0) {
-    throw new Error("No parts available to update.");
+  const part =
+    chunk.path.length === 1 ? message.parts[chunk.path[0]!] : undefined;
+  if (part === undefined) {
+    warnOnce(
+      `no-part:${chunk.type}`,
+      `Dropped ${chunk.type} chunk: no part at path [${chunk.path.join(", ")}]`,
+    );
+    return message;
   }
 
-  if (chunk.path.length !== 1)
-    throw new Error("Nested paths are not supported yet.");
-
   const partIndex = chunk.path[0]!;
-  const updatedPart = updater(message.parts[partIndex]!);
+  const updatedPart = updater(part);
+  if (updatedPart === part) return message;
   return {
     ...message,
     parts: [
@@ -67,6 +76,7 @@ const updatePartForPath = (
 const handlePartStart = (
   message: AssistantMessage,
   chunk: AssistantStreamChunk & { readonly type: "part-start" },
+  warnOnce: WarnOnce,
 ): AssistantMessage => {
   const partInit = chunk.part;
   if (partInit.type === "text" || partInit.type === "reasoning") {
@@ -147,7 +157,24 @@ const handlePartStart = (
       },
     };
   } else {
-    throw new Error(`Unsupported part type: ${partInit.type}`);
+    const unsupportedType = String((partInit as { type?: unknown }).type);
+    warnOnce(
+      `unsupported-part-start:${unsupportedType}`,
+      `Unsupported part-start type ${unsupportedType}: inserting an empty reasoning part to preserve part indices`,
+    );
+    // reasoning rather than a data sentinel: auiV0Encode rejects data parts, so a data placeholder would fail cloud persistence
+    const placeholderPart: ReasoningPart = {
+      type: "reasoning",
+      text: "",
+      status: { type: "running" },
+    };
+    return {
+      ...message,
+      parts: [...message.parts, placeholderPart],
+      get content() {
+        return this.parts;
+      },
+    };
   }
 };
 
@@ -156,14 +183,19 @@ const handleToolCallArgsTextFinish = (
   chunk: AssistantStreamChunk & {
     readonly type: "tool-call-args-text-finish";
   },
+  warnOnce: WarnOnce,
 ): AssistantMessage => {
-  return updatePartForPath(message, chunk, (part) => {
+  return updatePartForPath(message, chunk, warnOnce, (part) => {
     if (part.type !== "tool-call") {
-      throw new Error("Last is not a tool call");
+      warnOnce(
+        "wrong-part:tool-call-args-text-finish",
+        "Dropped tool-call-args-text-finish chunk: part is not a tool-call",
+      );
+      return part;
     }
 
     // TODO this should never be hit; this happens if args-text-finish is emitted after result
-    if (part.state !== "partial-call") return part;
+    if (part.state !== "partial-call") return { ...part };
     // throw new Error("Last is not a partial call");
 
     return {
@@ -176,8 +208,9 @@ const handleToolCallArgsTextFinish = (
 const handlePartFinish = (
   message: AssistantMessage,
   chunk: AssistantStreamChunk & { readonly type: "part-finish" },
+  warnOnce: WarnOnce,
 ): AssistantMessage => {
-  return updatePartForPath(message, chunk, (part) => ({
+  return updatePartForPath(message, chunk, warnOnce, (part) => ({
     ...part,
     status: { type: "complete", reason: "unknown" },
   }));
@@ -186,8 +219,9 @@ const handlePartFinish = (
 const handleTextDelta = (
   message: AssistantMessage,
   chunk: AssistantStreamChunk & { type: "text-delta" },
+  warnOnce: WarnOnce,
 ): AssistantMessage => {
-  return updatePartForPath(message, chunk, (part) => {
+  return updatePartForPath(message, chunk, warnOnce, (part) => {
     if (part.type === "text" || part.type === "reasoning") {
       return { ...part, text: part.text + chunk.textDelta };
     } else if (part.type === "tool-call") {
@@ -198,9 +232,11 @@ const handleTextDelta = (
 
       return { ...part, argsText: newArgsText, args: newArgs };
     } else {
-      throw new Error(
-        "text-delta received but part is neither text nor tool-call",
+      warnOnce(
+        "wrong-part:text-delta",
+        "Dropped text-delta chunk: part is neither text nor tool-call",
       );
+      return part;
     }
   });
 };
@@ -208,8 +244,9 @@ const handleTextDelta = (
 const handleResult = (
   message: AssistantMessage,
   chunk: AssistantStreamChunk & { type: "result" },
+  warnOnce: WarnOnce,
 ): AssistantMessage => {
-  return updatePartForPath(message, chunk, (part) => {
+  return updatePartForPath(message, chunk, warnOnce, (part) => {
     if (part.type === "tool-call") {
       return {
         ...part,
@@ -232,7 +269,11 @@ const handleResult = (
         status: { type: "complete", reason: "stop" },
       };
     } else {
-      throw new Error("Result chunk received but part is not a tool-call");
+      warnOnce(
+        "wrong-part:result",
+        "Dropped result chunk: part is not a tool-call",
+      );
+      return part;
     }
   });
 };
@@ -364,9 +405,23 @@ const handleErrorChunk = (
   message: AssistantMessage,
   chunk: AssistantStreamChunk & { type: "error" },
 ): AssistantMessage => {
+  const severity =
+    chunk.severity === "critical" ||
+    chunk.severity === "warning" ||
+    chunk.severity === "info"
+      ? chunk.severity
+      : undefined;
   return {
     ...message,
-    status: { type: "incomplete", reason: "error", error: chunk.error },
+    status: {
+      type: "incomplete",
+      reason: "error",
+      error: {
+        code: chunk.code ?? "unknown",
+        message: chunk.error ?? "unknown error",
+        ...(severity !== undefined && { severity }),
+      },
+    },
   };
 };
 
@@ -374,7 +429,7 @@ const handleUpdateState = (
   message: AssistantMessage,
   chunk: AssistantStreamChunk & { type: "update-state" },
 ): AssistantMessage => {
-  const acc = new ObjectStreamAccumulator(message.metadata.unstable_state);
+  const acc = new GorpStreamAccumulator(message.metadata.unstable_state);
   acc.append(chunk.operations);
 
   return {
@@ -437,6 +492,12 @@ export class AssistantMessageAccumulator extends TransformStream<
   } = {}) {
     let message = initialMessage ?? createInitialMessage();
     const tracker = new TimingTracker();
+    const warnedKeys = new Set<string>();
+    const warnOnce: WarnOnce = (key, warning) => {
+      if (warnedKeys.has(key) || warnedKeys.size >= MAX_WARNED_KEYS) return;
+      warnedKeys.add(key);
+      console.warn(warning);
+    };
     let controller:
       | TransformStreamDefaultController<AssistantMessage>
       | undefined;
@@ -456,26 +517,28 @@ export class AssistantMessageAccumulator extends TransformStream<
         const type = chunk.type;
         switch (type) {
           case "part-start":
-            message = handlePartStart(message, chunk);
+            message = handlePartStart(message, chunk, warnOnce);
             if (chunk.part.type === "tool-call") {
               tracker.recordToolCallStart(chunk.part.toolCallId);
             }
             break;
 
           case "tool-call-args-text-finish":
-            message = handleToolCallArgsTextFinish(message, chunk);
+            message = handleToolCallArgsTextFinish(message, chunk, warnOnce);
             break;
 
           case "part-finish":
-            message = handlePartFinish(message, chunk);
+            message = handlePartFinish(message, chunk, warnOnce);
             break;
 
-          case "text-delta":
-            message = handleTextDelta(message, chunk);
-            tracker.recordFirstToken();
+          case "text-delta": {
+            const next = handleTextDelta(message, chunk, warnOnce);
+            if (next !== message) tracker.recordFirstToken();
+            message = next;
             break;
+          }
           case "result":
-            message = handleResult(message, chunk);
+            message = handleResult(message, chunk, warnOnce);
             break;
           case "message-finish":
             message = handleMessageFinish(message, chunk);