@assistant-ui/react 0.11.5 → 0.11.7

package/src/legacy-runtime/runtime-cores/assistant-transport/useAssistantTransportRuntime.tsx

@@ -0,0 +1,232 @@
+"use client";
+
+import {
+  ReadonlyJSONObject,
+  ReadonlyJSONValue,
+  asAsyncIterableStream,
+} from "assistant-stream/utils";
+import { AppendMessage } from "../../../types";
+import { useExternalStoreRuntime } from "../external-store/useExternalStoreRuntime";
+import { AssistantRuntime } from "../../runtime/AssistantRuntime";
+import { AddToolResultOptions } from "../core";
+import { useState, useRef, useMemo } from "react";
+import {
+  AssistantMessageAccumulator,
+  DataStreamDecoder,
+  unstable_createInitialMessage as createInitialMessage,
+} from "assistant-stream";
+import {
+  AssistantTransportOptions,
+  AddMessageCommand,
+  AddToolResultCommand,
+  UserMessagePart,
+  QueuedCommand,
+} from "./types";
+import { useCommandQueue } from "./commandQueue";
+import { useRunManager } from "./runManager";
+import { useConvertedState } from "./useConvertedState";
+import { useToolInvocations } from "./useToolInvocations";
+import { toAISDKTools, getEnabledTools, createRequestHeaders } from "./utils";
+import { useRemoteThreadListRuntime } from "../remote-thread-list/useRemoteThreadListRuntime";
+import { InMemoryThreadListAdapter } from "../remote-thread-list/adapter/in-memory";
+
+const useAssistantTransportThreadRuntime = <T,>(
+  options: AssistantTransportOptions<T>,
+): AssistantRuntime => {
+  const agentStateRef = useRef(options.initialState);
+  const [, rerender] = useState(0);
+  const resumeFlagRef = useRef(false);
+  const commandQueue = useCommandQueue({
+    onQueue: () => runManager.schedule(),
+  });
+
+  const runManager = useRunManager({
+    onRun: async (signal: AbortSignal) => {
+      const isResume = resumeFlagRef.current;
+      resumeFlagRef.current = false;
+      const commands: QueuedCommand[] = isResume ? [] : commandQueue.flush();
+      if (commands.length === 0 && !isResume)
+        throw new Error("No commands to send");
+
+      const headers = await createRequestHeaders(options.headers);
+      const context = runtime.thread.getModelContext();
+
+      const response = await fetch(
+        isResume ? options.resumeApi! : options.api,
+        {
+          method: "POST",
+          headers,
+          body: JSON.stringify({
+            commands,
+            state: agentStateRef.current,
+            system: context.system,
+            tools: context.tools
+              ? toAISDKTools(getEnabledTools(context.tools))
+              : undefined,
+            ...context.callSettings,
+            ...context.config,
+            ...options.body,
+          }),
+          signal,
+        },
+      );
+
+      options.onResponse?.(response);
+
+      if (!response.ok) {
+        throw new Error(`Status ${response.status}: ${await response.text()}`);
+      }
+
+      if (!response.body) {
+        throw new Error("Response body is null");
+      }
+
+      const stream = response.body
+        .pipeThrough(new DataStreamDecoder())
+        .pipeThrough(
+          new AssistantMessageAccumulator({
+            initialMessage: createInitialMessage({
+              unstable_state:
+                (agentStateRef.current as ReadonlyJSONValue) ?? null,
+            }),
+          }),
+        );
+
+      let markedDelivered = false;
+
+      for await (const chunk of asAsyncIterableStream(stream)) {
+        if (chunk.metadata.unstable_state === agentStateRef.current) continue;
+
+        if (!markedDelivered) {
+          commandQueue.markDelivered();
+          markedDelivered = true;
+        }
+
+        agentStateRef.current = chunk.metadata.unstable_state as T;
+        rerender((prev) => prev + 1);
+      }
+    },
+    onFinish: options.onFinish,
+    onError: (error) => {
+      if (error instanceof Error && error.name === "AbortError") {
+        const cmds = [
+          ...commandQueue.state.inTransit,
+          ...commandQueue.state.queued,
+        ];
+        options.onCancel?.({
+          commands: cmds,
+          updateState: (updater) => {
+            agentStateRef.current = updater(agentStateRef.current);
+            rerender((prev) => prev + 1);
+          },
+        });
+
+        commandQueue.reset();
+      } else {
+        const cmds = [...commandQueue.state.inTransit];
+        options.onError?.(error as Error, {
+          commands: cmds,
+          updateState: (updater) => {
+            agentStateRef.current = updater(agentStateRef.current);
+            rerender((prev) => prev + 1);
+          },
+        });
+        commandQueue.markDelivered();
+      }
+    },
+  });
+
+  // Reactive conversion of agent state + connection metadata → UI state
+  const pendingCommands = useMemo(
+    () => [...commandQueue.state.inTransit, ...commandQueue.state.queued],
+    [commandQueue.state],
+  );
+  const converted = useConvertedState(
+    options.converter,
+    agentStateRef.current,
+    pendingCommands,
+    runManager.isRunning,
+  );
+
+  // Create runtime
+  const runtime = useExternalStoreRuntime({
+    isRunning: converted.isRunning,
+    messages: converted.messages,
+    adapters: options.adapters,
+    onNew: async (message: AppendMessage): Promise<void> => {
+      // Convert AppendMessage to AddMessageCommand
+      const parts: UserMessagePart[] = [];
+
+      for (const content of message.content) {
+        if (content.type === "text") {
+          parts.push({ type: "text", text: content.text });
+        } else if (content.type === "image") {
+          parts.push({ type: "image", image: content.image });
+        }
+      }
+
+      const command: AddMessageCommand = {
+        type: "add-message",
+        message: {
+          role: "user",
+          parts,
+        },
+      };
+
+      commandQueue.enqueue(command);
+    },
+    onCancel: async () => {
+      runManager.cancel();
+      toolInvocations.abort();
+    },
+    onResume: async () => {
+      if (!options.resumeApi)
+        throw new Error("Must pass resumeApi to options to resume runs");
+
+      resumeFlagRef.current = true;
+      runManager.schedule();
+    },
+    onAddToolResult: async (
+      toolOptions: AddToolResultOptions,
+    ): Promise<void> => {
+      const command: AddToolResultCommand = {
+        type: "add-tool-result",
+        toolCallId: toolOptions.toolCallId,
+        result: toolOptions.result as ReadonlyJSONObject,
+        toolName: toolOptions.toolName,
+        isError: toolOptions.isError,
+        ...(toolOptions.artifact && { artifact: toolOptions.artifact }),
+      };
+
+      commandQueue.enqueue(command);
+    },
+    onLoadExternalState: async (state) => {
+      agentStateRef.current = state as T;
+      toolInvocations.reset();
+      rerender((prev) => prev + 1);
+    },
+  });
+
+  const toolInvocations = useToolInvocations({
+    state: converted,
+    getTools: () => runtime.thread.getModelContext().tools,
+    onResult: commandQueue.enqueue,
+  });
+
+  return runtime;
+};
+
+/**
+ * @alpha This is an experimental API that is subject to change.
+ */
+export const useAssistantTransportRuntime = <T,>(
+  options: AssistantTransportOptions<T>,
+): AssistantRuntime => {
+  const runtime = useRemoteThreadListRuntime({
+    runtimeHook: function RuntimeHook() {
+      return useAssistantTransportThreadRuntime(options);
+    },
+    adapter: new InMemoryThreadListAdapter(),
+  });
+  return runtime;
+};

