@assistant-ui/react 0.14.28 → 0.15.0

package/src/client/ExternalThread.ts

@@ -1,4 +1,11 @@
-import { useState, useMemo, useEffect, useEffectEvent } from "react";
+import {
+  useState,
+  useMemo,
+  useEffect,
+  useEffectEvent,
+  useCallback,
+  useRef,
+} from "react";
 import { resource, withKey } from "@assistant-ui/tap";
 import {
   type ClientElement,
@@ -10,19 +17,28 @@ import {
 } from "@assistant-ui/store";
 
 import type {
+  AddToolResultOptions,
   AppendMessage,
   Attachment,
+  AttachmentAdapter,
   CreateAttachment,
+  PendingAttachment,
   RespondToToolApprovalOptions,
+  ResumeToolCallOptions,
+  MessagePartStatus,
   ThreadAssistantMessagePart,
   ThreadUserMessagePart,
   ThreadMessage,
+  ToolCallMessagePartStatus,
   ExternalThreadQueueAdapter,
   ExternalThreadBranchAdapter,
 } from "@assistant-ui/core";
+import { ToolResponse } from "assistant-stream";
+import type { ReadonlyJSONValue } from "assistant-stream/utils";
 import type { QueueItemState } from "@assistant-ui/core/store";
 import type { ComposerSendOptions } from "@assistant-ui/core/store";
 import {
+  fileMatchesAccept,
   getThreadMessageText,
   isCreateAttachment,
   resolveToolApprovalResponse,
@@ -38,9 +54,35 @@ export type ExternalThreadMessage = ThreadMessage & {
   id: string;
 };
 
+const COMPLETE_STATUS: MessagePartStatus = Object.freeze({
+  type: "complete",
+});
+
+// Legacy runtime parity (toMessagePartStatus): a tool call without a result
+// carries its message's status; the last part of a streaming message streams.
+const derivePartStatus = (
+  message: ExternalThreadMessage,
+  partIndex: number,
+  part: ThreadAssistantMessagePart | ThreadUserMessagePart,
+): ToolCallMessagePartStatus => {
+  if (message.role !== "assistant" || !message.status) return COMPLETE_STATUS;
+
+  if (part.type === "tool-call") {
+    if (!part.result) return message.status as ToolCallMessagePartStatus;
+    return COMPLETE_STATUS;
+  }
+
+  const isLastPart = partIndex === Math.max(0, message.content.length - 1);
+  if (message.status.type === "requires-action") return COMPLETE_STATUS;
+  return isLastPart ? (message.status as MessagePartStatus) : COMPLETE_STATUS;
+};
+
 export type ExternalThreadProps = {
   messages: readonly ExternalThreadMessage[];
   isRunning?: boolean;
+  isLoading?: boolean | undefined;
+  state?: ReadonlyJSONValue | undefined;
+  extras?: unknown;
   /**
    * Whether sending new messages is currently disabled. When `true`, the
    * thread composer's input remains usable but `send()` is a no-op and
@@ -57,6 +99,12 @@ export type ExternalThreadProps = {
   onReload?: (parentId: string | null) => void;
   onStartRun?: () => void;
   onCancel?: () => void;
+  onResume?: (() => void) | undefined;
+  onAddToolResult?: ((options: AddToolResultOptions) => void) | undefined;
+  /** Callback for resuming a tool call that is waiting for human input. */
+  onResumeToolCall?: ((options: ResumeToolCallOptions) => void) | undefined;
+  onLoadExternalState?: ((state: unknown) => void) | undefined;
+  attachmentAdapter?: AttachmentAdapter | undefined;
   /** Queue adapter for runtimes that support message queuing and steering. */
   queue?: ExternalThreadQueueAdapter;
   /** Branch adapter for runtimes that track sibling variants of messages. */
@@ -68,6 +116,7 @@ export type ExternalThreadProps = {
 type MessageClientProps = {
   message: ExternalThreadMessage;
   index: number;
+  parentId: string | null;
   onEdit?: (message: AppendMessage) => void;
   onReload?: () => void;
   queue?: ExternalThreadQueueAdapter | undefined;
@@ -75,17 +124,24 @@ type MessageClientProps = {
   onRespondToToolApproval?:
     | ((options: RespondToToolApprovalOptions) => void)
     | undefined;
+  onAddToolResult?: ((options: AddToolResultOptions) => void) | undefined;
+  onResumeToolCall?: ((options: ResumeToolCallOptions) => void) | undefined;
+  attachmentAdapter?: AttachmentAdapter | undefined;
 };
 
 // Message Client - minimal implementation
 const useMessageClient = ({
   message,
   index,
+  parentId,
   onEdit,
   onReload,
   queue,
   branches,
   onRespondToToolApproval,
+  onAddToolResult,
+  onResumeToolCall,
+  attachmentAdapter,
 }: MessageClientProps): ClientOutput<"message"> => {
   const [isCopied, setIsCopied] = useState(false);
   const [isHovering, setIsHovering] = useState(false);
@@ -93,7 +149,17 @@ const useMessageClient = ({
 
   const partClients = useClientLookup(
     message.content.map((part, idx) =>
-      withKey(idx, PartResource({ part, onRespondToToolApproval })),
+      withKey(
+        idx,
+        PartResource({
+          part,
+          status: derivePartStatus(message, idx, part),
+          messageId: message.id,
+          onRespondToToolApproval,
+          onAddToolResult,
+          onResumeToolCall,
+        }),
+      ),
     ),
   );
 
@@ -121,7 +187,7 @@ const useMessageClient = ({
     queue?.clear("edit");
     onEdit?.({
       ...msg,
-      parentId: message.id,
+      parentId,
       sourceId: message.id,
     });
     setIsEditing(false);
@@ -137,6 +203,7 @@ const useMessageClient = ({
       onSend: handleSendEdit,
       message,
       queue,
+      attachmentAdapter,
     }),
   );
 
@@ -149,7 +216,7 @@ const useMessageClient = ({
     return {
       ...message,
       attachments: message.attachments ?? [],
-      parentId: null,
+      parentId,
       isLast: false, // Will be set by thread
       branchNumber,
       branchCount,
@@ -162,6 +229,7 @@ const useMessageClient = ({
     };
   }, [
     message,
+    parentId,
     isCopied,
     isHovering,
     index,
@@ -220,28 +288,65 @@ const MessageClient = resource(useMessageClient);
 
 type PartResourceProps = {
   part: ThreadAssistantMessagePart | ThreadUserMessagePart;
+  status: ToolCallMessagePartStatus;
+  messageId: string;
   onRespondToToolApproval?:
     | ((options: RespondToToolApprovalOptions) => void)
     | undefined;
+  onAddToolResult?: ((options: AddToolResultOptions) => void) | undefined;
+  onResumeToolCall?: ((options: ResumeToolCallOptions) => void) | undefined;
 };
 
 // Part Client - minimal implementation
 const usePartResource = ({
   part,
+  status,
+  messageId,
   onRespondToToolApproval,
+  onAddToolResult,
+  onResumeToolCall,
 }: PartResourceProps): ClientOutput<"part"> => {
   const state = useMemo(
     () => ({
       ...part,
-      status: { type: "complete" as const },
+      status,
     }),
-    [part],
+    [part, status],
   );
 
   return {
     getState: () => state,
-    addToolResult: () => {},
-    resumeToolCall: () => {},
+    addToolResult: (result) => {
+      if (!onAddToolResult)
+        throw new Error(
+          "Runtime does not support tool results (onAddToolResult is not set).",
+        );
+      if (part.type !== "tool-call")
+        throw new Error("Tried to add tool result on non-tool message part");
+
+      const response = ToolResponse.toResponse(result);
+      onAddToolResult({
+        messageId,
+        toolName: part.toolName,
+        toolCallId: part.toolCallId,
+        result: response.result as ReadonlyJSONValue,
+        isError: response.isError,
+        ...(response.artifact !== undefined && { artifact: response.artifact }),
+        ...(response.modelContent !== undefined && {
+          modelContent: response.modelContent,
+        }),
+      });
+    },
+    resumeToolCall: (payload) => {
+      if (!onResumeToolCall)
+        throw new Error(
+          "Runtime does not support resuming tool calls (onResumeToolCall is not set).",
+        );
+      if (part.type !== "tool-call")
+        throw new Error("Tried to resume tool call on non-tool message part");
+
+      onResumeToolCall({ toolCallId: part.toolCallId, payload });
+    },
     respondToToolApproval: (response) => {
       if (!onRespondToToolApproval)
         throw new Error("Runtime does not support tool approvals.");
@@ -269,7 +374,7 @@ const PartResource = resource(usePartResource);
 
 type AttachmentResourceProps = {
   attachment: Attachment;
-  onRemove?: () => void;
+  onRemove?: () => void | Promise<void>;
 };
 
 // Attachment Client - minimal implementation
@@ -280,7 +385,7 @@ const useAttachmentResource = ({
   return {
     getState: () => attachment,
     remove: async () => {
-      onRemove?.();
+      await onRemove?.();
     },
   };
 };
@@ -297,6 +402,7 @@ type ComposerClientResourceProps = {
   onSend?: (message: AppendMessage) => void;
   message?: ExternalThreadMessage;
   queue?: ExternalThreadQueueAdapter | undefined;
+  attachmentAdapter?: AttachmentAdapter | undefined;
 };
 
 const useQueueItemClient = ({
@@ -317,6 +423,33 @@ const useQueueItemClient = ({
 
 const QueueItemClient = resource(useQueueItemClient);
 
+const drainAdapterAdd = async (
+  result: ReturnType<AttachmentAdapter["add"]>,
+  upsert: (attachment: Attachment) => void,
+) => {
+  if (Symbol.asyncIterator in result) {
+    for await (const attachment of result) {
+      upsert(attachment);
+    }
+  } else {
+    upsert(await result);
+  }
+};
+
+// State whose setter tracks the latest value in a ref, so imperative
+// call sequences (setText immediately followed by send) observe the write
+// before React re-renders — legacy composer parity.
+const useLiveState = <T>(initial: T) => {
+  const [state, setState] = useState(initial);
+  const ref = useRef(state);
+  const set = useCallback((next: T | ((prev: T) => T)) => {
+    ref.current =
+      typeof next === "function" ? (next as (prev: T) => T)(ref.current) : next;
+    setState(ref.current);
+  }, []);
+  return [state, set, ref] as const;
+};
+
 // Composer Client - minimal implementation
 const useComposerClientResource = ({
   type,
@@ -328,12 +461,19 @@ const useComposerClientResource = ({
   onSend,
   message,
   queue,
+  attachmentAdapter,
 }: ComposerClientResourceProps): ClientOutput<"composer"> => {
-  const [text, setText] = useState("");
-  const [role, setRole] = useState<"user" | "assistant" | "system">("user");
-  const [runConfig, setRunConfig] = useState<Record<string, unknown>>({});
-  const [attachments, setAttachments] = useState<readonly Attachment[]>([]);
-  const [quote, setQuote] = useState<
+  const [text, setText, textRef] = useLiveState("");
+  const [role, setRole, roleRef] = useLiveState<
+    "user" | "assistant" | "system"
+  >("user");
+  const [runConfig, setRunConfig, runConfigRef] = useLiveState<
+    Record<string, unknown>
+  >({});
+  const [attachments, setAttachments, attachmentsRef] = useLiveState<
+    readonly Attachment[]
+  >([]);
+  const [quote, setQuote, quoteRef] = useLiveState<
     { readonly text: string; readonly messageId: string } | undefined
   >(undefined);
 
@@ -359,19 +499,41 @@ const useComposerClientResource = ({
   }, [isEditing]);
 
   const attachmentClients = useClientLookup(
-    attachments.map((attachment, idx) =>
+    attachments.map((attachment) =>
       withKey(
         attachment.id,
         AttachmentResource({
           attachment,
-          onRemove: () => {
-            setAttachments(attachments.filter((_, i) => i !== idx));
+          onRemove: async () => {
+            await attachmentAdapter?.remove(attachment);
+            setAttachments((prev) =>
+              prev.filter((a) => a.id !== attachment.id),
+            );
           },
         }),
       ),
     ),
   );
 
+  const removePendingAttachments = async (removed: readonly Attachment[]) => {
+    if (!attachmentAdapter) return;
+    await Promise.all(
+      removed
+        .filter((a) => a.status.type !== "complete")
+        .map((a) => attachmentAdapter.remove(a)),
+    );
+  };
+
+  const upsertAttachment = (attachment: Attachment) => {
+    setAttachments((prev) => {
+      const idx = prev.findIndex((a) => a.id === attachment.id);
+      if (idx === -1) return [...prev, attachment];
+      const next = [...prev];
+      next[idx] = attachment;
+      return next;
+    });
+  };
+
   const queueItems = queue?.items ?? EMPTY_QUEUE_ITEMS;
   const queueItemClients = useClientLookup(
     queueItems.map((item) =>
@@ -396,7 +558,7 @@ const useComposerClientResource = ({
       isEditing,
       canCancel,
       canSend: isEditing && !isEmpty && !isSendDisabled,
-      attachmentAccept: "*",
+      attachmentAccept: attachmentAdapter?.accept ?? "*",
       isEmpty,
       type,
       dictation: undefined,
@@ -415,6 +577,7 @@ const useComposerClientResource = ({
     attachments.length,
     quote,
     queueItems,
+    attachmentAdapter?.accept,
   ]);
 
   return {
@@ -423,7 +586,24 @@ const useComposerClientResource = ({
     setRole,
     setRunConfig,
     addAttachment: async (fileOrAttachment: File | CreateAttachment) => {
-      if (!isCreateAttachment(fileOrAttachment)) {
+      if (attachmentAdapter) {
+        const file = isCreateAttachment(fileOrAttachment)
+          ? {
+              name: fileOrAttachment.name,
+              type: fileOrAttachment.contentType ?? "",
+            }
+          : { name: fileOrAttachment.name, type: fileOrAttachment.type };
+        if (!fileMatchesAccept(file, attachmentAdapter.accept))
+          throw new Error(
+            `File type ${file.type || "unknown"} is not accepted. Accepted types: ${attachmentAdapter.accept}`,
+          );
+      }
+      if (!isCreateAttachment(fileOrAttachment) && attachmentAdapter) {
+        await drainAdapterAdd(
+          attachmentAdapter.add({ file: fileOrAttachment }),
+          upsertAttachment,
+        );
+      } else if (!isCreateAttachment(fileOrAttachment)) {
         const newAttachment: Attachment = {
           id: Math.random().toString(36).substring(7),
           type: "file",
@@ -433,7 +613,7 @@ const useComposerClientResource = ({
           status: { type: "complete" },
           content: [],
         };
-        setAttachments([...attachments, newAttachment]);
+        setAttachments((prev) => [...prev, newAttachment]);
       } else {
         const newAttachment: Attachment = {
           id: fileOrAttachment.id ?? Math.random().toString(36).substring(7),
@@ -443,11 +623,13 @@ const useComposerClientResource = ({
           content: fileOrAttachment.content,
           status: { type: "complete" },
         };
-        setAttachments([...attachments, newAttachment]);
+        setAttachments((prev) => [...prev, newAttachment]);
       }
     },
     clearAttachments: async () => {
+      const removed = attachmentsRef.current;
       setAttachments([]);
+      await removePendingAttachments(removed);
     },
     attachment: (selector) => {
       if ("id" in selector) {
@@ -456,37 +638,69 @@ const useComposerClientResource = ({
       return attachmentClients.get(selector);
     },
     reset: async () => {
+      const removed = attachmentsRef.current;
       setText("");
       setRole("user");
       setRunConfig({});
       setAttachments([]);
       setQuote(undefined);
+      await removePendingAttachments(removed);
     },
     send: (opts?: ComposerSendOptions) => {
-      if (!state.canSend) return;
-
-      const currentQuote = quote;
-      const composedMessage: AppendMessage = {
-        role,
-        content: text ? [{ type: "text" as const, text }] : [],
-        attachments: attachments as any,
-        createdAt: new Date(),
-        parentId: null,
-        sourceId: null,
-        runConfig,
-        startRun: opts?.startRun,
-        metadata: {
-          custom: { ...(currentQuote ? { quote: currentQuote } : {}) },
-        },
-      };
-      if (queue) {
-        queue.enqueue(composedMessage, { steer: opts?.steer ?? false });
-      } else {
-        onSend?.(composedMessage);
-      }
+      const currentQuote = quoteRef.current;
+      const currentText = textRef.current;
+      const currentAttachments = attachmentsRef.current;
+      const isEmpty = !currentText.trim() && !currentAttachments.length;
+      if (!isEditing || isEmpty || isSendDisabled) return;
+
       setText("");
       setAttachments([]);
       setQuote(undefined);
+
+      const dispatch = (sendAttachments: readonly Attachment[]) => {
+        const composedMessage: AppendMessage = {
+          role: roleRef.current,
+          content: currentText
+            ? [{ type: "text" as const, text: currentText }]
+            : [],
+          attachments: sendAttachments as any,
+          createdAt: new Date(),
+          parentId: null,
+          sourceId: null,
+          runConfig: runConfigRef.current,
+          startRun: opts?.startRun,
+          metadata: {
+            custom: { ...(currentQuote ? { quote: currentQuote } : {}) },
+          },
+        };
+        if (queue) {
+          queue.enqueue(composedMessage, { steer: opts?.steer ?? false });
+        } else {
+          onSend?.(composedMessage);
+        }
+      };
+
+      if (attachmentAdapter && currentAttachments.length > 0) {
+        void Promise.all(
+          currentAttachments.map((attachment) =>
+            attachment.status.type === "complete"
+              ? attachment
+              : attachmentAdapter.send(attachment as PendingAttachment),
+          ),
+        ).then(dispatch, (error) => {
+          // Upload failed: merge the failed send back into the draft.
+          setText((prev) =>
+            currentText && prev
+              ? currentText + "\n" + prev
+              : currentText || prev,
+          );
+          setQuote((prev) => prev ?? currentQuote);
+          setAttachments((prev) => [...currentAttachments, ...prev]);
+          console.error("Failed to send attachments", error);
+        });
+      } else {
+        dispatch(currentAttachments);
+      }
     },
     cancel: onCancel,
     beginEdit: () => {
@@ -503,20 +717,50 @@ const useComposerClientResource = ({
 
 const ComposerClientResource = resource(useComposerClientResource);
 
+const dedupeMessagesById = (messages: readonly ExternalThreadMessage[]) => {
+  const seenIds = new Set<string>();
+  const deduped: ExternalThreadMessage[] = [];
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const message = messages[i]!;
+    if (seenIds.has(message.id)) {
+      console.warn(
+        `ExternalThread: duplicate message id "${message.id}" in the provided messages array; keeping the last occurrence.`,
+      );
+      continue;
+    }
+    seenIds.add(message.id);
+    deduped.push(message);
+  }
+  return deduped.length === messages.length ? messages : deduped.reverse();
+};
+
 // External Thread Client
 const useExternalThread = ({
-  messages,
+  messages: messagesProp,
   isRunning = false,
+  isLoading = false,
+  state: threadState,
+  extras,
   isSendDisabled = false,
   onNew,
   onEdit,
   onReload,
   onStartRun,
   onCancel,
+  onResume,
+  onAddToolResult,
+  onResumeToolCall,
+  onLoadExternalState,
+  attachmentAdapter,
   queue,
   branches,
   onRespondToToolApproval,
 }: ExternalThreadProps): ClientOutput<"thread"> => {
+  const messages = useMemo(
+    () => dedupeMessagesById(messagesProp),
+    [messagesProp],
+  );
+
   const handleReload = (messageId: string) => {
     const messageIndex = messages.findIndex((m) => m.id === messageId);
     if (messageIndex === -1) return;
@@ -531,10 +775,14 @@ const useExternalThread = ({
       const props: MessageClientProps = {
         message: msg,
         index,
+        parentId: index > 0 ? messages[index - 1]!.id : null,
         onReload: () => handleReload(msg.id),
         queue,
         branches,
         onRespondToToolApproval,
+        onAddToolResult,
+        onResumeToolCall,
+        attachmentAdapter,
       };
       if (onEdit) props.onEdit = onEdit;
       return withKey(msg.id, MessageClient(props));
@@ -547,9 +795,25 @@ const useExternalThread = ({
   };
 
   const handleSendNew = (message: AppendMessage) => {
-    onNew?.(message);
+    // The composer does not know the thread; stamp the current head as the
+    // parent (legacy composer parity).
+    onNew?.({ ...message, parentId: messages.at(-1)?.id ?? null });
   };
 
+  const headId = messages.at(-1)?.id ?? null;
+  const composerQueue = useMemo(
+    (): ExternalThreadQueueAdapter | undefined =>
+      queue && {
+        ...queue,
+        enqueue: (message, options) =>
+          queue.enqueue(
+            { ...message, parentId: message.parentId ?? headId },
+            options,
+          ),
+      },
+    [queue, headId],
+  );
+
   const composerClient = useClientResource(
     ComposerClientResource({
       type: "thread",
@@ -558,12 +822,16 @@ const useExternalThread = ({
       isSendDisabled,
       onCancel: handleCancelRun,
       onSend: handleSendNew,
-      queue,
+      queue: composerQueue,
+      attachmentAdapter,
     }),
   );
 
   const hasQueue = !!queue;
   const hasBranches = !!branches;
+  const hasEdit = !!onEdit;
+  const hasReload = !!onReload;
+  const hasAttachments = !!attachmentAdapter;
   const state = useMemo(() => {
     const messageStates = messageClients.state.map((s, idx, arr) => ({
       ...s,
@@ -573,15 +841,15 @@ const useExternalThread = ({
     return {
       isEmpty: messages.length === 0,
       isDisabled: false,
-      isLoading: false,
+      isLoading,
       isRunning,
       capabilities: {
-        edit: false,
+        edit: hasEdit,
         delete: false,
-        reload: false,
+        reload: hasReload,
         cancel: isRunning,
         speech: false,
-        attachments: false,
+        attachments: hasAttachments,
         feedback: false,
         voice: false,
         switchToBranch: hasBranches,
@@ -591,9 +859,9 @@ const useExternalThread = ({
         queue: hasQueue,
       },
       messages: messageStates,
-      state: {},
+      state: threadState ?? {},
       suggestions: [],
-      extras: undefined,
+      extras,
       speech: undefined,
       voice: undefined,
       composer: composerClient.state,
@@ -601,8 +869,14 @@ const useExternalThread = ({
   }, [
     messages,
     isRunning,
+    isLoading,
+    threadState,
+    extras,
     hasQueue,
     hasBranches,
+    hasEdit,
+    hasReload,
+    hasAttachments,
     messageClients.state,
     composerClient.state,
   ]);
@@ -644,8 +918,21 @@ const useExternalThread = ({
     startRun: () => {
       onStartRun?.();
     },
-    resumeRun: () => {},
+    resumeRun: () => {
+      if (!onResume)
+        throw new Error(
+          "Runtime does not support resuming runs (onResume is not set).",
+        );
+      onResume();
+    },
     cancelRun: handleCancelRun,
+    importExternalState: (state: unknown) => {
+      if (!onLoadExternalState)
+        throw new Error(
+          "Runtime does not support importing external states (onLoadExternalState is not set).",
+        );
+      onLoadExternalState(state);
+    },
     getModelContext: () => ({ tools: {}, config: {} }),
     export: () => ({ messages: [] }),
     import: () => {},
@@ -672,10 +959,12 @@ attachTransformScopes(useExternalThread, (scopes, parent) => {
   if (!scopes.threads && parent.threads.source === null) {
     const threadElement = scopes.thread as ClientElement<"thread">;
     scopes.threads = SingleThreadList({ thread: threadElement });
+    // scopes mount in key order; re-declare thread after the threads source it resolves from
+    delete scopes.thread;
     scopes.thread = Derived({
       source: "threads",
       query: { type: "main" },
-      get: (aui) => aui.threads().thread("main"),
+      get: (aui) => aui.threads.thread("main"),
     });
   }
 
@@ -683,14 +972,14 @@ attachTransformScopes(useExternalThread, (scopes, parent) => {
     scopes.threadListItem = Derived({
       source: "threads",
       query: { type: "main" },
-      get: (aui) => aui.threads().item("main"),
+      get: (aui) => aui.threads.item("main"),
     });
   }
 
   scopes.composer ??= Derived({
     source: "thread",
     query: {},
-    get: (aui) => aui.thread().composer(),
+    get: (aui) => aui.thread.composer(),
   });
 
   if (!scopes.modelContext && parent.modelContext.source === null) {