assistant-stream 0.3.26 → 0.3.27

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;
     }
   });
 };
@@ -377,7 +418,7 @@ const handleErrorChunk = (
       reason: "error",
       error: {
         code: chunk.code ?? "unknown",
-        message: chunk.error,
+        message: chunk.error ?? "unknown error",
         ...(severity !== undefined && { severity }),
       },
     },
@@ -388,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 {
@@ -451,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;
@@ -470,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);

package/src/core/{object/ObjectStream.test.ts → gorp/GorpStream.test.ts}

@@ -1,11 +1,8 @@
 import { describe, expect, it } from "vitest";
-import { createObjectStream } from "./createObjectStream";
-import {
-  ObjectStreamEncoder,
-  ObjectStreamDecoder,
-} from "./ObjectStreamResponse";
+import { createGorpStream } from "./createGorpStream";
+import { GorpStreamEncoder, GorpStreamDecoder } from "./GorpStreamResponse";
 import type { ReadonlyJSONValue } from "../../utils";
-import type { ObjectStreamChunk } from "./types";
+import type { GorpStreamChunk } from "./types";
 
 // Helper function to collect all chunks from a stream
 async function collectChunks<T>(stream: ReadableStream<T>): Promise<T[]> {
@@ -27,10 +24,10 @@ async function collectChunks<T>(stream: ReadableStream<T>): Promise<T[]> {
 
 // Helper function to encode and decode a stream
 async function encodeAndDecode(
-  stream: ReadableStream<ObjectStreamChunk>,
-): Promise<ReadableStream<ObjectStreamChunk>> {
+  stream: ReadableStream<GorpStreamChunk>,
+): Promise<ReadableStream<GorpStreamChunk>> {
   // Encode the stream to Uint8Array (simulating network transmission)
-  const encodedStream = stream.pipeThrough(new ObjectStreamEncoder());
+  const encodedStream = stream.pipeThrough(new GorpStreamEncoder());
 
   // Collect all encoded chunks
   const encodedChunks = await collectChunks(encodedStream);
@@ -45,13 +42,37 @@ async function encodeAndDecode(
     },
   });
 
-  // Decode the stream back to ObjectStreamChunk
-  return reconstructedStream.pipeThrough(new ObjectStreamDecoder());
+  // Decode the stream back to GorpStreamChunk
+  return reconstructedStream.pipeThrough(new GorpStreamDecoder());
 }
 