package/src/legacy-runtime/runtime-cores/assistant-transport/useConvertedState.ts

@@ -0,0 +1,18 @@
+import { useMemo } from "react";
+import type {
+  AssistantTransportCommand,
+  AssistantTransportState,
+  AssistantTransportStateConverter,
+} from "./types";
+
+export function useConvertedState<T>(
+  converter: AssistantTransportStateConverter<T>,
+  agentState: T,
+  pendingCommands: AssistantTransportCommand[],
+  isSending: boolean,
+): AssistantTransportState {
+  return useMemo(
+    () => converter(agentState, { pendingCommands, isSending }),
+    [converter, agentState, pendingCommands, isSending],
+  );
+}

package/src/legacy-runtime/runtime-cores/assistant-transport/useLatestRef.ts

@@ -0,0 +1,9 @@
+import { useEffect, useRef } from "react";
+
+export function useLatestRef<T>(value: T) {
+  const ref = useRef(value);
+  useEffect(() => {
+    ref.current = value;
+  }, [value]);
+  return ref as { current: T };
+}

package/src/legacy-runtime/runtime-cores/assistant-transport/useToolInvocations.ts

@@ -0,0 +1,173 @@
+import { useEffect, useRef, useState } from "react";
+import {
+  createAssistantStreamController,
+  ToolCallStreamController,
+  ToolResponse,
+  unstable_toolResultStream,
+  type Tool,
+} from "assistant-stream";
+import type {
+  AssistantTransportCommand,
+  AssistantTransportState,
+} from "./types";
+import {
+  AssistantMetaTransformStream,
+  ReadonlyJSONValue,
+} from "assistant-stream/utils";
+
+const isArgsTextComplete = (argsText: string) => {
+  try {
+    JSON.parse(argsText);
+    return true;
+  } catch {
+    return false;
+  }
+};
+
+type UseToolInvocationsParams = {
+  state: AssistantTransportState;
+  getTools: () => Record<string, Tool> | undefined;
+  onResult: (command: AssistantTransportCommand) => void;
+};
+
+export function useToolInvocations({
+  state,
+  getTools,
+  onResult,
+}: UseToolInvocationsParams) {
+  const lastToolStates = useRef<
+    Record<
+      string,
+      {
+        argsText: string;
+        hasResult: boolean;
+        controller: ToolCallStreamController;
+      }
+    >
+  >({});
+
+  const acRef = useRef<AbortController>(new AbortController());
+  const [controller] = useState(() => {
+    const [stream, controller] = createAssistantStreamController();
+    const transform = unstable_toolResultStream(
+      getTools,
+      () => acRef.current?.signal ?? new AbortController().signal,
+    );
+    stream
+      .pipeThrough(transform)
+      .pipeThrough(new AssistantMetaTransformStream())
+      .pipeTo(
+        new WritableStream({
+          write(chunk) {
+            if (chunk.type === "result") {
+              // the tool call result was already set by the backend
+              if (lastToolStates.current[chunk.meta.toolCallId]?.hasResult)
+                return;
+
+              onResult({
+                type: "add-tool-result",
+                toolCallId: chunk.meta.toolCallId,
+                toolName: chunk.meta.toolName,
+                result: chunk.result,
+                isError: chunk.isError,
+                ...(chunk.artifact && { artifact: chunk.artifact }),
+              });
+            }
+          },
+        }),
+      );
+
+    return controller;
+  });
+
+  const ignoredToolIds = useRef<Set<string>>(new Set());
+  const isInititialState = useRef(true);
+
+  useEffect(() => {
+    if (isInititialState.current) {
+      state.messages.forEach((message) => {
+        message.content.forEach((content) => {
+          if (content.type === "tool-call") {
+            ignoredToolIds.current.add(content.toolCallId);
+          }
+        });
+      });
+      isInititialState.current = false;
+    } else {
+      state.messages.forEach((message) => {
+        message.content.forEach((content) => {
+          if (content.type === "tool-call") {
+            if (ignoredToolIds.current.has(content.toolCallId)) {
+              return;
+            }
+            let lastState = lastToolStates.current[content.toolCallId];
+            if (!lastState) {
+              const toolCallController = controller.addToolCallPart({
+                toolName: content.toolName,
+                toolCallId: content.toolCallId,
+              });
+              lastState = {
+                argsText: "",
+                hasResult: false,
+                controller: toolCallController,
+              };
+              lastToolStates.current[content.toolCallId] = lastState;
+            }
+
+            if (content.argsText !== lastState.argsText) {
+              if (!content.argsText.startsWith(lastState.argsText)) {
+                throw new Error(
+                  `Tool call argsText can only be appended, not updated: ${content.argsText} does not start with ${lastState.argsText}`,
+                );
+              }
+
+              const argsTextDelta = content.argsText.slice(
+                lastState.argsText.length,
+              );
+              lastState.controller.argsText.append(argsTextDelta);
+
+              if (isArgsTextComplete(content.argsText)) {
+                lastState.controller.argsText.close();
+              }
+
+              lastToolStates.current[content.toolCallId] = {
+                argsText: content.argsText,
+                hasResult: lastState.hasResult,
+                controller: lastState.controller,
+              };
+            }
+
+            if (content.result !== undefined && !lastState.hasResult) {
+              lastState.controller.setResponse(
+                new ToolResponse({
+                  result: content.result as ReadonlyJSONValue,
+                  artifact: content.artifact as ReadonlyJSONValue | undefined,
+                  isError: content.isError,
+                }),
+              );
+              lastState.controller.close();
+
+              lastToolStates.current[content.toolCallId] = {
+                hasResult: true,
+                argsText: lastState.argsText,
+                controller: lastState.controller,
+              };
+            }
+          }
+        });
+      });
+    }
+  }, [state]);
+
+  return {
+    reset: () => {
+      acRef.current.abort();
+      acRef.current = new AbortController();
+      isInititialState.current = true;
+    },
+    abort: () => {
+      acRef.current.abort();
+      acRef.current = new AbortController();
+    },
+  };
+}

