assistant-stream 0.2.6 → 0.2.7

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

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

@@ -2,6 +2,7 @@ import {
   ReadonlyJSONObject,
   ReadonlyJSONValue,
 } from "../../../utils/json/json-value";
+import { ObjectStreamOperation } from "../../object/types";
 
 export type DataStreamChunk = {
   [K in DataStreamStreamChunkType]: {
@@ -41,6 +42,8 @@ export enum DataStreamStreamChunkType {
   RedactedReasoning = "i",
   ReasoningSignature = "j",
   File = "k",
+
+  AuiUpdateStateOperations = "aui-state",
 }
 type DataStreamStreamChunkValue = {
   [DataStreamStreamChunkType.TextDelta]: string;
@@ -90,4 +93,7 @@ type DataStreamStreamChunkValue = {
   [DataStreamStreamChunkType.RedactedReasoning]: { data: string };
   [DataStreamStreamChunkType.ReasoningSignature]: { signature: string };
   [DataStreamStreamChunkType.File]: { data: string; mimeType: string };
+
+  // aui-extensions
+  [DataStreamStreamChunkType.AuiUpdateStateOperations]: ObjectStreamOperation[];
 };

package/src/core/tool/ToolCallReader.ts

@@ -329,7 +329,6 @@ export class ToolCallArgsReaderImpl<T extends ReadonlyJSONObject>
       },
     });
 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     return asAsyncIterableStream(stream) as any;
   }
 
@@ -361,7 +360,6 @@ export class ToolCallArgsReaderImpl<T extends ReadonlyJSONObject>
       },
     });
 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     return asAsyncIterableStream(stream) as any;
   }
 
@@ -393,7 +391,6 @@ export class ToolCallArgsReaderImpl<T extends ReadonlyJSONObject>
       },
     });
 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     return asAsyncIterableStream(stream) as any;
   }
 }

package/src/core/tool/ToolResponse.ts

@@ -13,12 +13,14 @@ export class ToolResponse<TResult> {
     return true;
   }
 
-  readonly artifact?: ReadonlyJSONValue | undefined;
+  readonly artifact?: ReadonlyJSONValue;
   readonly result: TResult;
   readonly isError: boolean;
 
   constructor(options: ToolResponseInit<TResult>) {
-    this.artifact = options.artifact;
+    if (options.artifact !== undefined) {
+      this.artifact = options.artifact;
+    }
     this.result = options.result;
     this.isError = options.isError ?? false;
   }

package/src/core/tool/toolResultStream.ts

@@ -65,7 +65,6 @@ function getToolResponse(
 function getToolStreamResponse(
   tools: Record<string, Tool> | undefined,
   abortSignal: AbortSignal,
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
   reader: ToolCallReader<any, ReadonlyJSONValue>,
   context: {
     toolCallId: string;
@@ -94,7 +93,9 @@ export async function unstable_runPendingTools(
             return {
               ...p,
               state: "result" as const,
-              artifact: result.artifact,
+              ...(result.artifact !== undefined
+                ? { artifact: result.artifact }
+                : {}),
               result: result.result as ReadonlyJSONValue,
               isError: result.isError,
             };

package/src/core/tool/type-path-utils.ts

@@ -1,5 +1,3 @@
-/* eslint-disable @typescript-eslint/no-explicit-any */
-
 type AsNumber<K> = K extends `${infer N extends number}` ? N | K : never;
 type TupleIndex<T extends readonly any[]> = Exclude<keyof T, keyof any[]>;
 type ObjectKey<T> = keyof T & (string | number);

package/src/core/utils/stream/SSE.ts

@@ -0,0 +1,116 @@
+import { PipeableTransformStream } from "./PipeableTransformStream";
+import { LineDecoderStream } from "./LineDecoderStream";
+
+export class SSEEncoder<T> extends PipeableTransformStream<T, Uint8Array> {
+  static readonly headers = new Headers({
+    "Content-Type": "text/event-stream",
+    "Cache-Control": "no-cache",
+    Connection: "keep-alive",
+  });
+
+  headers = SSEEncoder.headers;
+
+  constructor() {
+    super((readable) =>
+      readable
+        .pipeThrough(
+          new TransformStream<T, string>({
+            transform(chunk, controller) {
+              controller.enqueue(`data: ${JSON.stringify(chunk)}\n\n`);
+            },
+          }),
+        )
+        .pipeThrough(new TextEncoderStream()),
+    );
+  }
+}
+
+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,
+          });
+        }
+      },
+    });
+  }
+}
+
+export class SSEDecoder<T> extends PipeableTransformStream<Uint8Array, T> {
+  constructor() {
+    super((readable) =>
+      readable
+        .pipeThrough(new TextDecoderStream())
+        .pipeThrough(new LineDecoderStream())
+        .pipeThrough(new SSEEventStream())
+        .pipeThrough(
+          new TransformStream<SSEEvent, T>({
+            transform(event, controller) {
+              switch (event.event) {
+                case "message":
+                  controller.enqueue(JSON.parse(event.data));
+                  break;
+                default:
+                  throw new Error(`Unknown SSE event type: ${event.event}`);
+              }
+            },
+          }),
+        ),
+    );
+  }
+}

package/src/core/utils/types.ts

@@ -50,19 +50,32 @@ type ToolCallStatus =
       reason: "cancelled" | "length" | "content-filter" | "other";
     };
 
-export type ToolCallPart = {
+type ToolCallPartBase = {
   type: "tool-call";
-  state: "partial-call" | "call" | "result";
   status: ToolCallStatus;
   toolCallId: string;
   toolName: string;
   argsText: string;
   args: ReadonlyJSONObject;
-  artifact?: unknown;
+  artifact?: ReadonlyJSONValue;
   result?: ReadonlyJSONValue;
   isError?: boolean;
 };
 
+type ToolCallPartWithoutResult = ToolCallPartBase & {
+  state: "partial-call" | "call";
+  result?: undefined;
+};
+
+type ToolCallPartWithResult = ToolCallPartBase & {
+  state: "result";
+  result: ReadonlyJSONValue;
+  artifact?: ReadonlyJSONValue;
+  isError?: boolean;
+};
+
+export type ToolCallPart = ToolCallPartWithoutResult | ToolCallPartWithResult;
+
 export type SourcePart = {
   type: "source";
   sourceType: "url";
@@ -141,7 +154,9 @@ export type AssistantMessage = {
    * @deprecated Use `parts` instead.
    */
   content: AssistantMessagePart[];
+
   metadata: {
+    unstable_state: ReadonlyJSONValue;
     unstable_data: ReadonlyJSONValue[];
     unstable_annotations: ReadonlyJSONValue[];
     steps: AssistantMessageStepMetadata[];

package/src/utils/index.ts

@@ -1,2 +1,5 @@
-export { asAsyncIterableStream, type AsyncIterableStream } from "./AsyncIterableStream";
+export {
+  asAsyncIterableStream,
+  type AsyncIterableStream,
+} from "./AsyncIterableStream";
 export { type ReadonlyJSONObject, type ReadonlyJSONValue } from "./json";