@mastra/react 0.4.3-alpha.5 → 0.4.3-alpha.8

package/dist/index.js

@@ -301,7 +301,7 @@ function signalContentsToUserMessages(contents, metadata) {
 var toUIMessage = ({ chunk, conversation, metadata }) => {
   const result = [...conversation];
   if (chunk.type.startsWith("data-")) {
-    if (chunk.type === "data-user-message" && "data" in chunk && chunk.data?.type === "user-message") {
+    if (chunk.type === "data-user-message" && "data" in chunk && (chunk.data?.type === "user-message" || chunk.data?.type === "user")) {
       const signalId = chunk.data.id;
       if (typeof signalId === "string" && result.some((message) => message.id === signalId)) {
         return result;
@@ -2020,6 +2020,27 @@ function convertSignalDataToBase64String(content) {
 }
 
 // src/agent/hooks.ts
+var extractPendingToolApprovalIdsFromMessages = (messages) => {
+  const pendingToolApprovalIds = /* @__PURE__ */ new Set();
+  for (const message of messages) {
+    const metadata = message.metadata;
+    const metadataSources = [
+      metadata?.pendingToolApprovals,
+      metadata?.requireApprovalMetadata,
+      metadata?.suspendedTools
+    ];
+    for (const source of metadataSources) {
+      if (!source || typeof source !== "object") continue;
+      for (const suspensionData of Object.values(source)) {
+        const toolCallId = suspensionData?.toolCallId;
+        if (typeof toolCallId === "string" && toolCallId.length > 0) {
+          pendingToolApprovalIds.add(toolCallId);
+        }
+      }
+    }
+  }
+  return pendingToolApprovalIds;
+};
 var isThreadSignalUnsupportedError = (error) => {
   const candidate = error;
   const status = candidate?.status;
@@ -2048,17 +2069,24 @@ var useChat = ({
   const _requestContext = useRef(propsRequestContext);
   const _streamAbortRef = useRef(null);
   const _threadSubscriptionAbortRef = useRef(null);
+  const _threadSubscriptionRef = useRef(
+    null
+  );
   const _threadSubscriptionKeyRef = useRef(void 0);
   const _threadSubscriptionPromiseRef = useRef(null);
   const _threadSignalsUnsupportedRef = useRef(false);
   const [messages, setMessages] = useState([]);
   const [toolCallApprovals, setToolCallApprovals] = useState({});
   const [networkToolCallApprovals, setNetworkToolCallApprovals] = useState({});
+  const pendingToolApprovalIdsRef = useRef(/* @__PURE__ */ new Set());
+  const [isAwaitingToolApproval, setIsAwaitingToolApproval] = useState(false);
   const baseClient = useMastraClient();
   const [isRunning, setIsRunning] = useState(false);
   useEffect(() => {
     const formattedMessages = resolveInitialMessages(initialMessages || []);
     setMessages(formattedMessages);
+    pendingToolApprovalIdsRef.current = extractPendingToolApprovalIdsFromMessages(formattedMessages);
+    setIsAwaitingToolApproval(pendingToolApprovalIdsRef.current.size > 0);
     _currentRunId.current = extractRunIdFromMessages(formattedMessages);
   }, [initialMessages]);
   useEffect(() => {
@@ -2114,7 +2142,13 @@ var useChat = ({
     return preview || "Attachment";
   };
   const closeThreadSubscription = useCallback(() => {
-    _threadSubscriptionAbortRef.current?.abort();
+    const subscription = _threadSubscriptionRef.current;
+    if (subscription?.unsubscribe) {
+      subscription.unsubscribe();
+    } else {
+      _threadSubscriptionAbortRef.current?.abort();
+    }
+    _threadSubscriptionRef.current = null;
     _threadSubscriptionAbortRef.current = null;
     _threadSubscriptionKeyRef.current = void 0;
     _threadSubscriptionPromiseRef.current = null;
@@ -2131,7 +2165,17 @@ var useChat = ({
           _currentRunId.current = chunk.runId;
         }
       }
+      if (chunk.type === "tool-call-approval" || chunk.type === "tool-call-suspended") {
+        const toolCallId = chunk.payload?.toolCallId;
+        if (typeof toolCallId === "string") {
+          pendingToolApprovalIdsRef.current.add(toolCallId);
+          setIsAwaitingToolApproval(true);
+        }
+        setIsRunning(false);
+      }
       if (chunk.type === "finish" || chunk.type === "abort" || chunk.type === "error") {
+        pendingToolApprovalIdsRef.current.clear();
+        setIsAwaitingToolApproval(false);
         setIsRunning(false);
       }
       void (onChunk ?? _onChunk.current)?.(chunk);
@@ -2145,7 +2189,7 @@ var useChat = ({
         await _threadSubscriptionPromiseRef.current;
         return;
       }
-      _threadSubscriptionAbortRef.current?.abort();
+      closeThreadSubscription();
       const subscriptionAbort = new AbortController();
       _threadSubscriptionAbortRef.current = subscriptionAbort;
       _threadSubscriptionKeyRef.current = subscriptionKey;
@@ -2155,7 +2199,17 @@ var useChat = ({
       });
       const subscriptionAgent = clientWithAbort.getAgent(agentId);
       _threadSubscriptionPromiseRef.current = subscriptionAgent.subscribeToThread({ resourceId: resourceId2, threadId: threadId2 }).then((response) => {
-        void response.processDataStream({
+        const subscription = response;
+        if (_threadSubscriptionAbortRef.current !== subscriptionAbort) {
+          if (subscription.unsubscribe) {
+            subscription.unsubscribe();
+          } else {
+            subscriptionAbort.abort();
+          }
+          return;
+        }
+        _threadSubscriptionRef.current = subscription;
+        void subscription.processDataStream({
           onChunk: (chunk) => processStreamChunk(chunk)
         }).catch((error) => {
           if (error.name !== "AbortError") {
@@ -2163,6 +2217,9 @@ var useChat = ({
             setIsRunning(false);
           }
         }).finally(() => {
+          if (_threadSubscriptionRef.current === subscription) {
+            _threadSubscriptionRef.current = null;
+          }
           if (_threadSubscriptionAbortRef.current === subscriptionAbort) {
             _threadSubscriptionAbortRef.current = null;
             _threadSubscriptionKeyRef.current = void 0;
@@ -2173,6 +2230,7 @@ var useChat = ({
         if (isThreadSignalUnsupportedError(error)) {
           markThreadSignalsUnsupported();
           if (_threadSubscriptionAbortRef.current === subscriptionAbort) {
+            _threadSubscriptionRef.current = null;
             _threadSubscriptionAbortRef.current = null;
             _threadSubscriptionKeyRef.current = void 0;
             _threadSubscriptionPromiseRef.current = null;
@@ -2187,7 +2245,7 @@ var useChat = ({
       });
       await _threadSubscriptionPromiseRef.current;
     },
-    [agentId, baseClient, markThreadSignalsUnsupported, processStreamChunk]
+    [agentId, baseClient, closeThreadSubscription, markThreadSignalsUnsupported, processStreamChunk]
   );
   useEffect(() => {
     _threadSignalsUnsupportedRef.current = false;
@@ -2408,6 +2466,9 @@ var useChat = ({
           }
         }
       });
+      if (pendingToolApprovalIdsRef.current.size > 0) {
+        setIsRunning(false);
+      }
     } catch (error) {
       onSignalEcho?.(resolvedSignalId);
       if (isThreadSignalUnsupportedError(error)) {
@@ -2471,8 +2532,14 @@ var useChat = ({
   const handleCancelRun = () => {
     _streamAbortRef.current?.abort();
     _streamAbortRef.current = null;
+    const threadSubscription = _threadSubscriptionRef.current;
+    void Promise.resolve(threadSubscription?.abort?.()).catch((error) => {
+      console.error("[useChat] Failed to abort thread subscription", error);
+    });
     closeThreadSubscription();
     setMessages((prev) => finishStreamingAssistantMessage(prev));
+    pendingToolApprovalIdsRef.current.clear();
+    setIsAwaitingToolApproval(false);
     setIsRunning(false);
     _currentRunId.current = void 0;
     _onChunk.current = void 0;
@@ -2488,18 +2555,37 @@ var useChat = ({
     setIsRunning(true);
     setToolCallApprovals((prev) => ({ ...prev, [toolCallId]: { status: "approved" } }));
     const agent = baseClient.getAgent(agentId);
+    if (_threadSubscriptionKeyRef.current && threadId) {
+      try {
+        await agent.sendToolApproval({
+          resourceId: resourceId || agentId,
+          threadId,
+          toolCallId,
+          approved: true,
+          requestContext: _requestContext.current
+        });
+        pendingToolApprovalIdsRef.current.delete(toolCallId);
+        setIsAwaitingToolApproval(pendingToolApprovalIdsRef.current.size > 0);
+        setIsRunning(false);
+      } catch (error) {
+        setToolCallApprovals((prev) => {
+          const next = { ...prev };
+          delete next[toolCallId];
+          return next;
+        });
+        setIsRunning(false);
+        throw error;
+      }
+      return;
+    }
     const response = await agent.approveToolCall({
       runId: currentRunId,
       toolCallId,
       requestContext: _requestContext.current
     });
-    if (_threadSubscriptionKeyRef.current) {
-      return;
-    }
     await response.processDataStream({
       onChunk: async (chunk) => {
-        setMessages((prev) => toUIMessage({ chunk, conversation: prev, metadata: { mode: "stream" } }));
-        void (onChunk ?? _onChunk.current)?.(chunk);
+        await processStreamChunk(chunk, onChunk);
       }
     });
     setIsRunning(false);
@@ -2512,18 +2598,37 @@ var useChat = ({
     setIsRunning(true);
     setToolCallApprovals((prev) => ({ ...prev, [toolCallId]: { status: "declined" } }));
     const agent = baseClient.getAgent(agentId);
+    if (_threadSubscriptionKeyRef.current && threadId) {
+      try {
+        await agent.sendToolApproval({
+          resourceId: resourceId || agentId,
+          threadId,
+          toolCallId,
+          approved: false,
+          requestContext: _requestContext.current
+        });
+        pendingToolApprovalIdsRef.current.delete(toolCallId);
+        setIsAwaitingToolApproval(pendingToolApprovalIdsRef.current.size > 0);
+        setIsRunning(false);
+      } catch (error) {
+        setToolCallApprovals((prev) => {
+          const next = { ...prev };
+          delete next[toolCallId];
+          return next;
+        });
+        setIsRunning(false);
+        throw error;
+      }
+      return;
+    }
     const response = await agent.declineToolCall({
       runId: currentRunId,
       toolCallId,
       requestContext: _requestContext.current
     });
-    if (_threadSubscriptionKeyRef.current) {
-      return;
-    }
     await response.processDataStream({
       onChunk: async (chunk) => {
-        setMessages((prev) => toUIMessage({ chunk, conversation: prev, metadata: { mode: "stream" } }));
-        void (onChunk ?? _onChunk.current)?.(chunk);
+        await processStreamChunk(chunk, onChunk);
       }
     });
     setIsRunning(false);
@@ -2653,6 +2758,7 @@ var useChat = ({
     setMessages,
     sendMessage,
     isRunning,
+    isAwaitingToolApproval,
     messages,
     approveToolCall,
     declineToolCall,