assistant-stream 0.2.5 → 0.2.7

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

@@ -0,0 +1,376 @@
+import { describe, expect, it } from "vitest";
+import { createObjectStream } from "./createObjectStream";
+import {
+  ObjectStreamEncoder,
+  ObjectStreamDecoder,
+} from "./ObjectStreamResponse";
+import { ReadonlyJSONValue } from "../../utils";
+import { ObjectStreamChunk } from "./types";
+
+// 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 encode and decode a stream
+async function encodeAndDecode(
+  stream: ReadableStream<ObjectStreamChunk>,
+): Promise<ReadableStream<ObjectStreamChunk>> {
+  // Encode the stream to Uint8Array (simulating network transmission)
+  const encodedStream = stream.pipeThrough(new ObjectStreamEncoder());
+
+  // Collect all encoded chunks
+  const encodedChunks = await collectChunks(encodedStream);
+
+  // Create a new stream from the encoded chunks
+  const reconstructedStream = new ReadableStream<Uint8Array>({
+    start(controller) {
+      for (const chunk of encodedChunks) {
+        controller.enqueue(chunk);
+      }
+      controller.close();
+    },
+  });
+
+  // Decode the stream back to ObjectStreamChunk
+  return reconstructedStream.pipeThrough(new ObjectStreamDecoder());
+}
+
+describe("ObjectStream serialization and deserialization", () => {
+  it("should correctly serialize and deserialize simple objects", async () => {
+    // Create an object stream with simple operations
+    const stream = createObjectStream({
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "set", path: ["name"], value: "John" },
+          { type: "set", path: ["age"], value: 30 },
+        ]);
+      },
+    });
+
+    // Encode and decode the stream
+    const decodedStream = await encodeAndDecode(stream);
+
+    // Collect all chunks from the decoded stream
+    const chunks = await collectChunks(decodedStream);
+
+    // Verify the final state
+    const finalChunk = chunks[chunks.length - 1]!;
+    expect(finalChunk.snapshot).toEqual({
+      name: "John",
+      age: 30,
+    });
+  });
+
+  it("should correctly handle nested objects", async () => {
+    const stream = createObjectStream({
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "set", path: ["user", "profile", "name"], value: "Jane" },
+          {
+            type: "set",
+            path: ["user", "profile", "email"],
+            value: "jane@example.com",
+          },
+          { type: "set", path: ["user", "settings", "theme"], value: "dark" },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+    const finalChunk = chunks[chunks.length - 1]!;
+
+    expect(finalChunk.snapshot).toEqual({
+      user: {
+        profile: {
+          name: "Jane",
+          email: "jane@example.com",
+        },
+        settings: {
+          theme: "dark",
+        },
+      },
+    });
+  });
+
+  it("should correctly handle arrays", async () => {
+    const stream = createObjectStream({
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "set", path: ["items"], value: [] },
+          { type: "set", path: ["items", "0"], value: "apple" },
+          { type: "set", path: ["items", "1"], value: "banana" },
+          { type: "set", path: ["items", "2"], value: "cherry" },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+    const finalChunk = chunks[chunks.length - 1]!;
+
+    expect(finalChunk.snapshot).toEqual({
+      items: ["apple", "banana", "cherry"],
+    });
+  });
+
+  it("should correctly handle mixed arrays and objects", async () => {
+    const stream = createObjectStream({
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "set", path: ["users"], value: [] },
+          { type: "set", path: ["users", "0"], value: {} },
+          { type: "set", path: ["users", "0", "id"], value: 1 },
+          { type: "set", path: ["users", "0", "name"], value: "Alice" },
+          { type: "set", path: ["users", "1"], value: {} },
+          { type: "set", path: ["users", "1", "id"], value: 2 },
+          { type: "set", path: ["users", "1", "name"], value: "Bob" },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+    const finalChunk = chunks[chunks.length - 1]!;
+
+    expect(finalChunk.snapshot).toEqual({
+      users: [
+        { id: 1, name: "Alice" },
+        { id: 2, name: "Bob" },
+      ],
+    });
+  });
+
+  it("should correctly handle append-text operations", async () => {
+    const stream = createObjectStream({
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "set", path: ["message"], value: "Hello" },
+          { type: "append-text", path: ["message"], value: " " },
+          { type: "append-text", path: ["message"], value: "World" },
+          { type: "append-text", path: ["message"], value: "!" },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+    const finalChunk = chunks[chunks.length - 1]!;
+
+    expect(finalChunk.snapshot).toEqual({
+      message: "Hello World!",
+    });
+  });
+
+  it("should correctly handle special characters and Unicode", async () => {
+    const stream = createObjectStream({
+      execute: (controller) => {
+        controller.enqueue([
+          {
+            type: "set",
+            path: ["special"],
+            value: "Special chars: !@#$%^&*()",
+          },
+          { type: "set", path: ["unicode"], value: "Unicode: 😀🌍🚀" },
+          { type: "set", path: ["quotes"], value: "Quotes: \"'`" },
+          {
+            type: "set",
+            path: ["newlines"],
+            value: "Line 1\nLine 2\r\nLine 3",
+          },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+    const finalChunk = chunks[chunks.length - 1]!;
+
+    expect(finalChunk.snapshot).toEqual({
+      special: "Special chars: !@#$%^&*()",
+      unicode: "Unicode: 😀🌍🚀",
+      quotes: "Quotes: \"'`",
+      newlines: "Line 1\nLine 2\r\nLine 3",
+    });
+  });
+
+  it("should correctly handle null and undefined values", async () => {
+    const stream = createObjectStream({
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "set", path: ["nullValue"], value: null },
+          { type: "set", path: ["emptyObject"], value: {} },
+          { type: "set", path: ["emptyArray"], value: [] },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+    const finalChunk = chunks[chunks.length - 1]!;
+
+    expect(finalChunk.snapshot).toEqual({
+      nullValue: null,
+      emptyObject: {},
+      emptyArray: [],
+    });
+  });
+
+  it("should correctly handle large nested structures", async () => {
+    // Create a deep nested structure
+    const stream = createObjectStream({
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "set", path: ["level1"], value: {} },
+          { type: "set", path: ["level1", "level2"], value: {} },
+          { type: "set", path: ["level1", "level2", "level3"], value: {} },
+          {
+            type: "set",
+            path: ["level1", "level2", "level3", "level4"],
+            value: {},
+          },
+          {
+            type: "set",
+            path: ["level1", "level2", "level3", "level4", "level5"],
+            value: "deep value",
+          },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+    const finalChunk = chunks[chunks.length - 1]!;
+
+    expect(finalChunk.snapshot).toEqual({
+      level1: {
+        level2: {
+          level3: {
+            level4: {
+              level5: "deep value",
+            },
+          },
+        },
+      },
+    });
+  });
+
+  it("should correctly handle operations in multiple enqueue calls", async () => {
+    const stream = createObjectStream({
+      execute: (controller) => {
+        // First batch of operations
+        controller.enqueue([
+          { type: "set", path: ["user"], value: { name: "Initial" } },
+        ]);
+
+        // Second batch of operations
+        controller.enqueue([
+          { type: "set", path: ["user", "name"], value: "Updated" },
+          { type: "set", path: ["user", "email"], value: "user@example.com" },
+        ]);
+
+        // Third batch of operations
+        controller.enqueue([
+          { type: "set", path: ["status"], value: "complete" },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+    const finalChunk = chunks[chunks.length - 1]!;
+
+    expect(finalChunk.snapshot).toEqual({
+      user: {
+        name: "Updated",
+        email: "user@example.com",
+      },
+      status: "complete",
+    });
+
+    // Verify that we got the correct number of chunks
+    expect(chunks.length).toBe(3);
+
+    // Verify intermediate states
+    expect(chunks[0]!.snapshot).toEqual({
+      user: { name: "Initial" },
+    });
+
+    expect(chunks[1]!.snapshot).toEqual({
+      user: {
+        name: "Updated",
+        email: "user@example.com",
+      },
+    });
+  });
+
+  it("should correctly handle overwriting existing values", async () => {
+    const stream = createObjectStream({
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "set", path: ["value"], value: "initial" },
+          { type: "set", path: ["nested"], value: { prop: "initial" } },
+        ]);
+
+        controller.enqueue([
+          { type: "set", path: ["value"], value: "updated" },
+          { type: "set", path: ["nested"], value: "completely replaced" },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+    const finalChunk = chunks[chunks.length - 1]!;
+
+    expect(finalChunk.snapshot).toEqual({
+      value: "updated",
+      nested: "completely replaced",
+    });
+  });
+
+  it("should correctly handle custom initial values", async () => {
+    const initialValue: ReadonlyJSONValue = {
+      existing: "value",
+      nested: {
+        prop: 123,
+      },
+    };
+
+    const stream = createObjectStream({
+      defaultValue: initialValue,
+      execute: (controller) => {
+        controller.enqueue([
+          { type: "set", path: ["new"], value: "added" },
+          { type: "set", path: ["nested", "prop"], value: 456 },
+        ]);
+      },
+    });
+
+    const decodedStream = await encodeAndDecode(stream);
+    const chunks = await collectChunks(decodedStream);
+    const finalChunk = chunks[chunks.length - 1]!;
+
+    expect(finalChunk.snapshot).toEqual({
+      existing: "value",
+      nested: {
+        prop: 456,
+      },
+      new: "added",
+    });
+  });
+});

package/src/core/object/ObjectStreamAccumulator.ts

@@ -0,0 +1,80 @@
+import { ReadonlyJSONValue, ReadonlyJSONObject } from "../../utils";
+import { ObjectStreamOperation } from "./types";
+
+export class ObjectStreamAccumulator {
+  private _state: ReadonlyJSONValue;
+
+  constructor(initialValue: ReadonlyJSONValue = null) {
+    this._state = initialValue;
+  }
+
+  get state() {
+    return this._state;
+  }
+
+  append(ops: readonly ObjectStreamOperation[]) {
+    this._state = ops.reduce(
+      (state, op) => ObjectStreamAccumulator.apply(state, op),
+      this._state,
+    );
+  }
+
+  private static apply(state: ReadonlyJSONValue, op: ObjectStreamOperation) {
+    const type = op.type;
+    switch (type) {
+      case "set":
+        return ObjectStreamAccumulator.updatePath(
+          state,
+          op.path,
+          () => op.value,
+        );
+      case "append-text":
+        return ObjectStreamAccumulator.updatePath(state, op.path, (current) => {
+          if (typeof current !== "string")
+            throw new Error(`Expected string at path [${op.path.join(", ")}]`);
+          return current + op.value;
+        });
+
+      default: {
+        const _exhaustiveCheck: never = type;
+        throw new Error(`Invalid operation type: ${_exhaustiveCheck}`);
+      }
+    }
+  }
+
+  private static updatePath(
+    state: ReadonlyJSONValue | undefined,
+    path: readonly string[],
+    updater: (current: ReadonlyJSONValue | undefined) => ReadonlyJSONValue,
+  ): ReadonlyJSONValue {
+    if (path.length === 0) return updater(state);
+
+    // Initialize state as empty object if it's null and we're trying to set a property
+    if (state === null) {
+      state = {};
+    }
+
+    if (typeof state !== "object") {
+      throw new Error(`Invalid path: [${path.join(", ")}]`);
+    }
+
+    const [key, ...rest] = path as [string, ...(readonly string[])];
+    if (Array.isArray(state)) {
+      const idx = Number(key);
+      if (isNaN(idx))
+        throw new Error(`Expected array index at [${path.join(", ")}]`);
+      if (idx > state.length || idx < 0)
+        throw new Error(`Insert array index out of bounds`);
+
+      const nextState = [...state];
+      nextState[idx] = this.updatePath(nextState[idx], rest, updater);
+
+      return nextState;
+    }
+
+    const nextState = { ...(state as ReadonlyJSONObject) };
+    nextState[key] = this.updatePath(nextState[key], rest, updater);
+
+    return nextState;
+  }
+}

package/src/core/object/ObjectStreamResponse.ts

@@ -0,0 +1,81 @@
+import { PipeableTransformStream } from "../utils/stream/PipeableTransformStream";
+import { ObjectStreamAccumulator } from "./ObjectStreamAccumulator";
+import { SSEDecoder, SSEEncoder } from "../utils/stream/SSE";
+import { ObjectStreamChunk, ObjectStreamOperation } from "./types";
+
+export class ObjectStreamEncoder extends PipeableTransformStream<
+  ObjectStreamChunk,
+  Uint8Array
+> {
+  constructor() {
+    super((readable) =>
+      readable
+        .pipeThrough(
+          new TransformStream<
+            ObjectStreamChunk,
+            readonly ObjectStreamOperation[]
+          >({
+            transform(chunk, controller) {
+              controller.enqueue(chunk.operations);
+            },
+          }),
+        )
+        .pipeThrough(new SSEEncoder()),
+    );
+  }
+}
+
+export class ObjectStreamDecoder extends PipeableTransformStream<
+  Uint8Array,
+  ObjectStreamChunk
+> {
+  constructor() {
+    const accumulator = new ObjectStreamAccumulator();
+    super((readable) =>
+      readable
+        .pipeThrough(new SSEDecoder<readonly ObjectStreamOperation[]>())
+        .pipeThrough(
+          new TransformStream<
+            readonly ObjectStreamOperation[],
+            ObjectStreamChunk
+          >({
+            transform(operations, controller) {
+              accumulator.append(operations);
+              controller.enqueue({
+                snapshot: accumulator.state,
+                operations,
+              });
+            },
+          }),
+        ),
+    );
+  }
+}
+
+export class ObjectStreamResponse extends Response {
+  constructor(body: ReadableStream<ObjectStreamChunk>) {
+    super(body.pipeThrough(new ObjectStreamEncoder()), {
+      headers: new Headers({
+        "Content-Type": "text/event-stream",
+        "Cache-Control": "no-cache",
+        Connection: "keep-alive",
+        "Assistant-Stream-Format": "object-stream/v0",
+      }),
+    });
+  }
+}
+
+export const fromObjectStreamResponse = (
+  response: Response,
+): ReadableStream<ObjectStreamChunk> => {
+  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") {
+    throw new Error("Response is not an event stream");
+  }
+  if (response.headers.get("Assistant-Stream-Format") !== "object-stream/v0") {
+    throw new Error("Unsupported Assistant-Stream-Format header");
+  }
+  return response.body.pipeThrough(new ObjectStreamDecoder());
+};

package/src/core/object/createObjectStream.ts

@@ -0,0 +1,87 @@
+import { ReadonlyJSONValue } from "../../utils";
+import { withPromiseOrValue } from "../utils/withPromiseOrValue";
+import { ObjectStreamAccumulator } from "./ObjectStreamAccumulator";
+import { ObjectStreamOperation, ObjectStreamChunk } from "./types";
+
+type ObjectStreamController = {
+  readonly abortSignal: AbortSignal;
+
+  enqueue(operations: readonly ObjectStreamOperation[]): void;
+};
+
+class ObjectStreamControllerImpl implements ObjectStreamController {
+  private _controller: ReadableStreamDefaultController<ObjectStreamChunk>;
+  private _abortController = new AbortController();
+  private _accumulator: ObjectStreamAccumulator;
+
+  get abortSignal() {
+    return this._abortController.signal;
+  }
+
+  constructor(
+    controller: ReadableStreamDefaultController<ObjectStreamChunk>,
+    defaultValue: ReadonlyJSONValue,
+  ) {
+    this._controller = controller;
+    this._accumulator = new ObjectStreamAccumulator(defaultValue);
+  }
+
+  enqueue(operations: readonly ObjectStreamOperation[]) {
+    this._accumulator.append(operations);
+
+    this._controller.enqueue({
+      snapshot: this._accumulator.state,
+      operations,
+    });
+  }
+
+  __internalError(error: unknown) {
+    this._controller.error(error);
+  }
+
+  __internalClose() {
+    this._controller.close();
+  }
+
+  __internalCancel(reason?: unknown) {
+    this._abortController.abort(reason);
+  }
+}
+
+const getStreamControllerPair = (defaultValue: ReadonlyJSONValue) => {
+  let controller!: ObjectStreamControllerImpl;
+  const stream = new ReadableStream<ObjectStreamChunk>({
+    start(c) {
+      controller = new ObjectStreamControllerImpl(c, defaultValue);
+    },
+    cancel(reason: unknown) {
+      controller.__internalCancel(reason);
+    },
+  });
+
+  return [stream, controller] as const;
+};
+
+type CreateObjectStreamOptions = {
+  execute: (controller: ObjectStreamController) => void | PromiseLike<void>;
+  defaultValue?: ReadonlyJSONValue;
+};
+
+export const createObjectStream = ({
+  execute,
+  defaultValue = {},
+}: CreateObjectStreamOptions) => {
+  const [stream, controller] = getStreamControllerPair(defaultValue);
+
+  withPromiseOrValue(
+    () => execute(controller),
+    () => {
+      controller.__internalClose();
+    },
+    (e: unknown) => {
+      controller.__internalError(e);
+    },
+  );
+
+  return stream;
+};

package/src/core/object/types.ts

@@ -0,0 +1,18 @@
+import { ReadonlyJSONValue } from "../../utils";
+
+export type ObjectStreamOperation =
+  | {
+      readonly type: "set";
+      readonly path: readonly string[];
+      readonly value: ReadonlyJSONValue;
+    }
+  | {
+      readonly type: "append-text";
+      readonly path: readonly string[];
+      readonly value: string;
+    };
+
+export type ObjectStreamChunk = {
+  readonly snapshot: ReadonlyJSONValue;
+  readonly operations: readonly ObjectStreamOperation[];
+};

package/src/core/serialization/PlainText.ts

@@ -35,7 +35,8 @@ export class PlainTextEncoder
                 | "tool-call-begin"
                 | "tool-call-delta"
                 | "result"
-                | "error" = type;
+                | "error"
+                | "update-state" = type;
               throw new Error(`unsupported chunk type: ${unsupportedType}`);
           }
         },

package/src/core/serialization/data-stream/DataStream.ts

@@ -150,6 +150,14 @@ export class DataStreamEncoder
               break;
             }
 
+            case "update-state": {
+              controller.enqueue({
+                type: DataStreamStreamChunkType.AuiUpdateStateOperations,
+                value: chunk.operations,
+              });
+              break;
+            }
+
             // TODO ignore for now
             // in the future, we should create a handler that waits for text parts to finish before continuing
             case "tool-call-args-text-finish":
@@ -326,6 +334,14 @@ export class DataStreamDecoder extends PipeableTransformStream<
               });
               break;
 
+            case DataStreamStreamChunkType.AuiUpdateStateOperations:
+              controller.enqueue({
+                type: "update-state",
+                path: [],
+                operations: value,
+              });
+              break;
+
             case DataStreamStreamChunkType.ReasoningSignature:
             case DataStreamStreamChunkType.RedactedReasoning:
               // ignore these for now