From ae22e648c1433091ed6228769667886af1bb6407 Mon Sep 17 00:00:00 2001
From: Abdul Rahman <ar5438376@gmail.com>
Date: Mon, 13 Oct 2025 23:28:20 +0530
Subject: [PATCH] Refactor AIChat components to streamline error handling

- Removed isRetrying prop from AIChatErrorMessage and related components to simplify error management.
- Updated AIChatErrorMessage to disable the retry button based on the streaming status instead of isRetrying.
- Enhanced AIChatErrorMessageWithRecordsContext to directly pass error without isRetrying prop, improving clarity in error handling.
---
 .../src/modules/ai/components/AIChatErrorMessage.tsx   |  6 ++----
 .../src/modules/ai/components/AIChatMessage.tsx        |  9 ++-------
 .../src/modules/ai/components/AIChatTab.tsx            |  1 -
 .../internal/AIChatErrorMessageWithRecordsContext.tsx  | 10 +---------
 4 files changed, 5 insertions(+), 21 deletions(-)

diff --git a/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx b/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx
index 32bbde8b715..b41c15db8f3 100644
--- a/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx
+++ b/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx
@@ -45,19 +45,17 @@ const StyledErrorMessage = styled.div`
 
 type AIChatErrorMessageProps = {
   error: Error;
-  isRetrying?: boolean;
   records?: ObjectRecord[];
 };
 
 export const AIChatErrorMessage = ({
   error,
-  isRetrying = false,
   records,
 }: AIChatErrorMessageProps) => {
   const theme = useTheme();
   const { chat } = useAgentChatContextOrThrow();
   const { buildRequestBody } = useAgentChatRequestBody();
-  const { regenerate } = useChat({ chat });
+  const { regenerate, status } = useChat({ chat });
 
   const handleRetry = () => {
     regenerate({
@@ -81,7 +79,7 @@ export const AIChatErrorMessage = ({
         size="small"
         Icon={IconRefresh}
         onClick={handleRetry}
-        disabled={isRetrying}
+        disabled={status === 'streaming'}
         title={t`Retry`}
       />
     </StyledErrorContainer>
diff --git a/packages/twenty-front/src/modules/ai/components/AIChatMessage.tsx b/packages/twenty-front/src/modules/ai/components/AIChatMessage.tsx
index dcb8a8ebe52..bd7acb72395 100644
--- a/packages/twenty-front/src/modules/ai/components/AIChatMessage.tsx
+++ b/packages/twenty-front/src/modules/ai/components/AIChatMessage.tsx
@@ -145,12 +145,10 @@ export const AIChatMessage = ({
   message,
   isLastMessageStreaming,
   error,
-  isRetrying,
 }: {
   message: UIMessageWithMetadata;
   isLastMessageStreaming: boolean;
   error?: Error | null;
-  isRetrying?: boolean;
 }) => {
   const theme = useTheme();
   const { localeCatalog } = useRecoilValue(dateLocaleState);
@@ -204,12 +202,9 @@ export const AIChatMessage = ({
           )}
           {showError &&
             (contextStoreCurrentObjectMetadataItemId ? (
-              <AIChatErrorMessageWithRecordsContext
-                error={error}
-                isRetrying={isRetrying}
-              />
+              <AIChatErrorMessageWithRecordsContext error={error} />
             ) : (
-              <AIChatErrorMessage error={error} isRetrying={isRetrying} />
+              <AIChatErrorMessage error={error} />
             ))}
           {message.parts.length > 0 && message.metadata?.createdAt && (
             <StyledMessageFooter className="message-footer">
diff --git a/packages/twenty-front/src/modules/ai/components/AIChatTab.tsx b/packages/twenty-front/src/modules/ai/components/AIChatTab.tsx
index bc46027397c..6248bfd3466 100644
--- a/packages/twenty-front/src/modules/ai/components/AIChatTab.tsx
+++ b/packages/twenty-front/src/modules/ai/components/AIChatTab.tsx
@@ -106,7 +106,6 @@ export const AIChatTab = ({ agentId }: { agentId: string }) => {
                     message={message}
                     key={message.id}
                     error={shouldShowError ? error : null}
-                    isRetrying={isStreaming}
                   />
                 );
               })}
diff --git a/packages/twenty-front/src/modules/ai/components/internal/AIChatErrorMessageWithRecordsContext.tsx b/packages/twenty-front/src/modules/ai/components/internal/AIChatErrorMessageWithRecordsContext.tsx
index 80094d89559..08a6bc3aa43 100644
--- a/packages/twenty-front/src/modules/ai/components/internal/AIChatErrorMessageWithRecordsContext.tsx
+++ b/packages/twenty-front/src/modules/ai/components/internal/AIChatErrorMessageWithRecordsContext.tsx
@@ -3,20 +3,12 @@ import { useFindManyRecordsSelectedInContextStore } from '@/context-store/hooks/
 
 export const AIChatErrorMessageWithRecordsContext = ({
   error,
-  isRetrying,
 }: {
   error: Error;
-  isRetrying?: boolean;
 }) => {
   const { records } = useFindManyRecordsSelectedInContextStore({
     limit: 10,
   });
 
-  return (
-    <AIChatErrorMessage
-      error={error}
-      isRetrying={isRetrying}
-      records={records}
-    />
-  );
+  return <AIChatErrorMessage error={error} records={records} />;
 };