-describe("ObjectStream serialization and deserialization", () => {
+describe("GorpStream serialization and deserialization", () => {
+  it("does not settle the stream after cancellation", async () => {
+    let finishExecution!: () => void;
+    let abortSignal!: AbortSignal;
+    const execution = new Promise<void>((resolve) => {
+      finishExecution = resolve;
+    });
+    const stream = createGorpStream({
+      execute: (controller) => {
+        abortSignal = controller.abortSignal;
+        return execution;
+      },
+    });
+    const reader = stream.getReader();
+
+    await reader.cancel("consumer stopped");
+
+    expect(abortSignal.aborted).toBe(true);
+    expect(abortSignal.reason).toBe("consumer stopped");
+
+    finishExecution();
+    await execution;
+    await new Promise((resolve) => setTimeout(resolve, 0));
+  });
+
   it("should discard an unterminated SSE event", async () => {
-    const operations: ObjectStreamChunk["operations"] = [
+    const operations: GorpStreamChunk["operations"] = [
       { type: "set", path: ["status"], value: "complete" },
     ];
     const stream = new ReadableStream<Uint8Array>({
@@ -64,7 +85,7 @@ describe("ObjectStream serialization and deserialization", () => {
     });
 
     const chunks = await collectChunks(
-      stream.pipeThrough(new ObjectStreamDecoder()),
+      stream.pipeThrough(new GorpStreamDecoder()),
     );
 
     expect(chunks).toEqual([]);
@@ -72,7 +93,7 @@ describe("ObjectStream serialization and deserialization", () => {
 
   it("should correctly serialize and deserialize simple objects", async () => {
     // Create an object stream with simple operations
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       execute: (controller) => {
         controller.enqueue([
           { type: "set", path: ["name"], value: "John" },
@@ -96,7 +117,7 @@ describe("ObjectStream serialization and deserialization", () => {
   });
 
   it("should correctly handle nested objects", async () => {
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       execute: (controller) => {
         controller.enqueue([
           { type: "set", path: ["user", "profile", "name"], value: "Jane" },
@@ -128,7 +149,7 @@ describe("ObjectStream serialization and deserialization", () => {
   });
 
   it("should correctly handle arrays", async () => {
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       execute: (controller) => {
         controller.enqueue([
           { type: "set", path: ["items"], value: [] },
@@ -149,7 +170,7 @@ describe("ObjectStream serialization and deserialization", () => {
   });
 
   it("should correctly handle mixed arrays and objects", async () => {
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       execute: (controller) => {
         controller.enqueue([
           { type: "set", path: ["users"], value: [] },
@@ -176,7 +197,7 @@ describe("ObjectStream serialization and deserialization", () => {
   });
 
   it("should correctly handle append-text operations", async () => {
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       execute: (controller) => {
         controller.enqueue([
           { type: "set", path: ["message"], value: "Hello" },
@@ -196,8 +217,48 @@ describe("ObjectStream serialization and deserialization", () => {
     });
   });
 
+  it("does not apply first-chunk operations twice", async () => {
+    const stream = createGorpStream({
+      defaultValue: { message: "Hello" },
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "append-text", path: ["message"], value: " world" },
+        ]);
+        controller.enqueue([
+          { type: "append-text", path: ["message"], value: "!" },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+
+    expect(chunks).toEqual([
+      {
+        snapshot: { message: "Hello world" },
+        operations: [
+          {
+            type: "set",
+            path: [],
+            value: { message: "Hello world" },
+          },
+        ],
+      },
+      {
+        snapshot: { message: "Hello world!" },
+        operations: [
+          {
+            type: "append-text",
+            path: ["message"],
+            value: "!",
+          },
+        ],
+      },
+    ]);
+  });
+
   it("should correctly handle special characters and Unicode", async () => {
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       execute: (controller) => {
         controller.enqueue([
           {
@@ -229,7 +290,7 @@ describe("ObjectStream serialization and deserialization", () => {
   });
 
   it("should correctly handle null and undefined values", async () => {
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       execute: (controller) => {
         controller.enqueue([
           { type: "set", path: ["nullValue"], value: null },
@@ -252,7 +313,7 @@ describe("ObjectStream serialization and deserialization", () => {
 
   it("should correctly handle large nested structures", async () => {
     // Create a deep nested structure
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       execute: (controller) => {
         controller.enqueue([
           { type: "set", path: ["level1"], value: {} },
@@ -290,7 +351,7 @@ describe("ObjectStream serialization and deserialization", () => {
   });
 
   it("should correctly handle operations in multiple enqueue calls", async () => {
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       execute: (controller) => {
         // First batch of operations
         controller.enqueue([
@@ -339,7 +400,7 @@ describe("ObjectStream serialization and deserialization", () => {
   });
 
   it("should correctly handle overwriting existing values", async () => {
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       execute: (controller) => {
         controller.enqueue([
           { type: "set", path: ["value"], value: "initial" },
@@ -371,7 +432,7 @@ describe("ObjectStream serialization and deserialization", () => {
       },
     };
 
-    const stream = createObjectStream({
+    const stream = createGorpStream({
       defaultValue: initialValue,
       execute: (controller) => {
         controller.enqueue([

package/src/core/gorp/GorpStreamAccumulator.test.ts

@@ -0,0 +1,28 @@
+import { describe, expect, it } from "vitest";
+import { GorpStreamAccumulator } from "./GorpStreamAccumulator";
+
+describe("GorpStreamAccumulator", () => {
+  it("rejects unsafe path segments", () => {
+    for (const path of [
+      ["__proto__", "polluted"],
+      ["constructor", "prototype", "polluted"],
+      ["a", "__proto__"],
+    ]) {
+      const acc = new GorpStreamAccumulator({});
+      expect(() => acc.append([{ type: "set", path, value: true }])).toThrow(
+        /Unsafe gorp path segment/,
+      );
+    }
+    expect(({} as { polluted?: unknown }).polluted).toBeUndefined();
+  });
+
+  it("treats inherited keys as absent when navigating", () => {
+    const acc = new GorpStreamAccumulator({});
+    acc.append([{ type: "set", path: ["toString", "x"], value: 1 }]);
+    expect(acc.state).toEqual({ toString: { x: 1 } });
+
+    const append = new GorpStreamAccumulator({ toString: "hi" });
+    append.append([{ type: "append-text", path: ["toString"], value: "!" }]);
+    expect(append.state).toEqual({ toString: "hi!" });
+  });
+});

package/src/core/{object/ObjectStreamAccumulator.ts → gorp/GorpStreamAccumulator.ts}

@@ -1,7 +1,8 @@
 import type { ReadonlyJSONValue, ReadonlyJSONObject } from "../../utils";
-import type { ObjectStreamOperation } from "./types";
+import { assertSafePathSegment } from "./changeTree";
+import type { GorpStreamOperation } from "./types";
 
-export class ObjectStreamAccumulator {
+export class GorpStreamAccumulator {
   private _state: ReadonlyJSONValue;
 
   constructor(initialValue: ReadonlyJSONValue = null) {
@@ -12,24 +13,20 @@ export class ObjectStreamAccumulator {
     return this._state;
   }
 
-  append(ops: readonly ObjectStreamOperation[]) {
+  append(ops: readonly GorpStreamOperation[]) {
     this._state = ops.reduce(
-      (state, op) => ObjectStreamAccumulator.apply(state, op),
+      (state, op) => GorpStreamAccumulator.apply(state, op),
       this._state,
     );
   }
 
-  private static apply(state: ReadonlyJSONValue, op: ObjectStreamOperation) {
+  private static apply(state: ReadonlyJSONValue, op: GorpStreamOperation) {
     const type = op.type;
     switch (type) {
       case "set":
-        return ObjectStreamAccumulator.updatePath(
-          state,
-          op.path,
-          () => op.value,
-        );
+        return GorpStreamAccumulator.updatePath(state, op.path, () => op.value);
       case "append-text":
-        return ObjectStreamAccumulator.updatePath(state, op.path, (current) => {
+        return GorpStreamAccumulator.updatePath(state, op.path, (current) => {
           if (typeof current !== "string")
             throw new Error(`Expected string at path [${op.path.join(", ")}]`);
           return current + op.value;
@@ -57,6 +54,7 @@ export class ObjectStreamAccumulator {
     }
 
     const [key, ...rest] = path as [string, ...(readonly string[])];
+    assertSafePathSegment(key);
     if (Array.isArray(state)) {
       const idx = Number(key);
       if (Number.isNaN(idx))
@@ -65,8 +63,8 @@ export class ObjectStreamAccumulator {
         throw new Error(`Insert array index out of bounds`);
 
       const nextState = [...state];
-      nextState[idx] = ObjectStreamAccumulator.updatePath(
-        nextState[idx],
+      nextState[idx] = GorpStreamAccumulator.updatePath(
+        Object.hasOwn(nextState, idx) ? nextState[idx] : undefined,
         rest,
         updater,
       );
@@ -75,8 +73,8 @@ export class ObjectStreamAccumulator {
     }
 
     const nextState = { ...(state as ReadonlyJSONObject) };
-    nextState[key] = ObjectStreamAccumulator.updatePath(
-      nextState[key],
+    nextState[key] = GorpStreamAccumulator.updatePath(
+      Object.hasOwn(nextState, key) ? nextState[key] : undefined,
       rest,
       updater,
     );