@assistant-ui/eve 0.0.6 → 0.0.8

package/src/useEveAgentRuntime.ts

@@ -1,6 +1,6 @@
 "use client";
 
-import { useMemo, useRef, useState } from "react";
+import { useEffect, useMemo, useRef, useState } from "react";
 import {
   fromThreadMessageLike,
   generateId,
@@ -35,16 +35,6 @@ const USER_STAGED_STATUS = {
   reason: "unknown",
 } as const;
 
-const truncateThreadMessages = (
-  messages: readonly ThreadMessage[],
-  parentId: string | null,
-) => {
-  if (parentId === null) return [];
-  const parentIndex = messages.findIndex((message) => message.id === parentId);
-  if (parentIndex === -1) return [];
-  return messages.slice(0, parentIndex + 1);
-};
-
 export type UseEveAgentRuntimeOptions = Omit<
   UseEveAgentOptions<EveMessageData>,
   "reducer"
@@ -117,6 +107,7 @@ export const useEveAgentRuntime = (options: UseEveAgentRuntimeOptions = {}) => {
 
     return convertEveMessages(agent.data, {
       isRunning,
+      error: agent.error,
       getCreatedAt: (message) => {
         const existing = createdAtByMessageId.get(message.id);
         if (existing) return existing;
@@ -126,12 +117,38 @@ export const useEveAgentRuntime = (options: UseEveAgentRuntimeOptions = {}) => {
         return createdAt;
       },
     });
-  }, [agent.data, isRunning]);
+  }, [agent.data, agent.error, isRunning]);
 
   const messages = stagedMessages ?? convertedMessages;
   const messagesRef = useRef(messages);
   messagesRef.current = messages;
 
+  useEffect(() => {
+    if (stagedInputsRef.current.size === 0) return;
+    const baseIds = new Set(convertedMessages.map((message) => message.id));
+    const remaining = messagesRef.current.filter(
+      (message) =>
+        stagedInputsRef.current.has(message.id) && !baseIds.has(message.id),
+    );
+    setStagedMessages(
+      remaining.length > 0 ? [...convertedMessages, ...remaining] : null,
+    );
+  }, [convertedMessages]);
+
+  const getStagedRun = (parentId: string | null) => {
+    if (!parentId || !stagedInputsRef.current.has(parentId)) return null;
+    const staged: {
+      message: ThreadMessage;
+      input: { message: AppendMessage; runConfig: AppendMessage["runConfig"] };
+    }[] = [];
+    for (const message of messagesRef.current) {
+      const input = stagedInputsRef.current.get(message.id);
+      if (input) staged.push({ message, input });
+      if (message.id === parentId) break;
+    }
+    return staged;
+  };
+
   const stageUserMessage = (message: AppendMessage) => {
     const threadMessage = fromThreadMessageLike(
       message,
@@ -142,10 +159,9 @@ export const useEveAgentRuntime = (options: UseEveAgentRuntimeOptions = {}) => {
       message,
       runConfig: message.runConfig,
     });
-    setStagedMessages([
-      ...truncateThreadMessages(messagesRef.current, message.parentId),
-      threadMessage,
-    ]);
+    const nextMessages = [...messagesRef.current, threadMessage];
+    messagesRef.current = nextMessages;
+    setStagedMessages(nextMessages);
   };
 
   return useExternalStoreRuntime({
@@ -171,14 +187,30 @@ export const useEveAgentRuntime = (options: UseEveAgentRuntimeOptions = {}) => {
     ...(stagedMessages
       ? {
           onReload: async (parentId: string | null) => {
-            const staged = parentId
-              ? stagedInputsRef.current.get(parentId)
-              : null;
-            if (!staged)
+            const stagedRun = getStagedRun(parentId);
+            if (!stagedRun)
               throw new Error("Runtime does not support reloading messages.");
-            stagedInputsRef.current.delete(parentId!);
-            setStagedMessages(null);
-            await agent.send({ message: getEveMessageContent(staged.message) });
+            for (const { message: stagedMessage, input } of stagedRun) {
+              const previousMessages = messagesRef.current;
+              stagedInputsRef.current.delete(stagedMessage.id);
+              const nextMessages = previousMessages.filter(
+                (message) => message.id !== stagedMessage.id,
+              );
+              messagesRef.current = nextMessages;
+              setStagedMessages(
+                stagedInputsRef.current.size > 0 ? nextMessages : null,
+              );
+              try {
+                await agent.send({
+                  message: getEveMessageContent(input.message),
+                });
+              } catch (error) {
+                stagedInputsRef.current.set(stagedMessage.id, input);
+                messagesRef.current = previousMessages;
+                setStagedMessages(previousMessages);
+                throw error;
+              }
+            }
           },
         }
       : {}),