@assistant-ui/react 0.8.5 → 0.8.7

package/src/runtimes/adapters/thread-history/ThreadHistoryAdapter.ts

@@ -1,9 +1,13 @@
+import { ChatModelRunOptions, ChatModelRunResult } from "../../local";
 import {
   ExportedMessageRepository,
   ExportedMessageRepositoryItem,
 } from "../../utils/MessageRepository";
 
 export type ThreadHistoryAdapter = {
-  load(): Promise<ExportedMessageRepository>;
+  load(): Promise<ExportedMessageRepository & { unstable_resume?: boolean }>;
+  resume?(
+    options: ChatModelRunOptions,
+  ): AsyncGenerator<ChatModelRunResult, void, unknown>;
   append(item: ExportedMessageRepositoryItem): Promise<void>;
 };

package/src/runtimes/composer/BaseComposerRuntimeCore.tsx

@@ -159,7 +159,9 @@ export abstract class BaseComposerRuntimeCore
     if (!adapter) throw new Error("Attachments are not supported");
 
     const upsertAttachment = (a: PendingAttachment) => {
-      const idx = this._attachments.findIndex((attachment) => attachment.id === a.id);
+      const idx = this._attachments.findIndex(
+        (attachment) => attachment.id === a.id,
+      );
       if (idx !== -1)
         this._attachments = [
           ...this._attachments.slice(0, idx),

package/src/runtimes/dangerous-in-browser/DangerousInBrowserAdapter.ts

@@ -5,13 +5,13 @@ import {
 import { toCoreMessages } from "../edge/converters/toCoreMessages";
 import { toLanguageModelTools } from "../edge/converters/toLanguageModelTools";
 import { EdgeRuntimeRequestOptions } from "../edge/EdgeRuntimeRequestOptions";
-import { runResultStream } from "../edge/streams/runResultStream";
 import { toolResultStream } from "../edge/streams/toolResultStream";
 import { asAsyncIterable } from "../edge/EdgeChatAdapter";
 import {
   CreateEdgeRuntimeAPIOptions,
   getEdgeRuntimeStream,
 } from "../edge/createEdgeRuntimeAPI";
+import { AssistantMessageAccumulator } from "assistant-stream";
 
 export type DangerousInBrowserAdapterOptions = CreateEdgeRuntimeAPIOptions;
 
@@ -33,7 +33,7 @@ export class DangerousInBrowserAdapter implements ChatModelAdapter {
 
     const stream = res
       .pipeThrough(toolResultStream(context.tools, abortSignal))
-      .pipeThrough(runResultStream());
+      .pipeThrough(new AssistantMessageAccumulator());
 
     for await (const update of asAsyncIterable(stream)) {
       yield update;

package/src/runtimes/edge/EdgeChatAdapter.ts

@@ -2,13 +2,9 @@ import {
   ChatModelAdapter,
   ChatModelRunOptions,
 } from "../local/ChatModelAdapter";
-import { ChatModelRunResult } from "../local/ChatModelAdapter";
 import { toCoreMessages } from "./converters/toCoreMessages";
 import { toLanguageModelTools } from "./converters/toLanguageModelTools";
 import { EdgeRuntimeRequestOptions } from "./EdgeRuntimeRequestOptions";
-import { assistantDecoderStream } from "./streams/assistantDecoderStream";
-import { streamPartDecoderStream } from "./streams/utils/streamPartDecoderStream";
-import { runResultStream } from "./streams/runResultStream";
 import { toolResultStream } from "./streams/toolResultStream";
 import { toLanguageModelMessages } from "./converters";
 import { ThreadMessage } from "../../types";
@@ -16,6 +12,10 @@ import { Tool } from "../../model-context";
 import { z } from "zod";
 import zodToJsonSchema from "zod-to-json-schema";
 import { JSONSchema7 } from "json-schema";
+import {
+  AssistantMessageAccumulator,
+  DataStreamDecoder,
+} from "assistant-stream";
 
 export function asAsyncIterable<T>(
   source: ReadableStream<T>,
@@ -35,6 +35,8 @@ export function asAsyncIterable<T>(
   };
 }
 
+type HeadersValue = Record<string, string> | Headers;
+
 export type EdgeChatAdapterOptions = {
   api: string;
 
@@ -65,7 +67,13 @@ export type EdgeChatAdapterOptions = {
   onError?: (error: Error) => void;
 
   credentials?: RequestCredentials;
-  headers?: Record<string, string> | Headers;
+
+  /**
+   * Headers to be sent with the request.
+   * Can be a static headers object or a function that returns a Promise of headers.
+   */
+  headers?: HeadersValue | (() => Promise<HeadersValue>);
+
   body?: object;
 
   /**
@@ -112,7 +120,12 @@ export class EdgeChatAdapter implements ChatModelAdapter {
     unstable_assistantMessageId,
     unstable_getMessage,
   }: ChatModelRunOptions) {
-    const headers = new Headers(this.options.headers);
+    const headersValue =
+      typeof this.options.headers === "function"
+        ? await this.options.headers()
+        : this.options.headers;
+
+    const headers = new Headers(headersValue);
     headers.set("Content-Type", "application/json");
 
     const result = await fetch(this.options.api, {
@@ -153,21 +166,19 @@ export class EdgeChatAdapter implements ChatModelAdapter {
       if (!result.ok) {
         throw new Error(`Status ${result.status}: ${await result.text()}`);
       }
+      if (!result.body) {
+        throw new Error("Response body is null");
+      }
 
-      const stream = result
-        .body!.pipeThrough(streamPartDecoderStream())
-        .pipeThrough(assistantDecoderStream())
+      const stream = result.body
+        .pipeThrough(new DataStreamDecoder())
         .pipeThrough(toolResultStream(context.tools, abortSignal))
-        .pipeThrough(runResultStream());
+        .pipeThrough(new AssistantMessageAccumulator());
 
-      let update: ChatModelRunResult | undefined;
-      for await (update of asAsyncIterable(stream)) {
+      for await (const update of asAsyncIterable(stream)) {
         yield update;
       }
 
-      if (update === undefined)
-        throw new Error("No data received from Edge Runtime");
-
       this.options.onFinish?.(unstable_getMessage());
     } catch (error: unknown) {
       this.options.onError?.(error as Error);

package/src/runtimes/edge/converters/fromLanguageModelMessages.ts

@@ -4,12 +4,12 @@ import { Writable } from "stream";
 import { ReadonlyJSONObject } from "../../../utils/json/json-value";
 
 type fromLanguageModelMessagesOptions = {
-  mergeSteps: boolean;
+  mergeSteps?: boolean;
 };
 
 export const fromLanguageModelMessages = (
   lm: LanguageModelV1Message[],
-  { mergeSteps }: fromLanguageModelMessagesOptions,
+  { mergeSteps }: fromLanguageModelMessagesOptions = {},
 ): CoreMessage[] => {
   const messages: CoreMessage[] = [];
 
@@ -70,18 +70,28 @@ export const fromLanguageModelMessages = (
         break;
       }
       case "assistant": {
-        const newContent = lmMessage.content.map((part) => {
-          if (part.type === "tool-call") {
-            return {
-              type: "tool-call",
-              toolCallId: part.toolCallId,
-              toolName: part.toolName,
-              argsText: JSON.stringify(part.args),
-              args: part.args as ReadonlyJSONObject,
-            } satisfies ToolCallContentPart;
-          }
-          return part;
-        });
+        const newContent = lmMessage.content
+          .map((part) => {
+            if (part.type === "tool-call") {
+              return {
+                type: "tool-call",
+                toolCallId: part.toolCallId,
+                toolName: part.toolName,
+                argsText: JSON.stringify(part.args),
+                args: part.args as ReadonlyJSONObject,
+              } satisfies ToolCallContentPart;
+            }
+            // TODO handle these
+            if (
+              part.type === "redacted-reasoning" ||
+              part.type === "file" ||
+              part.type === "reasoning"
+            ) {
+              return null;
+            }
+            return part;
+          })
+          .filter((p) => !!p);
 
         if (mergeSteps) {
           const previousMessage = messages[messages.length - 1];

package/src/runtimes/edge/converters/toCoreMessages.ts

@@ -23,8 +23,12 @@ export const toCoreMessage = <T extends boolean = false>(
         role,
         content: message.content
           .map((part) => {
-            if (part.type === "reasoning" || part.type === "source")
-              return null; // reasoning and source parts are omitted
+            if (
+              part.type === "reasoning" ||
+              part.type === "source" ||
+              part.type === "file"
+            )
+              return null; // reasoning, source, and file parts are omitted
             if (part.type === "tool-call") {
               const { argsText, ...rest } = part;
               return rest;

package/src/runtimes/edge/converters/toLanguageModelMessages.ts

@@ -151,8 +151,9 @@ export function toLanguageModelMessages(
           const type = part.type;
           switch (type) {
             case "reasoning":
-            case "source": {
-              break; // reasoning and source parts are omitted
+            case "source":
+            case "file": {
+              break; // reasoning, source, and file parts are omitted
             }
 
             case "text": {

package/src/runtimes/edge/createEdgeRuntimeAPI.ts

@@ -5,28 +5,32 @@ import {
   LanguageModelV1Prompt,
   LanguageModelV1CallOptions,
 } from "@ai-sdk/provider";
-import { CoreMessage, ThreadMessage, ThreadStep } from "../../types/AssistantTypes";
-import { assistantEncoderStream } from "./streams/assistantEncoderStream";
+import {
+  CoreMessage,
+  ThreadMessage,
+  ThreadStep,
+} from "../../types/AssistantTypes";
 import { EdgeRuntimeRequestOptionsSchema } from "./EdgeRuntimeRequestOptions";
 import { toLanguageModelMessages } from "./converters/toLanguageModelMessages";
 import { toLanguageModelTools } from "./converters/toLanguageModelTools";
-import {
-  toolResultStream,
-  ToolResultStreamPart,
-} from "./streams/toolResultStream";
-import { runResultStream } from "./streams/runResultStream";
+import { toolResultStream } from "./streams/toolResultStream";
 import {
   LanguageModelConfig,
   LanguageModelV1CallSettings,
   LanguageModelV1CallSettingsSchema,
   Tool,
 } from "../../model-context/ModelContextTypes";
-import { CoreChatModelRunResult } from "../local/ChatModelAdapter";
-import { streamPartEncoderStream } from "./streams/utils/streamPartEncoderStream";
 import { z } from "zod";
+import {
+  AssistantMessage,
+  AssistantMessageAccumulator,
+  AssistantStreamChunk,
+  DataStreamEncoder,
+} from "assistant-stream";
+import { LanguageModelV1StreamDecoder } from "assistant-stream/ai-sdk";
 
 type FinishResult = {
-  messages: readonly (CoreMessage | ThreadMessage)[];
+  messages: readonly (CoreMessage | ThreadMessage)[];
   metadata: {
     steps: readonly ThreadStep[];
   };
@@ -44,14 +48,6 @@ export type CreateEdgeRuntimeAPIOptions = LanguageModelV1CallSettings & {
   onFinish?: (result: FinishResult) => void;
 };
 
-const voidStream = () => {
-  return new WritableStream({
-    abort(reason) {
-      console.error("Server stream processing aborted:", reason);
-    },
-  });
-};
-
 type GetEdgeRuntimeStreamOptions = {
   abortSignal: AbortSignal;
   requestData: z.infer<typeof EdgeRuntimeRequestOptionsSchema>;
@@ -102,7 +98,7 @@ export const getEdgeRuntimeStream = async ({
       ? await modelOrCreator({ apiKey, baseUrl, modelName })
       : modelOrCreator;
 
-  let stream: ReadableStream<ToolResultStreamPart>;
+  let stream: ReadableStream<AssistantStreamChunk>;
   const streamResult = await streamMessage({
     ...(settings as Partial<StreamMessageOptions>),
     ...callSettings,
@@ -115,7 +111,7 @@ export const getEdgeRuntimeStream = async ({
     tools: lmServerTools.concat(clientTools as LanguageModelV1FunctionTool[]),
     ...(toolChoice ? { toolChoice } : undefined),
   });
-  stream = streamResult.stream;
+  stream = streamResult.stream.pipeThrough(new LanguageModelV1StreamDecoder());
 
   // add tool results if we have server tools
   const canExecuteTools = hasServerTools && toolChoice?.type !== "none";
@@ -130,13 +126,13 @@ export const getEdgeRuntimeStream = async ({
     let serverStream = tees[1];
 
     if (onFinish) {
-      let lastChunk: CoreChatModelRunResult | undefined;
-      serverStream = serverStream.pipeThrough(runResultStream()).pipeThrough(
-        new TransformStream({
-          transform(chunk) {
+      let lastChunk: AssistantMessage | undefined;
+      serverStream.pipeThrough(new AssistantMessageAccumulator()).pipeTo(
+        new WritableStream({
+          write(chunk) {
             lastChunk = chunk;
           },
-          flush() {
+          close() {
             if (!lastChunk?.status || lastChunk.status.type === "running")
               return;
 
@@ -148,31 +144,22 @@ export const getEdgeRuntimeStream = async ({
                 role: "assistant",
                 content: lastChunk.content,
                 status: lastChunk.status,
-                metadata: {
-                  unstable_data: lastChunk.metadata?.unstable_data ?? [],
-                  unstable_annotations: lastChunk.metadata?.unstable_annotations ?? [],
-                  steps: lastChunk.metadata?.steps ?? [],
-                  custom: lastChunk.metadata?.custom ?? {},
-                }
+                metadata: lastChunk.metadata,
               } satisfies ThreadMessage,
             ];
             onFinish({
               messages: resultingMessages,
               metadata: {
-                // TODO
-                // eslint-disable-next-line @typescript-eslint/no-non-null-asserted-optional-chain
-                steps: lastChunk.metadata?.steps!,
+                steps: lastChunk.metadata.steps,
               },
             });
           },
+          abort(e) {
+            console.error("Server stream processing error:", e);
+          },
         }),
       );
     }
-
-    // drain the server stream
-    serverStream.pipeTo(voidStream()).catch((e) => {
-      console.error("Server stream processing error:", e);
-    });
   }
 
   return stream;
@@ -186,16 +173,12 @@ export const getEdgeRuntimeResponse = async (
   options: getEdgeRuntimeResponse.Options,
 ) => {
   const stream = await getEdgeRuntimeStream(options);
-  return new Response(
-    stream
-      .pipeThrough(assistantEncoderStream())
-      .pipeThrough(streamPartEncoderStream()),
-    {
-      headers: {
-        "Content-Type": "text/plain; charset=utf-8",
-      },
+  return new Response(stream.pipeThrough(new DataStreamEncoder()), {
+    headers: {
+      "Content-Type": "text/plain; charset=utf-8",
+      "x-vercel-ai-data-stream": "v1",
     },
-  );
+  });
 };
 
 export const createEdgeRuntimeAPI = (options: CreateEdgeRuntimeAPIOptions) => ({

package/src/runtimes/edge/index.ts

@@ -1,7 +1,5 @@
 export * from "./converters";
 
-export * from "./streams/utils";
-
 export { useEdgeRuntime, type EdgeRuntimeOptions } from "./useEdgeRuntime";
 export { EdgeChatAdapter } from "./EdgeChatAdapter";
 export type { EdgeRuntimeRequestOptions } from "./EdgeRuntimeRequestOptions";

package/src/runtimes/edge/streams/toolResultStream.ts

@@ -1,163 +1,39 @@
 import { Tool } from "../../../model-context/ModelContextTypes";
-import { LanguageModelV1StreamPart } from "@ai-sdk/provider";
 import { z } from "zod";
-import sjson from "secure-json-parse";
-import { ReadonlyJSONValue } from "../../../utils/json/json-value";
-
-export type ToolResultStreamPart =
-  | LanguageModelV1StreamPart
-  | {
-      type: "annotations";
-      annotations: ReadonlyJSONValue[];
-    }
-  | {
-      type: "data";
-      data: ReadonlyJSONValue[];
-    }
-  | {
-      type: "source";
-      source: {
-        readonly sourceType: "url";
-        readonly id: string;
-        readonly url: string;
-        readonly title?: string;
-      };
-    }
-  | {
-      type: "tool-result";
-      toolCallType: "function";
-      toolCallId: string;
-      toolName: string;
-      result: unknown;
-      isError?: boolean;
-    }
-  | {
-      type: "step-finish";
-      finishReason:
-        | "stop"
-        | "length"
-        | "content-filter"
-        | "tool-calls"
-        | "error"
-        | "other"
-        | "unknown";
-      usage: {
-        promptTokens: number;
-        completionTokens: number;
-      };
-      isContinued: boolean;
-    };
+import { ToolExecutionStream } from "assistant-stream";
 
 export function toolResultStream(
   tools: Record<string, Tool<any, any>> | undefined,
   abortSignal: AbortSignal,
 ) {
-  const toolCallExecutions = new Map<string, Promise<any>>();
-
-  return new TransformStream<ToolResultStreamPart, ToolResultStreamPart>({
-    transform(chunk, controller) {
-      // forward everything
-      controller.enqueue(chunk);
-
-      // handle tool calls
-      const chunkType = chunk.type;
-      switch (chunkType) {
-        case "tool-call": {
-          const { toolCallId, toolCallType, toolName, args: argsText } = chunk;
-          const tool = tools?.[toolName];
-          if (!tool || !tool.execute) return;
-
-          let args;
-          try {
-            args = sjson.parse(argsText);
-          } catch (e) {
-            controller.enqueue({
-              type: "tool-result",
-              toolCallType,
-              toolCallId,
-              toolName,
-              result:
-                "Function parameter parsing failed. " +
-                JSON.stringify((e as Error).message),
-              isError: true,
-            });
-            return;
-          }
-
-          toolCallExecutions.set(
-            toolCallId,
-            (async () => {
-              if (!tool.execute) return;
-
-              let executeFn = tool.execute;
-
-              if (tool.parameters instanceof z.ZodType) {
-                const result = tool.parameters.safeParse(args);
-                if (!result.success) {
-                  executeFn =
-                    tool.experimental_onSchemaValidationError ??
-                    (() => {
-                      throw (
-                        "Function parameter validation failed. " +
-                        JSON.stringify(result.error.issues)
-                      );
-                    });
-                }
-              }
-
-              try {
-                const result = await executeFn(args, {
-                  toolCallId,
-                  abortSignal,
-                });
-
-                controller.enqueue({
-                  type: "tool-result",
-                  toolCallType,
-                  toolCallId,
-                  toolName,
-                  result: result === undefined ? "<no result>" : result,
-                });
-              } catch (error) {
-                controller.enqueue({
-                  type: "tool-result",
-                  toolCallType,
-                  toolCallId,
-                  toolName,
-                  result: "Error: " + error,
-                  isError: true,
-                });
-              } finally {
-                toolCallExecutions.delete(toolCallId);
-              }
-            })(),
-          );
-          break;
-        }
-
-        // ignore other parts
-        case "text-delta":
-        case "reasoning":
-        case "source":
-        case "tool-call-delta":
-        case "tool-result":
-        case "step-finish":
-        case "finish":
-        case "error":
-        case "response-metadata":
-        case "annotations":
-        case "data":
-          break;
-
-        default: {
-          const unhandledType: never = chunkType;
-          throw new Error(`Unhandled chunk type: ${unhandledType}`);
-        }
+  return new ToolExecutionStream(({ toolCallId, toolName, args }) => {
+    const tool = tools?.[toolName];
+    if (!tool || !tool.execute) return undefined;
+
+    let executeFn = tool.execute;
+
+    if (tool.parameters instanceof z.ZodType) {
+      const result = tool.parameters.safeParse(args);
+      if (!result.success) {
+        executeFn =
+          tool.experimental_onSchemaValidationError ??
+          (() => {
+            throw (
+              "Function parameter validation failed. " +
+              JSON.stringify(result.error.issues)
+            );
+          });
       }
-    },
+    }
+
+    const getResult = async () => {
+      const result = await executeFn(args, {
+        toolCallId,
+        abortSignal,
+      });
+      return result === undefined ? "<no result>" : result;
+    };
 
-    async flush() {
-      await Promise.all(toolCallExecutions.values());
-    },
+    return getResult();
   });
 }

package/src/runtimes/external-store/ThreadMessageLike.tsx

@@ -100,6 +100,7 @@ export const fromThreadMessageLike = (
                 if (part.text.trim().length === 0) return null;
                 return part;
 
+              case "file":
               case "source":
                 return part;
 
@@ -121,7 +122,7 @@ export const fromThreadMessageLike = (
               }
 
               default: {
-                const unhandledType: "image" | "audio" | "file" = type;
+                const unhandledType: "image" | "audio" = type;
                 throw new Error(
                   `Unsupported assistant content part type: ${unhandledType}`,
                 );

package/src/runtimes/external-store/external-message-converter.tsx

@@ -317,7 +317,7 @@ export const useExternalMessageConverter = <T extends WeakKey>({
     ] = messages;
 
     return threadMessages;
-  }, [state, messages, isRunning]);
+  }, [state, messages, isRunning, joinStrategy]);
 };
 
 const shallowArrayEqual = (a: unknown[], b: unknown[]) => {

package/src/runtimes/local/ChatModelAdapter.tsx

@@ -1,4 +1,5 @@
 import type {
+  FileContentPart,
   MessageStatus,
   ReasoningContentPart,
   RunConfig,
@@ -34,6 +35,7 @@ export type CoreChatModelRunResult = Omit<ChatModelRunResult, "content"> & {
     | ReasoningContentPart
     | ToolCallContentPart
     | SourceContentPart
+    | FileContentPart
   )[];
 };