@assistant-ui/react 0.11.6 → 0.11.8

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: "*",