@assistant-ui/react 0.8.14 → 0.8.16

package/src/runtimes/edge/{EdgeChatAdapter.ts → EdgeModelAdapter.ts}

@@ -37,7 +37,7 @@ export function asAsyncIterable<T>(
 
 type HeadersValue = Record<string, string> | Headers;
 
-export type EdgeChatAdapterOptions = {
+export type EdgeModelAdapterOptions = {
   api: string;
 
   // experimental_prepareRequestBody?: (options: {
@@ -109,8 +109,8 @@ const toAISDKTools = (tools: Record<string, Tool<any, any>>) => {
   );
 };
 
-export class EdgeChatAdapter implements ChatModelAdapter {
-  constructor(private options: EdgeChatAdapterOptions) {}
+export class EdgeModelAdapter implements ChatModelAdapter {
+  constructor(private options: EdgeModelAdapterOptions) {}
 
   async *run({
     messages,

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

@@ -11,6 +11,7 @@ import {
   ThreadMessage,
   TextContentPart,
   CoreToolCallContentPart,
+  ToolCallContentPart,
 } from "../../../types/AssistantTypes";
 
 const assistantMessageSplitter = () => {
@@ -46,7 +47,7 @@ const assistantMessageSplitter = () => {
 
       assistantMessage.content.push(part);
     },
-    addToolCallPart: (part: CoreToolCallContentPart) => {
+    addToolCallPart: (part: CoreToolCallContentPart | ToolCallContentPart) => {
       assistantMessage.content.push({
         type: "tool-call",
         toolCallId: part.toolCallId,
@@ -58,6 +59,7 @@ const assistantMessageSplitter = () => {
         type: "tool-result",
         toolCallId: part.toolCallId,
         toolName: part.toolName,
+        ...("artifact" in part ? { artifact: part.artifact } : {}),
         result:
           part.result === undefined
             ? "Error: tool is has no configured code to run"

package/src/runtimes/edge/index.ts

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

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

@@ -1,39 +1,88 @@
 import { Tool } from "../../../model-context/ModelContextTypes";
 import { z } from "zod";
-import { ToolExecutionStream } from "assistant-stream";
+import { AssistantMessage, ToolExecutionStream } from "assistant-stream";
+import { ToolResponse } from "assistant-stream/core/effects/ToolResponse";
 
-export function toolResultStream(
+function getToolResponse(
   tools: Record<string, Tool<any, any>> | undefined,
   abortSignal: AbortSignal,
+  toolCall: {
+    toolCallId: string;
+    toolName: string;
+    args: unknown;
+  },
 ) {
-  return new ToolExecutionStream(({ toolCallId, toolName, args }) => {
-    const tool = tools?.[toolName];
-    if (!tool || !tool.execute) return undefined;
+  const tool = tools?.[toolCall.toolName];
+  if (!tool || !tool.execute) return undefined;
 
-    let executeFn = tool.execute;
+  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)
-            );
-          });
-      }
+  if (tool.parameters instanceof z.ZodType) {
+    const result = tool.parameters.safeParse(toolCall.args);
+    if (!result.success) {
+      executeFn =
+        tool.experimental_onSchemaValidationError ??
+        (() => {
+          throw new Error(
+            `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;
-    };
+  const getResult = async () => {
+    const result = await executeFn(toolCall.args, {
+      toolCallId: toolCall.toolCallId,
+      abortSignal,
+    });
+    if (result instanceof ToolResponse) return result;
+    return new ToolResponse({
+      result: result === undefined ? "<no result>" : result,
+    });
+  };
+
+  return getResult();
+}
 
-    return getResult();
-  });
+export async function unstable_runPendingTools(
+  message: AssistantMessage,
+  tools: Record<string, Tool<any, any>> | undefined,
+  abortSignal: AbortSignal,
+) {
+  // TODO parallel tool calling
+  for (const part of message.parts) {
+    if (part.type === "tool-call") {
+      const promiseOrUndefined = getToolResponse(tools, abortSignal, part);
+      if (promiseOrUndefined) {
+        const result = await promiseOrUndefined;
+        const updatedParts = message.parts.map((p) => {
+          if (p.type === "tool-call" && p.toolCallId === part.toolCallId) {
+            return {
+              ...p,
+              state: "result" as const,
+              artifact: result.artifact,
+              result: result.result,
+              isError: result.isError,
+            };
+          }
+          return p;
+        });
+        message = {
+          ...message,
+          parts: updatedParts,
+          content: updatedParts,
+        };
+      }
+    }
+  }
+  return message;
+}
+
+export function toolResultStream(
+  tools: Record<string, Tool<any, any>> | undefined,
+  abortSignal: AbortSignal,
+) {
+  return new ToolExecutionStream((toolCall) =>
+    getToolResponse(tools, abortSignal, toolCall),
+  );
 }

package/src/runtimes/edge/useEdgeRuntime.ts

@@ -1,17 +1,17 @@
 "use client";
 
 import { LocalRuntimeOptions, useLocalRuntime } from "..";
-import { EdgeChatAdapterOptions, EdgeChatAdapter } from "./EdgeChatAdapter";
+import { EdgeModelAdapterOptions, EdgeModelAdapter } from "./EdgeModelAdapter";
 import { splitLocalRuntimeOptions } from "../local/LocalRuntimeOptions";
 
-export type EdgeRuntimeOptions = EdgeChatAdapterOptions & LocalRuntimeOptions;
+export type EdgeRuntimeOptions = EdgeModelAdapterOptions & LocalRuntimeOptions;
 
 export const useEdgeRuntime = (options: EdgeRuntimeOptions) => {
   const { localRuntimeOptions, otherOptions } =
     splitLocalRuntimeOptions(options);
 
   return useLocalRuntime(
-    new EdgeChatAdapter(otherOptions),
+    new EdgeModelAdapter(otherOptions),
     localRuntimeOptions,
   );
 };

package/src/runtimes/remote-thread-list/EMPTY_THREAD_CORE.tsx

@@ -24,6 +24,10 @@ export const EMPTY_THREAD_CORE: ThreadRuntimeCore = {
     throw EMPTY_THREAD_ERROR;
   },
 
+  resumeRun() {
+    throw EMPTY_THREAD_ERROR;
+  },
+
   cancelRun() {
     throw EMPTY_THREAD_ERROR;
   },

package/src/runtimes/utils/MessageRepository.tsx

@@ -5,31 +5,62 @@ import { ThreadMessageLike } from "../external-store";
 import { getAutoStatus } from "../external-store/auto-status";
 import { fromThreadMessageLike } from "../external-store/ThreadMessageLike";
 
+/**
+ * Represents a parent node in the repository tree structure.
+ */
 type RepositoryParent = {
+  /** IDs of child messages */
   children: string[];
+  /** Reference to the next message in the active branch */
   next: RepositoryMessage | null;
 };
 
+/**
+ * Represents a message node in the repository tree structure.
+ */
 type RepositoryMessage = RepositoryParent & {
+  /** Reference to the parent message */
   prev: RepositoryMessage | null;
+  /** The actual message data */
   current: ThreadMessage;
+  /** The depth level in the tree (0 for root messages) */
   level: number;
 };
 
+/**
+ * Represents a message item that can be exported from the repository.
+ */
 export type ExportedMessageRepositoryItem = {
+  /** The message data */
   message: ThreadMessage;
+  /** ID of the parent message, or null for root messages */
   parentId: string | null;
 };
 
+/**
+ * Represents the entire repository state for export/import.
+ */
 export type ExportedMessageRepository = {
+  /** ID of the head message, or null/undefined if no head */
   headId?: string | null;
+  /** Array of all messages with their parent references */
   messages: Array<{
     message: ThreadMessage;
     parentId: string | null;
   }>;
 };
 
+/**
+ * Utility functions for working with exported message repositories.
+ */
 export const ExportedMessageRepository = {
+  /**
+   * Converts an array of messages to an ExportedMessageRepository format.
+   * Creates parent-child relationships based on the order of messages in the array.
+   *
+   * @param messages - Array of message-like objects to convert
+   * @returns ExportedMessageRepository with parent-child relationships established
+   */
   fromArray: (
     messages: readonly ThreadMessageLike[],
   ): ExportedMessageRepository => {
@@ -46,6 +77,12 @@ export const ExportedMessageRepository = {
   },
 };
 
+/**
+ * Recursively finds the head (leaf) message in a branch.
+ *
+ * @param message - The starting message or parent node
+ * @returns The leaf message of the branch, or null if not found
+ */
 const findHead = (
   message: RepositoryMessage | RepositoryParent,
 ): RepositoryMessage | null => {
@@ -54,11 +91,20 @@ const findHead = (
   return null;
 };
 
+/**
+ * A utility class for caching computed values and invalidating the cache when needed.
+ */
 class CachedValue<T> {
   private _value: T | null = null;
 
+  /**
+   * @param func - The function that computes the cached value
+   */
   constructor(private func: () => T) {}
 
+  /**
+   * Gets the cached value, computing it if necessary.
+   */
   get value() {
     if (this._value === null) {
       this._value = this.func();
@@ -66,19 +112,40 @@ class CachedValue<T> {
     return this._value;
   }
 
+  /**
+   * Invalidates the cache, forcing recomputation on next access.
+   */
   dirty() {
     this._value = null;
   }
 }
 
+/**
+ * A repository that manages a tree of messages with branching capabilities.
+ * Supports operations like adding, updating, and deleting messages, as well as
+ * managing multiple conversation branches.
+ */
 export class MessageRepository {
-  private messages = new Map<string, RepositoryMessage>(); // message_id -> item
+  /** Map of message IDs to repository message objects */
+  private messages = new Map<string, RepositoryMessage>();
+  /** Reference to the current head (most recent) message in the active branch */
   private head: RepositoryMessage | null = null;
+  /** Root node of the tree structure */
   private root: RepositoryParent = {
     children: [],
     next: null,
   };
 
+  /**
+   * Performs link/unlink operations between messages in the tree.
+   *
+   * @param newParent - The new parent message, or null
+   * @param child - The child message to operate on
+   * @param operation - The type of operation to perform:
+   *   - "cut": Remove the child from its current parent
+   *   - "link": Add the child to a new parent
+   *   - "relink": Both cut and link operations
+   */
   private performOp(
     newParent: RepositoryMessage | null,
     child: RepositoryMessage,
@@ -139,6 +206,7 @@ export class MessageRepository {
     }
   }
 
+  /** Cached array of messages in the current active branch, from root to head */
   private _messages = new CachedValue<readonly ThreadMessage[]>(() => {
     const messages = new Array<ThreadMessage>(this.head?.level ?? 0);
     for (let current = this.head; current; current = current.prev) {
@@ -147,14 +215,31 @@ export class MessageRepository {
     return messages;
   });
 
+  /**
+   * Gets the ID of the current head message.
+   * @returns The ID of the head message, or null if no messages exist
+   */
   get headId() {
     return this.head?.current.id ?? null;
   }
 
+  /**
+   * Gets all messages in the current active branch, from root to head.
+   * @returns Array of messages in the current branch
+   */
   getMessages() {
     return this._messages.value;
   }
 
+  /**
+   * Adds a new message or updates an existing one in the repository.
+   * If the message ID already exists, the message is updated and potentially relinked to a new parent.
+   * If the message is new, it's added as a child of the specified parent.
+   *
+   * @param parentId - ID of the parent message, or null for root messages
+   * @param message - The message to add or update
+   * @throws Error if the parent message is not found
+   */
   addOrUpdateMessage(parentId: string | null, message: ThreadMessage) {
     const existingItem = this.messages.get(message.id);
     const prev = parentId ? this.messages.get(parentId) : null;
@@ -190,6 +275,13 @@ export class MessageRepository {
     this._messages.dirty();
   }
 
+  /**
+   * Gets a message and its parent ID by message ID.
+   *
+   * @param messageId - ID of the message to retrieve
+   * @returns Object containing the message and its parent ID
+   * @throws Error if the message is not found
+   */
   getMessage(messageId: string) {
     const message = this.messages.get(messageId);
     if (!message)
@@ -203,6 +295,14 @@ export class MessageRepository {
     };
   }
 
+  /**
+   * Adds an optimistic message to the repository.
+   * An optimistic message is a temporary placeholder that will be replaced by a real message later.
+   *
+   * @param parentId - ID of the parent message, or null for root messages
+   * @param message - The core message to convert to an optimistic message
+   * @returns The generated optimistic ID
+   */
   appendOptimisticMessage(parentId: string | null, message: CoreMessage) {
     let optimisticId: string;
     do {
@@ -220,6 +320,15 @@ export class MessageRepository {
     return optimisticId;
   }
 
+  /**
+   * Deletes a message from the repository and relinks its children.
+   *
+   * @param messageId - ID of the message to delete
+   * @param replacementId - Optional ID of the message to become the new parent of the children,
+   *                       undefined means use the deleted message's parent,
+   *                       null means use the root
+   * @throws Error if the message or replacement is not found
+   */
   deleteMessage(messageId: string, replacementId?: string | null | undefined) {
     const message = this.messages.get(messageId);
 
@@ -258,6 +367,13 @@ export class MessageRepository {
     this._messages.dirty();
   }
 
+  /**
+   * Gets all branch IDs (sibling messages) at the level of a specified message.
+   *
+   * @param messageId - ID of the message to find branches for
+   * @returns Array of message IDs representing branches
+   * @throws Error if the message is not found
+   */
   getBranches(messageId: string) {
     const message = this.messages.get(messageId);
     if (!message)
@@ -269,6 +385,12 @@ export class MessageRepository {
     return children;
   }
 
+  /**
+   * Switches the active branch to the one containing the specified message.
+   *
+   * @param messageId - ID of the message in the branch to switch to
+   * @throws Error if the branch is not found
+   */
   switchToBranch(messageId: string) {
     const message = this.messages.get(messageId);
     if (!message)
@@ -284,6 +406,12 @@ export class MessageRepository {
     this._messages.dirty();
   }
 
+  /**
+   * Resets the head to a specific message or null.
+   *
+   * @param messageId - ID of the message to set as head, or null to clear the head
+   * @throws Error if the message is not found
+   */
   resetHead(messageId: string | null) {
     if (messageId === null) {
       this.head = null;
@@ -311,6 +439,9 @@ export class MessageRepository {
     this._messages.dirty();
   }
 
+  /**
+   * Clears all messages from the repository.
+   */
   clear(): void {
     this.messages.clear();
     this.head = null;
@@ -321,6 +452,11 @@ export class MessageRepository {
     this._messages.dirty();
   }
 
+  /**
+   * Exports the repository state for persistence.
+   *
+   * @returns Exportable repository state
+   */
   export(): ExportedMessageRepository {
     const exportItems: ExportedMessageRepository["messages"] = [];
 
@@ -339,6 +475,11 @@ export class MessageRepository {
     };
   }
 
+  /**
+   * Imports repository state from an exported repository.
+   *
+   * @param repository - The exported repository state to import
+   */
   import({ headId, messages }: ExportedMessageRepository) {
     for (const { message, parentId } of messages) {
       this.addOrUpdateMessage(parentId, message);