package/src/legacy-runtime/runtime-cores/assistant-transport/utils.ts

@@ -0,0 +1,42 @@
+import { Tool } from "assistant-stream";
+import { JSONSchema7 } from "json-schema";
+import { z } from "zod";
+
+// Convert tools to AI SDK format
+export const toAISDKTools = (tools: Record<string, Tool>) => {
+  return Object.fromEntries(
+    Object.entries(tools).map(([name, tool]) => [
+      name,
+      {
+        ...(tool.description ? { description: tool.description } : undefined),
+        parameters: (tool.parameters instanceof z.ZodType
+          ? z.toJSONSchema(tool.parameters)
+          : tool.parameters) as JSONSchema7,
+      },
+    ]),
+  );
+};
+
+// Filter enabled tools
+export const getEnabledTools = (tools: Record<string, Tool>) => {
+  return Object.fromEntries(
+    Object.entries(tools).filter(
+      ([, tool]) => !tool.disabled && tool.type !== "backend",
+    ),
+  );
+};
+
+// Create headers for fetch request
+export const createRequestHeaders = async (
+  headersValue:
+    | Record<string, string>
+    | Headers
+    | (() => Promise<Record<string, string> | Headers>),
+): Promise<Headers> => {
+  const resolvedHeaders =
+    typeof headersValue === "function" ? await headersValue() : headersValue;
+
+  const headers = new Headers(resolvedHeaders);
+  headers.set("Content-Type", "application/json");
+  return headers;
+};

