[pull] next from novuhq:next

apps/dashboard/src/components/ai-elements/chain-of-thought.tsx

@@ -83,7 +83,7 @@ export type ChainOfThoughtStepProps = ComponentProps<'div'> & {
   icon?: IconType | LucideIcon;
   label?: ReactNode;
   description?: ReactNode;
-  status?: 'complete' | 'active' | 'pending';
+  status?: 'complete' | 'active' | 'pending' | 'error';
   collapsible?: boolean;
   defaultOpen?: boolean;
   autoCollapse?: boolean;
@@ -105,7 +105,7 @@ export const ChainOfThoughtStep = memo(
     const [isOpen, setIsOpen] = useState(defaultOpen);
 
     useEffect(() => {
-      if (autoCollapse && status === 'complete') {
+      if (autoCollapse && (status === 'complete' || status === 'error')) {
         setIsOpen(false);
       }
     }, [autoCollapse, status]);
@@ -114,6 +114,7 @@ export const ChainOfThoughtStep = memo(
       complete: 'text-muted-foreground',
       active: 'text-foreground',
       pending: 'text-muted-foreground/50',
+      error: 'text-muted-foreground',
     };
 
     return (

apps/dashboard/src/components/ai-sidekick/ai-chat-context.tsx

@@ -1,5 +1,5 @@
 import { AiAgentTypeEnum, AiMessageRoleEnum, AiResourceTypeEnum } from '@novu/shared';
-import { ChatStatus, DataUIPart, generateId, UIMessage } from 'ai';
+import { ChatStatus, DataUIPart, DynamicToolUIPart, generateId, UIMessage } from 'ai';
 import { createContext, useCallback, useContext, useEffect, useMemo, useRef, useState } from 'react';
 import { useLocation } from 'react-router-dom';
 import { cancelStream } from '@/api/ai';
@@ -12,6 +12,7 @@ import { useFetchLatestAiChat } from '@/hooks/use-fetch-latest-ai-chat';
 import { useKeepAiChanges } from '@/hooks/use-keep-ai-changes';
 import { useRevertMessage } from '@/hooks/use-revert-message';
 import { showErrorToast } from '../primitives/sonner-helpers';
+import { isCancelledToolCall } from './message-utils';
 
 export type ReasoningDataPart = DataUIPart<{ reasoning: { toolCallId: string; text: string } }>;
 
@@ -58,6 +59,57 @@ export type AiChatResourceConfig = {
 
 const AiChatContext = createContext<AiChatContextValue | null>(null);
 
+/**
+ * Strip incomplete tool-call parts and step-start markers from all assistant messages.
+ * Dangling parts are kept in the DB (so toUIMessageStream can match them to the correct
+ * assistant message via the values stream), but hidden from the user in the UI.
+ */
+const cleanupIncompleteToolCalls = <T extends UIMessage>(currentMessages: T[]): T[] => {
+  let changed = false;
+
+  const result = currentMessages.reduce<T[]>((acc, msg) => {
+    if (msg.role !== 'assistant') {
+      acc.push(msg);
+
+      return acc;
+    }
+
+    const cleanedParts = msg.parts.filter((part) => {
+      if (part.type === 'step-start') return false;
+      if (part.type.startsWith('dynamic-tool')) {
+        const tool = part as DynamicToolUIPart;
+        if (isCancelledToolCall(tool)) return false;
+
+        return tool.state === 'output-available' || tool.state === 'output-error';
+      }
+
+      return true;
+    });
+
+    if (cleanedParts.length !== msg.parts.length) {
+      changed = true;
+    }
+
+    const hasContent = cleanedParts.some(
+      (p) =>
+        p.type === 'text' ||
+        (p.type.startsWith('dynamic-tool') &&
+          !isCancelledToolCall(p as DynamicToolUIPart) &&
+          ((p as DynamicToolUIPart).state === 'output-available' || (p as DynamicToolUIPart).state === 'output-error'))
+    );
+
+    if (hasContent) {
+      acc.push(changed ? ({ ...msg, parts: cleanedParts } as T) : msg);
+    } else {
+      changed = true;
+    }
+
+    return acc;
+  }, []);
+
+  return changed ? result : currentMessages;
+};
+
 export function AiChatProvider({ children, config }: { children: React.ReactNode; config: AiChatResourceConfig }) {
   const {
     resourceType,
@@ -79,6 +131,8 @@ export function AiChatProvider({ children, config }: { children: React.ReactNode
   } | null>(null);
   const isMountedRef = useRef(false);
   const hasHandledInitialResumeRef = useRef(false);
+  const isStoppingRef = useRef(false);
+  const skipMessageSyncRef = useRef(false);
   const location = useLocation();
   const { areEnvironmentsInitialLoading, currentEnvironment } = useEnvironment();
 
@@ -113,11 +167,14 @@ export function AiChatProvider({ children, config }: { children: React.ReactNode
           onData({ type: dataType });
         }
       },
-      onFinish: ({ isAbort, isDisconnect, isError }) => {
+      onFinish: ({ isAbort, isDisconnect, isError, messages }) => {
+        setMessages(cleanupIncompleteToolCalls(messages));
+
         if (isAbort || isDisconnect || isError) {
           return;
         }
 
+        skipMessageSyncRef.current = true;
         refetchLatestChat();
       },
     });
@@ -129,11 +186,18 @@ export function AiChatProvider({ children, config }: { children: React.ReactNode
   const isActionPending = isKeepPending || isRevertPending;
 
   useEffect(() => {
-    if (!latestChat || isGenerating) {
+    if (!latestChat || isGenerating || isStoppingRef.current) {
       return;
     }
 
-    setMessages(latestChat.messages as typeof messages);
+    if (skipMessageSyncRef.current) {
+      skipMessageSyncRef.current = false;
+
+      return;
+    }
+
+    const latestChatMessages = latestChat.messages as typeof messages;
+    setMessages(cleanupIncompleteToolCalls(latestChatMessages));
   }, [latestChat, isGenerating, setMessages]);
 
   useEffect(() => {
@@ -175,7 +239,7 @@ export function AiChatProvider({ children, config }: { children: React.ReactNode
 
   const handleSendMessage = useCallback(
     async (message: string) => {
-      const { resourceType, resourceId, isAborted, latestChat, messages } = dataRef.current;
+      const { resourceType, resourceId, latestChat, messages } = dataRef.current;
       const isLastUserMessage = messages.length > 0 && messages[messages.length - 1].role === AiMessageRoleEnum.USER;
 
       const messageToSend = message.trim();
@@ -184,7 +248,7 @@ export function AiChatProvider({ children, config }: { children: React.ReactNode
       if (!latestChat) {
         const newChat = await createAiChat({ resourceType, resourceId });
         sendPrompt({ chatId: newChat._id, prompt: messageToSend });
-      } else if (isLastUserMessage || isAborted) {
+      } else if (isLastUserMessage) {
         const lastUserMessage = messages.filter((m) => m.role === AiMessageRoleEnum.USER).pop();
         sendPrompt({ messageId: lastUserMessage?.id, chatId: latestChat._id, prompt: messageToSend });
       } else if (messageToSend) {
@@ -342,16 +406,15 @@ export function AiChatProvider({ children, config }: { children: React.ReactNode
   }, [firstMessageRevert]);
 
   const handleStop = useCallback(async () => {
+    isStoppingRef.current = true;
+    await stop();
     if (latestChat && currentEnvironment && isGenerating) {
-      cancelStream({ environment: currentEnvironment, chatId: latestChat._id });
-    }
-    stop();
-    refetchLatestChat();
-    const lastUserMessage = messages.filter((m) => m.role === AiMessageRoleEnum.USER).pop();
-    if (lastUserMessage) {
-      setInputText(lastUserMessage.parts.find((p) => p.type === 'text')?.text ?? '');
+      await cancelStream({ environment: currentEnvironment, chatId: latestChat._id });
     }
-  }, [latestChat, currentEnvironment, isGenerating, stop, messages, refetchLatestChat]);
+
+    await refetchLatestChat();
+    isStoppingRef.current = false;
+  }, [latestChat, currentEnvironment, isGenerating, stop, refetchLatestChat]);
 
   const isLoading = isResourceLoading || isFetchingAiChat || areEnvironmentsInitialLoading;
 

apps/dashboard/src/components/ai-sidekick/assistant-message.tsx

@@ -1,10 +1,10 @@
-import { UIMessage } from 'ai';
+import { DynamicToolUIPart, UIMessage } from 'ai';
 import { useMemo } from 'react';
 import { Message } from '../ai-elements/message';
 import { ChatChainOfThought } from './chat-chain-of-thought';
 import { ChatMessageActions } from './chat-message-actions';
 import { StyledMessageResponse } from './chat-message-response';
-import { hasKnownMessageParts } from './message-utils';
+import { hasKnownMessageParts, isCancelledToolCall } from './message-utils';
 
 export const AssistantMessage = ({
   message,
@@ -29,7 +29,7 @@ export const AssistantMessage = ({
 }) => {
   const isAssistantMessageWithKnownParts = useMemo(() => hasKnownMessageParts(message), [message]);
   const hasDynamicToolParts = useMemo(
-    () => message.parts.filter((p) => p.type.startsWith('dynamic-tool')).length > 0,
+    () => message.parts.some((p) => p.type.startsWith('dynamic-tool') && !isCancelledToolCall(p as DynamicToolUIPart)),
     [message]
   );
   const textParts = useMemo(() => {
@@ -48,7 +48,7 @@ export const AssistantMessage = ({
   }
 
   return (
-    <Message from={message.role} key={message.id}>
+    <Message from={message.role}>
       {hasDynamicToolParts && <ChatChainOfThought message={message} />}
       {textParts.map((text, i) => (
         <StyledMessageResponse key={`text-${message.id}-${i}`}>{text}</StyledMessageResponse>