assistant-stream 0.2.5 → 0.2.7

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

@@ -8,18 +8,24 @@ import {
   ToolCallReader,
   ToolCallResponseReader,
 } from "./tool-types";
-import { TypeAtPath, TypePath } from "./type-path-utils";
+import { DeepPartial, TypeAtPath, TypePath } from "./type-path-utils";
 import { ToolResponse } from "./ToolResponse";
+import {
+  asAsyncIterableStream,
+  AsyncIterableStream,
+  ReadonlyJSONObject,
+  ReadonlyJSONValue,
+} from "../../utils";
 
 // TODO: remove dispose
 
-function getField<T>(obj: T, fieldPath: (string | number)[]): any {
-  let current: any = obj;
+function getField<T>(obj: T, fieldPath: (string | number)[]): unknown {
+  let current: unknown = obj;
   for (const key of fieldPath) {
     if (current === undefined || current === null) {
       return undefined;
     }
-    current = current[key as string | number];
+    current = current[key as keyof typeof current];
   }
   return current;
 }
@@ -29,14 +35,14 @@ interface Handle {
   dispose(): void;
 }
 
-class GetHandle<T> implements Handle {
-  private resolve: (value: any) => void;
+class GetHandle<T, TValue> implements Handle {
+  private resolve: (value: TValue) => void;
   private reject: (reason: unknown) => void;
   private disposed = false;
   private fieldPath: (string | number)[];
 
   constructor(
-    resolve: (value: any) => void,
+    resolve: (value: TValue) => void,
     reject: (reason: unknown) => void,
     fieldPath: (string | number)[],
   ) {
@@ -58,7 +64,7 @@ class GetHandle<T> implements Handle {
       ) {
         const value = getField(args as T, this.fieldPath);
         if (value !== undefined) {
-          this.resolve(value);
+          this.resolve(value as TValue);
           this.dispose();
         }
       }
@@ -74,12 +80,12 @@ class GetHandle<T> implements Handle {
 }
 
 class StreamValuesHandle<T> implements Handle {
-  private controller: ReadableStreamDefaultController<any>;
+  private controller: ReadableStreamDefaultController<unknown>;
   private disposed = false;
   private fieldPath: (string | number)[];
 
   constructor(
-    controller: ReadableStreamDefaultController<any>,
+    controller: ReadableStreamDefaultController<unknown>,
     fieldPath: (string | number)[],
   ) {
     this.controller = controller;
@@ -118,13 +124,13 @@ class StreamValuesHandle<T> implements Handle {
 }
 
 class StreamTextHandle<T> implements Handle {
-  private controller: ReadableStreamDefaultController<any>;
+  private controller: ReadableStreamDefaultController<unknown>;
   private disposed = false;
   private fieldPath: (string | number)[];
-  private lastValue: any = undefined;
+  private lastValue: string | undefined = undefined;
 
   constructor(
-    controller: ReadableStreamDefaultController<any>,
+    controller: ReadableStreamDefaultController<unknown>,
     fieldPath: (string | number)[],
   ) {
     this.controller = controller;
@@ -165,13 +171,13 @@ class StreamTextHandle<T> implements Handle {
 }
 
 class ForEachHandle<T> implements Handle {
-  private controller: ReadableStreamDefaultController<any>;
+  private controller: ReadableStreamDefaultController<unknown>;
   private disposed = false;
   private fieldPath: (string | number)[];
   private processedIndexes = new Set<number>();
 
   constructor(
-    controller: ReadableStreamDefaultController<any>,
+    controller: ReadableStreamDefaultController<unknown>,
     fieldPath: (string | number)[],
   ) {
     this.controller = controller;
@@ -182,7 +188,7 @@ class ForEachHandle<T> implements Handle {
     if (this.disposed) return;
 
     try {
-      const array = getField(args as T, this.fieldPath) as unknown as any[];
+      const array = getField(args as T, this.fieldPath);
 
       if (!Array.isArray(array)) {
         return;
@@ -226,10 +232,12 @@ class ForEachHandle<T> implements Handle {
 }
 
 // Implementation of ToolCallReader that uses stream of partial JSON
-export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
+export class ToolCallArgsReaderImpl<T extends ReadonlyJSONObject>
+  implements ToolCallArgsReader<T>
+{
   private argTextDeltas: ReadableStream<string>;
   private handles: Set<Handle> = new Set();
-  private args: any = parsePartialJsonObject("");
+  private args: unknown = parsePartialJsonObject("");
 
   constructor(argTextDeltas: ReadableStream<string>) {
     this.argTextDeltas = argTextDeltas;
@@ -268,8 +276,12 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
   get<PathT extends TypePath<T>>(
     ...fieldPath: PathT
   ): Promise<TypeAtPath<T, PathT>> {
-    return new Promise<any>((resolve, reject) => {
-      const handle = new GetHandle<T>(resolve, reject, fieldPath);
+    return new Promise<TypeAtPath<T, PathT>>((resolve, reject) => {
+      const handle = new GetHandle<T, TypeAtPath<T, PathT>>(
+        resolve,
+        reject,
+        fieldPath,
+      );
 
       // Check if the field is already complete in current args
       if (
@@ -281,7 +293,7 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
       ) {
         const value = getField(this.args as T, fieldPath);
         if (value !== undefined) {
-          resolve(value);
+          resolve(value as TypeAtPath<T, PathT>);
           return;
         }
       }
@@ -291,11 +303,13 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
     });
   }
 
-  streamValues<PathT extends TypePath<T>>(...fieldPath: PathT): any {
+  streamValues<PathT extends TypePath<T>>(
+    ...fieldPath: PathT
+  ): AsyncIterableStream<DeepPartial<TypeAtPath<T, PathT>>> {
     // Use a type assertion to convert the complex TypePath to a simple array
     const simplePath = fieldPath as unknown as (string | number)[];
 
-    const stream = new ReadableStream<any>({
+    const stream = new ReadableStream<DeepPartial<TypeAtPath<T, PathT>>>({
       start: (controller) => {
         const handle = new StreamValuesHandle<T>(controller, simplePath);
         this.handles.add(handle);
@@ -315,15 +329,18 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
       },
     });
 
-    // For type compatibility, cast the stream to the required type
-    return stream as any;
+    return asAsyncIterableStream(stream) as any;
   }
 
-  streamText<PathT extends TypePath<T>>(...fieldPath: PathT): any {
+  streamText<PathT extends TypePath<T>>(
+    ...fieldPath: PathT
+  ): TypeAtPath<T, PathT> extends string & infer U
+    ? AsyncIterableStream<U>
+    : never {
     // Use a type assertion to convert the complex TypePath to a simple array
     const simplePath = fieldPath as unknown as (string | number)[];
 
-    const stream = new ReadableStream<any>({
+    const stream = new ReadableStream<unknown>({
       start: (controller) => {
         const handle = new StreamTextHandle<T>(controller, simplePath);
         this.handles.add(handle);
@@ -343,15 +360,18 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
       },
     });
 
-    // For type compatibility, cast the stream to the required type
-    return stream as any;
+    return asAsyncIterableStream(stream) as any;
   }
 
-  forEach<PathT extends TypePath<T>>(...fieldPath: PathT): any {
+  forEach<PathT extends TypePath<T>>(
+    ...fieldPath: PathT
+  ): TypeAtPath<T, PathT> extends Array<infer U>
+    ? AsyncIterableStream<U>
+    : never {
     // Use a type assertion to convert the complex TypePath to a simple array
     const simplePath = fieldPath as unknown as (string | number)[];
 
-    const stream = new ReadableStream<any>({
+    const stream = new ReadableStream<unknown>({
       start: (controller) => {
         const handle = new ForEachHandle<T>(controller, simplePath);
         this.handles.add(handle);
@@ -371,12 +391,11 @@ export class ToolCallArgsReaderImpl<T> implements ToolCallArgsReader<T> {
       },
     });
 
-    // For type compatibility, cast the stream to the required type
-    return stream as any;
+    return asAsyncIterableStream(stream) as any;
   }
 }
 
-export class ToolCallResponseReaderImpl<TResult>
+export class ToolCallResponseReaderImpl<TResult extends ReadonlyJSONValue>
   implements ToolCallResponseReader<TResult>
 {
   constructor(private readonly promise: Promise<ToolResponse<TResult>>) {}
@@ -386,8 +405,10 @@ export class ToolCallResponseReaderImpl<TResult>
   }
 }
 
-export class ToolCallReaderImpl<TArgs, TResult>
-  implements ToolCallReader<TArgs, TResult>
+export class ToolCallReaderImpl<
+  TArgs extends ReadonlyJSONObject,
+  TResult extends ReadonlyJSONValue,
+> implements ToolCallReader<TArgs, TResult>
 {
   public readonly args: ToolCallArgsReaderImpl<TArgs>;
   public readonly response: ToolCallResponseReaderImpl<TResult>;

package/src/core/tool/ToolExecutionStream.ts

@@ -5,7 +5,10 @@ import {
   AssistantMetaTransformStream,
 } from "../utils/stream/AssistantMetaTransformStream";
 import { PipeableTransformStream } from "../utils/stream/PipeableTransformStream";
-import { ReadonlyJSONValue } from "../../utils/json/json-value";
+import {
+  ReadonlyJSONObject,
+  ReadonlyJSONValue,
+} from "../../utils/json/json-value";
 import { ToolResponse } from "./ToolResponse";
 import { withPromiseOrValue } from "../utils/withPromiseOrValue";
 import { ToolCallReaderImpl } from "./ToolCallReader";
@@ -14,13 +17,16 @@ import { ToolCallReader } from "./tool-types";
 type ToolCallback = (toolCall: {
   toolCallId: string;
   toolName: string;
-  args: unknown;
+  args: ReadonlyJSONObject;
 }) =>
   | Promise<ToolResponse<ReadonlyJSONValue>>
   | ToolResponse<ReadonlyJSONValue>
   | undefined;
 
-type ToolStreamCallback = <TArgs, TResult>(toolCall: {
+type ToolStreamCallback = <
+  TArgs extends ReadonlyJSONObject = ReadonlyJSONObject,
+  TResult extends ReadonlyJSONValue = ReadonlyJSONValue,
+>(toolCall: {
   reader: ToolCallReader<TArgs, TResult>;
   toolCallId: string;
   toolName: string;
@@ -39,7 +45,7 @@ export class ToolExecutionStream extends PipeableTransformStream<
     const toolCallPromises = new Map<string, PromiseLike<void>>();
     const toolCallControllers = new Map<
       string,
-      ToolCallReaderImpl<unknown, unknown>
+      ToolCallReaderImpl<ReadonlyJSONObject, ReadonlyJSONValue>
     >();
 
     super((readable) => {
@@ -58,7 +64,10 @@ export class ToolExecutionStream extends PipeableTransformStream<
           switch (type) {
             case "part-start":
               if (chunk.part.type === "tool-call") {
-                const reader = new ToolCallReaderImpl<unknown, unknown>();
+                const reader = new ToolCallReaderImpl<
+                  ReadonlyJSONObject,
+                  ReadonlyJSONValue
+                >();
                 toolCallControllers.set(chunk.part.toolCallId, reader);
 
                 options.streamCall({

package/src/core/tool/ToolResponse.ts

@@ -13,17 +13,21 @@ 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;
   }
 
-  static [Symbol.hasInstance](obj: unknown): obj is ToolResponse<unknown> {
+  static [Symbol.hasInstance](
+    obj: unknown,
+  ): obj is ToolResponse<ReadonlyJSONValue> {
     return (
       typeof obj === "object" && obj !== null && TOOL_RESPONSE_SYMBOL in obj
     );

package/src/core/tool/tool-types.ts

@@ -1,6 +1,5 @@
 import { JSONSchema7 } from "json-schema";
-import { TypeAtPath, TypePath } from "./type-path-utils";
-import { DeepPartial } from "ai";
+import { DeepPartial, TypeAtPath, TypePath } from "./type-path-utils";
 import { AsyncIterableStream } from "../../utils";
 import type { StandardSchemaV1 } from "@standard-schema/spec";
 import { ToolResponse } from "./ToolResponse";
@@ -13,7 +12,7 @@ import { ToolResponse } from "./ToolResponse";
  *
  * @template TArgs The type of arguments being read.
  */
-export interface ToolCallArgsReader<TArgs> {
+export interface ToolCallArgsReader<TArgs extends Record<string, unknown>> {
   /**
    * Returns a promise that will resolve to the value at the given path,
    * as soon as that path is generated by the LLM.
@@ -63,7 +62,10 @@ export interface ToolCallResponseReader<TResult> {
   get: () => Promise<ToolResponse<TResult>>;
 }
 
-export interface ToolCallReader<TArgs, TResult> {
+export interface ToolCallReader<
+  TArgs extends Record<string, unknown> = Record<string, unknown>,
+  TResult = unknown,
+> {
   args: ToolCallArgsReader<TArgs>;
   response: ToolCallResponseReader<TResult>;
 
@@ -85,7 +87,10 @@ export type ToolExecuteFunction<TArgs, TResult> = (
   context: ToolExecutionContext,
 ) => TResult | Promise<TResult>;
 
-export type ToolStreamCallFunction<TArgs, TResult> = (
+export type ToolStreamCallFunction<
+  TArgs extends Record<string, unknown> = Record<string, unknown>,
+  TResult = unknown,
+> = (
   reader: ToolCallReader<TArgs, TResult>,
   context: ToolExecutionContext,
 ) => void;

package/src/core/tool/toolResultStream.ts

@@ -3,6 +3,7 @@ import { StandardSchemaV1 } from "@standard-schema/spec";
 import { ToolResponse } from "./ToolResponse";
 import { ToolExecutionStream } from "./ToolExecutionStream";
 import { AssistantMessage } from "../utils/types";
+import { ReadonlyJSONObject, ReadonlyJSONValue } from "../../utils";
 
 const isStandardSchemaV1 = (
   schema: unknown,
@@ -16,18 +17,20 @@ const isStandardSchemaV1 = (
 };
 
 function getToolResponse(
-  tools: Record<string, Tool<any, any>> | undefined,
+  tools: Record<string, Tool> | undefined,
   abortSignal: AbortSignal,
   toolCall: {
     toolCallId: string;
     toolName: string;
-    args: unknown;
+    args: ReadonlyJSONObject;
   },
 ) {
   const tool = tools?.[toolCall.toolName];
   if (!tool || !tool.execute) return undefined;
 
-  const getResult = async (toolExecute: ToolExecuteFunction<any, any>) => {
+  const getResult = async (
+    toolExecute: ToolExecuteFunction<ReadonlyJSONObject, unknown>,
+  ): Promise<ToolResponse<ReadonlyJSONValue>> => {
     let executeFn = toolExecute;
 
     if (isStandardSchemaV1(tool.parameters)) {
@@ -45,11 +48,12 @@ function getToolResponse(
       }
     }
 
-    const result = await executeFn(toolCall.args, {
+    const result = (await executeFn(toolCall.args, {
       toolCallId: toolCall.toolCallId,
       abortSignal,
-    });
-    if (result instanceof ToolResponse) return result;
+    })) as unknown as ReadonlyJSONValue;
+    if (result instanceof ToolResponse)
+      return result as ToolResponse<ReadonlyJSONValue>;
     return new ToolResponse({
       result: result === undefined ? "<no result>" : result,
     });
@@ -58,10 +62,10 @@ function getToolResponse(
   return getResult(tool.execute);
 }
 
-function getToolStreamResponse<TArgs, TResult>(
-  tools: Record<string, Tool<any, any>> | undefined,
+function getToolStreamResponse(
+  tools: Record<string, Tool> | undefined,
   abortSignal: AbortSignal,
-  reader: ToolCallReader<TArgs, TResult>,
+  reader: ToolCallReader<any, ReadonlyJSONValue>,
   context: {
     toolCallId: string;
     toolName: string;
@@ -75,7 +79,7 @@ function getToolStreamResponse<TArgs, TResult>(
 
 export async function unstable_runPendingTools(
   message: AssistantMessage,
-  tools: Record<string, Tool<any, any>> | undefined,
+  tools: Record<string, Tool> | undefined,
   abortSignal: AbortSignal,
 ) {
   // TODO parallel tool calling
@@ -89,8 +93,10 @@ export async function unstable_runPendingTools(
             return {
               ...p,
               state: "result" as const,
-              artifact: result.artifact,
-              result: result.result,
+              ...(result.artifact !== undefined
+                ? { artifact: result.artifact }
+                : {}),
+              result: result.result as ReadonlyJSONValue,
               isError: result.isError,
             };
           }
@@ -108,7 +114,7 @@ export async function unstable_runPendingTools(
 }
 
 export function toolResultStream(
-  tools: Record<string, Tool<any, any>> | undefined,
+  tools: Record<string, Tool> | undefined,
   abortSignal: AbortSignal,
 ) {
   return new ToolExecutionStream({

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

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

package/src/utils/json/index.ts

@@ -0,0 +1 @@
+export { type ReadonlyJSONObject, type ReadonlyJSONValue } from "./json-value";