package/src/legacy-runtime/runtime-cores/core/BaseThreadRuntimeCore.tsx

@@ -47,6 +47,7 @@ export abstract class BaseThreadRuntimeCore implements ThreadRuntimeCore {
   public abstract resumeRun(config: ResumeRunConfig): void;
   public abstract addToolResult(options: AddToolResultOptions): void;
   public abstract cancelRun(): void;
+  public abstract unstable_loadExternalState(state: any): void;
 
   public get messages() {
     return this.repository.getMessages();

package/src/legacy-runtime/runtime-cores/core/ThreadRuntimeCore.tsx

@@ -63,7 +63,7 @@ export type StartRunConfig = {
 };
 
 export type ResumeRunConfig = StartRunConfig & {
-  stream: (
+  stream?: (
     options: ChatModelRunOptions,
   ) => AsyncGenerator<ChatModelRunResult, void, unknown>;
 };
@@ -122,4 +122,5 @@ export type ThreadRuntimeCore = Readonly<{
   reset(initialMessages?: readonly ThreadMessageLike[]): void;
 
   unstable_on(event: ThreadRuntimeEventType, callback: () => void): Unsubscribe;
+  unstable_loadExternalState: (state: any) => void;
 }>;

package/src/legacy-runtime/runtime-cores/external-store/ExternalStoreAdapter.tsx

@@ -2,6 +2,7 @@ import { AppendMessage, ThreadMessage } from "../../../types";
 import { AttachmentAdapter } from "../adapters/attachment";
 import {
   AddToolResultOptions,
+  ResumeRunConfig,
   StartRunConfig,
   ThreadSuggestion,
 } from "../core/ThreadRuntimeCore";
@@ -63,11 +64,13 @@ type ExternalStoreAdapterBase<T> = {
 
   setMessages?: ((messages: readonly T[]) => void) | undefined;
   onImport?: ((messages: readonly ThreadMessage[]) => void) | undefined;
+  onLoadExternalState?: ((state: any) => void) | undefined;
   onNew: (message: AppendMessage) => Promise<void>;
   onEdit?: ((message: AppendMessage) => Promise<void>) | undefined;
   onReload?: // TODO: remove parentId in 0.8.0
   | ((parentId: string | null, config: StartRunConfig) => Promise<void>)
     | undefined;
+  onResume?: ((config: ResumeRunConfig) => Promise<void>) | undefined;
   onCancel?: (() => Promise<void>) | undefined;
   onAddToolResult?:
     | ((options: AddToolResultOptions) => Promise<void> | void)

package/src/legacy-runtime/runtime-cores/external-store/ExternalStoreThreadRuntimeCore.tsx

@@ -1,5 +1,6 @@
 import {
   AddToolResultOptions,
+  ResumeRunConfig,
   StartRunConfig,
   ThreadSuggestion,
 } from "../core/ThreadRuntimeCore";
@@ -244,8 +245,18 @@ export class ExternalStoreThreadRuntimeCore
     await this._store.onReload(config.parentId, config);
   }
 
-  public async resumeRun(): Promise<void> {
-    throw new Error("Runtime does not support resuming runs.");
+  public async resumeRun(config: ResumeRunConfig): Promise<void> {
+    if (!this._store.onResume)
+      throw new Error("Runtime does not support resuming runs.");
+
+    await this._store.onResume(config);
+  }
+
+  public unstable_loadExternalState(state: any): void {
+    if (!this._store.onLoadExternalState)
+      throw new Error("Runtime does not support importing states.");
+
+    this._store.onLoadExternalState(state);
   }
 
   public cancelRun(): void {

package/src/legacy-runtime/runtime-cores/external-store/ThreadMessageLike.tsx

@@ -133,15 +133,12 @@ export const fromThreadMessageLike = (
                   return {
                     ...commonProps,
                     args: part.args,
-                    argsText: JSON.stringify(part.args),
+                    argsText: part.argsText ?? JSON.stringify(part.args),
                   };
                 }
                 return {
                   ...commonProps,
-                  args:
-                    part.args ??
-                    parsePartialJsonObject(part.argsText ?? "") ??
-                    {},
+                  args: parsePartialJsonObject(part.argsText ?? "") ?? {},
                   argsText: part.argsText ?? "",
                 };
               }

package/src/legacy-runtime/runtime-cores/index.ts

@@ -1,4 +1,5 @@
 export * from "./adapters";
+export * from "./assistant-transport";
 export * from "./core";
 export * from "./external-store";
 export * from "./local";

package/src/legacy-runtime/runtime-cores/local/LocalThreadRuntimeCore.tsx

@@ -166,9 +166,15 @@ export class LocalThreadRuntimeCore
   }
 
   public resumeRun({ stream, ...startConfig }: ResumeRunConfig): Promise<void> {
+    if (!stream)
+      throw new Error("You must pass a stream parameter to resume runs.");
     return this.startRun(startConfig, stream);
   }
 
+  public unstable_loadExternalState(): void {
+    throw new Error("Runtime does not support importing external states.");
+  }
+
   public async startRun(
     { parentId, runConfig }: StartRunConfig,
     runCallback?: ChatModelAdapter["run"],

package/src/legacy-runtime/runtime-cores/remote-thread-list/EMPTY_THREAD_CORE.tsx

@@ -56,6 +56,10 @@ export const EMPTY_THREAD_CORE: ThreadRuntimeCore = {
     return {};
   },
 
+  unstable_loadExternalState() {
+    throw EMPTY_THREAD_ERROR;
+  },
+
   composer: {
     attachments: [],
     attachmentAccept: "*",

package/src/utils/tap-store/tap-api.ts

@@ -34,15 +34,32 @@ class ReadonlyApiHandler<TApi extends ApiObject> implements ProxyHandler<TApi> {
   }
 }
 
-export const tapApi = <TApi extends ApiObject>(api: TApi) => {
+export const tapApi = <TApi extends ApiObject & { getState: () => any }>(
+  api: TApi,
+  options?: {
+    key?: string | undefined;
+  },
+) => {
   const ref = tapRef(() => api);
   tapEffect(() => {
     ref.current = api;
   });
 
-  return tapMemo(
+  const apiProxy = tapMemo(
     () =>
       new Proxy<TApi>({} as TApi, new ReadonlyApiHandler(() => ref.current)),
     [],
   );
+
+  const key = options?.key;
+  const state = api.getState();
+
+  return tapMemo(
+    () => ({
+      key,
+      state,
+      api: apiProxy,
+    }),
+    [state, key],
+  );
 };