From dfb796ec5e4dc49580263d627d4e8aac351d77f2 Mon Sep 17 00:00:00 2001
From: Abdul Rahman <81605929+abdulrahmancodes@users.noreply.github.com>
Date: Wed, 15 Oct 2025 14:01:18 +0530
Subject: [PATCH] Improvement AI chat error handling (#15035)

- Fixed AI chat failures when tool calls were present in conversation
history
- Improved error handling and user feedback for AI streaming errors



https://github.com/user-attachments/assets/ca85820f-32c0-4f42-86ee-98f4543e6038
---
 .../AIChatAssistantMessageRenderer.tsx        |  36 +++---
 .../ai/components/AIChatErrorMessage.tsx      |  87 +++++++++++++
 .../modules/ai/components/AIChatMessage.tsx   |  21 +++
 .../src/modules/ai/components/AIChatTab.tsx   |  33 +++--
 .../ai/components/AgentChatProvider.tsx       |  58 +++++++--
 .../components/AgentChatSessionProvider.tsx   |  54 --------
 .../ai/components/ErrorStepRenderer.tsx       |  62 ---------
 .../AIChatErrorMessageWithRecordsContext.tsx  |  14 ++
 .../components/internal/SendMessageButton.tsx |  39 +++++-
 .../SendMessageWithRecordsContextButton.tsx   |   8 +-
 .../modules/ai/contexts/AgentChatContext.ts   |  20 +--
 .../src/modules/ai/hooks/useAgentChat.ts      | 121 ++----------------
 .../ai/hooks/useAgentChatRequestBody.ts       |  45 +++++++
 .../ai/states/agentChatContextState.ts        |  13 --
 .../agent/utils/mapUIMessagePartsToDBParts.ts |   3 +-
 15 files changed, 310 insertions(+), 304 deletions(-)
 create mode 100644 packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx
 delete mode 100644 packages/twenty-front/src/modules/ai/components/AgentChatSessionProvider.tsx
 delete mode 100644 packages/twenty-front/src/modules/ai/components/ErrorStepRenderer.tsx
 create mode 100644 packages/twenty-front/src/modules/ai/components/internal/AIChatErrorMessageWithRecordsContext.tsx
 create mode 100644 packages/twenty-front/src/modules/ai/hooks/useAgentChatRequestBody.ts
 delete mode 100644 packages/twenty-front/src/modules/ai/states/agentChatContextState.ts

diff --git a/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx b/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx
index 577030eaea9..c35c0b91b23 100644
--- a/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx
+++ b/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx
@@ -12,13 +12,13 @@ import {
   type UITools,
 } from 'ai';
 
-const StyledStepsContainer = styled.div`
+const StyledMessagePartsContainer = styled.div`
   display: flex;
   flex-direction: column;
   gap: ${({ theme }) => theme.spacing(1)};
 `;
 
-const StyledDotsIconContainer = styled.div`
+const StyledLoadingIconContainer = styled.div`
   align-items: center;
   border: ${({ theme }) => `1px solid ${theme.border.color.light}`};
   border-radius: ${({ theme }) => theme.border.radius.md};
@@ -27,46 +27,48 @@ const StyledDotsIconContainer = styled.div`
   padding-inline: ${({ theme }) => theme.spacing(1)};
 `;
 
-const StyledDotsIcon = styled(IconDotsVertical)`
+const StyledLoadingIcon = styled(IconDotsVertical)`
   color: ${({ theme }) => theme.font.color.light};
   transform: rotate(90deg);
 `;
 
-const dots = keyframes`
+const streamingDotsAnimation = keyframes`
   0% { content: ''; }
   33% { content: '.'; }
   66% { content: '..'; }
   100% { content: '...'; }
 `;
 
-const StyledToolCallContainer = styled.div`
+const StyledStreamingIndicator = styled.div`
   &::after {
     display: inline-block;
     content: '';
-    animation: ${dots} 750ms steps(3, end) infinite;
+    animation: ${streamingDotsAnimation} 750ms steps(3, end) infinite;
     width: 2ch;
     text-align: left;
   }
 `;
 
-const LoadingDotsIcon = () => {
+const InitialLoadingIndicator = () => {
   const theme = useTheme();
 
   return (
-    <StyledDotsIconContainer>
-      <StyledDotsIcon size={theme.icon.size.xl} />
-    </StyledDotsIconContainer>
+    <StyledLoadingIconContainer>
+      <StyledLoadingIcon size={theme.icon.size.xl} />
+    </StyledLoadingIconContainer>
   );
 };
 
 export const AIChatAssistantMessageRenderer = ({
   messageParts,
   isLastMessageStreaming,
+  hasError,
 }: {
   messageParts: UIMessagePart<UIDataTypes, UITools>[];
   isLastMessageStreaming: boolean;
+  hasError?: boolean;
 }) => {
-  const renderStep = (
+  const renderMessagePart = (
     part: UIMessagePart<UIDataTypes, UITools>,
     index: number,
   ) => {
@@ -99,16 +101,16 @@ export const AIChatAssistantMessageRenderer = ({
     }
   };
 
-  if (!messageParts.length) {
-    return <LoadingDotsIcon />;
+  if (!messageParts.length && !hasError) {
+    return <InitialLoadingIndicator />;
   }
 
   return (
     <div>
-      <StyledStepsContainer>
-        {messageParts.map(renderStep)}
-      </StyledStepsContainer>
-      {isLastMessageStreaming && <StyledToolCallContainer />}
+      <StyledMessagePartsContainer>
+        {messageParts.map(renderMessagePart)}
+      </StyledMessagePartsContainer>
+      {isLastMessageStreaming && !hasError && <StyledStreamingIndicator />}
     </div>
   );
 };
diff --git a/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx b/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx
new file mode 100644
index 00000000000..b41c15db8f3
--- /dev/null
+++ b/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx
@@ -0,0 +1,87 @@
+import { useAgentChatContextOrThrow } from '@/ai/hooks/useAgentChatContextOrThrow';
+import { useAgentChatRequestBody } from '@/ai/hooks/useAgentChatRequestBody';
+import { type ObjectRecord } from '@/object-record/types/ObjectRecord';
+import { useChat } from '@ai-sdk/react';
+import { useTheme } from '@emotion/react';
+import styled from '@emotion/styled';
+import { t } from '@lingui/core/macro';
+import { IconAlertCircle, IconRefresh } from 'twenty-ui/display';
+import { Button } from 'twenty-ui/input';
+
+const StyledErrorContainer = styled.div`
+  align-items: center;
+  background: ${({ theme }) => theme.background.danger};
+  border: 1px solid ${({ theme }) => theme.border.color.danger};
+  border-radius: ${({ theme }) => theme.border.radius.md};
+  display: flex;
+  gap: ${({ theme }) => theme.spacing(2)};
+  padding: ${({ theme }) => theme.spacing(2, 3)};
+`;
+
+const StyledErrorIcon = styled.div`
+  align-items: center;
+  color: ${({ theme }) => theme.color.red};
+  display: flex;
+`;
+
+const StyledErrorContent = styled.div`
+  display: flex;
+  flex-direction: column;
+  gap: ${({ theme }) => theme.spacing(0.5)};
+  flex: 1;
+`;
+
+const StyledErrorTitle = styled.div`
+  color: ${({ theme }) => theme.font.color.primary};
+  font-size: ${({ theme }) => theme.font.size.sm};
+  font-weight: ${({ theme }) => theme.font.weight.medium};
+`;
+
+const StyledErrorMessage = styled.div`
+  color: ${({ theme }) => theme.font.color.secondary};
+  font-size: ${({ theme }) => theme.font.size.xs};
+  word-break: break-word;
+`;
+
+type AIChatErrorMessageProps = {
+  error: Error;
+  records?: ObjectRecord[];
+};
+
+export const AIChatErrorMessage = ({
+  error,
+  records,
+}: AIChatErrorMessageProps) => {
+  const theme = useTheme();
+  const { chat } = useAgentChatContextOrThrow();
+  const { buildRequestBody } = useAgentChatRequestBody();
+  const { regenerate, status } = useChat({ chat });
+
+  const handleRetry = () => {
+    regenerate({
+      body: buildRequestBody(records),
+    });
+  };
+
+  return (
+    <StyledErrorContainer>
+      <StyledErrorIcon>
+        <IconAlertCircle size={theme.icon.size.md} />
+      </StyledErrorIcon>
+      <StyledErrorContent>
+        <StyledErrorTitle>{t`Failed to get response`}</StyledErrorTitle>
+        <StyledErrorMessage>
+          {error.message || t`An error occurred while processing your message`}
+        </StyledErrorMessage>
+      </StyledErrorContent>
+      <Button
+        variant="secondary"
+        size="small"
+        Icon={IconRefresh}
+        onClick={handleRetry}
+        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 54236df66af..bd7acb72395 100644
--- a/packages/twenty-front/src/modules/ai/components/AIChatMessage.tsx
+++ b/packages/twenty-front/src/modules/ai/components/AIChatMessage.tsx
@@ -7,8 +7,13 @@ import { AgentChatFilePreview } from '@/ai/components/internal/AgentChatFilePrev
 import { AgentChatMessageRole } from '@/ai/constants/AgentChatMessageRole';
 
 import { AIChatAssistantMessageRenderer } from '@/ai/components/AIChatAssistantMessageRenderer';
+import { AIChatErrorMessage } from '@/ai/components/AIChatErrorMessage';
+import { AIChatErrorMessageWithRecordsContext } from '@/ai/components/internal/AIChatErrorMessageWithRecordsContext';
 import { type UIMessageWithMetadata } from '@/ai/types/UIMessageWithMetadata';
+import { contextStoreCurrentObjectMetadataItemIdComponentState } from '@/context-store/states/contextStoreCurrentObjectMetadataItemIdComponentState';
 import { LightCopyIconButton } from '@/object-record/record-field/ui/components/LightCopyIconButton';
+import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
+import { isDefined } from 'twenty-shared/utils';
 import { dateLocaleState } from '~/localization/states/dateLocaleState';
 import { beautifyPastDateRelativeToNow } from '~/utils/date-utils';
 const StyledMessageBubble = styled.div<{ isUser?: boolean }>`
@@ -139,13 +144,22 @@ const StyledFilesContainer = styled.div`
 export const AIChatMessage = ({
   message,
   isLastMessageStreaming,
+  error,
 }: {
   message: UIMessageWithMetadata;
   isLastMessageStreaming: boolean;
+  error?: Error | null;
 }) => {
   const theme = useTheme();
   const { localeCatalog } = useRecoilValue(dateLocaleState);
 
+  const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
+    contextStoreCurrentObjectMetadataItemIdComponentState,
+  );
+
+  const showError =
+    isDefined(error) && message.role === AgentChatMessageRole.ASSISTANT;
+
   return (
     <StyledMessageBubble
       key={message.id}
@@ -174,6 +188,7 @@ export const AIChatMessage = ({
             <AIChatAssistantMessageRenderer
               isLastMessageStreaming={isLastMessageStreaming}
               messageParts={message.parts}
+              hasError={showError}
             />
           </StyledMessageText>
           {message.parts.length > 0 && (
@@ -185,6 +200,12 @@ export const AIChatMessage = ({
                 ))}
             </StyledFilesContainer>
           )}
+          {showError &&
+            (contextStoreCurrentObjectMetadataItemId ? (
+              <AIChatErrorMessageWithRecordsContext error={error} />
+            ) : (
+              <AIChatErrorMessage error={error} />
+            ))}
           {message.parts.length > 0 && message.metadata?.createdAt && (
             <StyledMessageFooter className="message-footer">
               <span>
diff --git a/packages/twenty-front/src/modules/ai/components/AIChatTab.tsx b/packages/twenty-front/src/modules/ai/components/AIChatTab.tsx
index 2676cef8fb4..6248bfd3466 100644
--- a/packages/twenty-front/src/modules/ai/components/AIChatTab.tsx
+++ b/packages/twenty-front/src/modules/ai/components/AIChatTab.tsx
@@ -17,7 +17,7 @@ import { SendMessageButton } from '@/ai/components/internal/SendMessageButton';
 import { SendMessageWithRecordsContextButton } from '@/ai/components/internal/SendMessageWithRecordsContextButton';
 import { AI_CHAT_INPUT_ID } from '@/ai/constants/AiChatInputId';
 import { useAIChatFileUpload } from '@/ai/hooks/useAIChatFileUpload';
-import { useAgentChatContextOrThrow } from '@/ai/hooks/useAgentChatContextOrThrow';
+import { useAgentChat } from '@/ai/hooks/useAgentChat';
 import { contextStoreCurrentObjectMetadataItemIdComponentState } from '@/context-store/states/contextStoreCurrentObjectMetadataItemIdComponentState';
 import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
 import { t } from '@lingui/core/macro';
@@ -69,7 +69,8 @@ export const AIChatTab = ({ agentId }: { agentId: string }) => {
     scrollWrapperId,
     messages,
     isStreaming,
-  } = useAgentChatContextOrThrow();
+    error,
+  } = useAgentChat(agentId);
 
   const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
     contextStoreCurrentObjectMetadataItemIdComponentState,
@@ -94,16 +95,20 @@ export const AIChatTab = ({ agentId }: { agentId: string }) => {
         <>
           {messages.length !== 0 && (
             <StyledScrollWrapper componentInstanceId={scrollWrapperId}>
-              {messages.map((message) => (
-                <AIChatMessage
-                  isLastMessageStreaming={
-                    isStreaming &&
-                    message.id === messages[messages.length - 1].id
-                  }
-                  message={message}
-                  key={message.id}
-                />
-              ))}
+              {messages.map((message, index) => {
+                const isLastMessage = index === messages.length - 1;
+                const isLastMessageStreaming = isStreaming && isLastMessage;
+                const shouldShowError = error && isLastMessage;
+
+                return (
+                  <AIChatMessage
+                    isLastMessageStreaming={isLastMessageStreaming}
+                    message={message}
+                    key={message.id}
+                    error={shouldShowError ? error : null}
+                  />
+                );
+              })}
             </StyledScrollWrapper>
           )}
           {messages.length === 0 && <AIChatEmptyState />}
@@ -140,9 +145,9 @@ export const AIChatTab = ({ agentId }: { agentId: string }) => {
               />
               <AgentChatFileUploadButton />
               {contextStoreCurrentObjectMetadataItemId ? (
-                <SendMessageWithRecordsContextButton />
+                <SendMessageWithRecordsContextButton agentId={agentId} />
               ) : (
-                <SendMessageButton />
+                <SendMessageButton agentId={agentId} />
               )}
             </StyledButtonsContainer>
           </StyledInputArea>
diff --git a/packages/twenty-front/src/modules/ai/components/AgentChatProvider.tsx b/packages/twenty-front/src/modules/ai/components/AgentChatProvider.tsx
index 68ccc49edae..45e74526d66 100644
--- a/packages/twenty-front/src/modules/ai/components/AgentChatProvider.tsx
+++ b/packages/twenty-front/src/modules/ai/components/AgentChatProvider.tsx
@@ -1,12 +1,27 @@
 import { AgentChatContext } from '@/ai/contexts/AgentChatContext';
-import { useAgentChat } from '@/ai/hooks/useAgentChat';
 import { useAgentChatData } from '@/ai/hooks/useAgentChatData';
+import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
+import { type UIMessageWithMetadata } from '@/ai/types/UIMessageWithMetadata';
+import { REST_API_BASE_URL } from '@/apollo/constant/rest-api-base-url';
+import { getTokenPair } from '@/apollo/utils/getTokenPair';
 import { currentWorkspaceState } from '@/auth/states/currentWorkspaceState';
 import { useIsFeatureEnabled } from '@/workspace/hooks/useIsFeatureEnabled';
+import { Chat } from '@ai-sdk/react';
+import { DefaultChatTransport } from 'ai';
 import { Suspense } from 'react';
 import { useRecoilValue } from 'recoil';
 import { FeatureFlagKey } from '~/generated/graphql';
 
+const createLoadingChat = () =>
+  new Chat<UIMessageWithMetadata>({
+    transport: new DefaultChatTransport({
+      api: `${REST_API_BASE_URL}/agent-chat/stream`,
+      headers: () => ({}),
+    }),
+    messages: [],
+    id: 'loading',
+  });
+
 const AgentChatProviderContent = ({
   agentId,
   children,
@@ -15,15 +30,24 @@ const AgentChatProviderContent = ({
   children: React.ReactNode;
 }) => {
   const { uiMessages, isLoading } = useAgentChatData(agentId);
-  const chatState = useAgentChat(agentId, uiMessages);
-  const combinedIsLoading = chatState.isLoading || isLoading;
+  const currentAIChatThread = useRecoilValue(currentAIChatThreadState);
+
+  const chatConfig = isLoading
+    ? createLoadingChat()
+    : new Chat<UIMessageWithMetadata>({
+        transport: new DefaultChatTransport({
+          api: `${REST_API_BASE_URL}/agent-chat/stream`,
+          headers: () => ({
+            Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
+          }),
+        }),
+        messages: uiMessages,
+        id: `${currentAIChatThread}-${uiMessages.length}`,
+      });
 
   return (
     <AgentChatContext.Provider
-      value={{
-        ...chatState,
-        isLoading: combinedIsLoading,
-      }}
+      value={{ chat: chatConfig, isLoadingData: isLoading }}
     >
       {children}
     </AgentChatContext.Provider>
@@ -41,14 +65,30 @@ export const AgentChatProvider = ({
 
   if (!isAiEnabled || !agentId) {
     return (
-      <AgentChatContext.Provider value={null}>
+      <AgentChatContext.Provider
+        value={{
+          chat: createLoadingChat(),
+          isLoadingData: false,
+        }}
+      >
         {children}
       </AgentChatContext.Provider>
     );
   }
 
   return (
-    <Suspense fallback={null}>
+    <Suspense
+      fallback={
+        <AgentChatContext.Provider
+          value={{
+            chat: createLoadingChat(),
+            isLoadingData: true,
+          }}
+        >
+          {children}
+        </AgentChatContext.Provider>
+      }
+    >
       <AgentChatProviderContent agentId={agentId}>
         {children}
       </AgentChatProviderContent>
diff --git a/packages/twenty-front/src/modules/ai/components/AgentChatSessionProvider.tsx b/packages/twenty-front/src/modules/ai/components/AgentChatSessionProvider.tsx
deleted file mode 100644
index 8db3836edaa..00000000000
--- a/packages/twenty-front/src/modules/ai/components/AgentChatSessionProvider.tsx
+++ /dev/null
@@ -1,54 +0,0 @@
-import { AgentChatContext } from '@/ai/contexts/AgentChatContext';
-import { useAgentChat } from '@/ai/hooks/useAgentChat';
-import { useAgentChatData } from '@/ai/hooks/useAgentChatData';
-import { currentWorkspaceState } from '@/auth/states/currentWorkspaceState';
-import { Suspense } from 'react';
-import { useRecoilValue } from 'recoil';
-
-const AgentChatProviderContent = ({
-  agentId,
-  children,
-}: {
-  agentId: string;
-  children: React.ReactNode;
-}) => {
-  const { uiMessages, isLoading } = useAgentChatData(agentId);
-  const chatState = useAgentChat(agentId, uiMessages);
-  const combinedIsLoading = chatState.isLoading || isLoading;
-
-  return (
-    <AgentChatContext.Provider
-      value={{
-        ...chatState,
-        isLoading: combinedIsLoading,
-      }}
-    >
-      {children}
-    </AgentChatContext.Provider>
-  );
-};
-
-export const AgentChatProvider = ({
-  children,
-}: {
-  children: React.ReactNode;
-}) => {
-  const currentWorkspace = useRecoilValue(currentWorkspaceState);
-  const agentId = currentWorkspace?.defaultAgent?.id;
-
-  if (!agentId) {
-    return (
-      <AgentChatContext.Provider value={null}>
-        {children}
-      </AgentChatContext.Provider>
-    );
-  }
-
-  return (
-    <Suspense fallback={null}>
-      <AgentChatProviderContent agentId={agentId}>
-        {children}
-      </AgentChatProviderContent>
-    </Suspense>
-  );
-};
diff --git a/packages/twenty-front/src/modules/ai/components/ErrorStepRenderer.tsx b/packages/twenty-front/src/modules/ai/components/ErrorStepRenderer.tsx
deleted file mode 100644
index 628c4a11b95..00000000000
--- a/packages/twenty-front/src/modules/ai/components/ErrorStepRenderer.tsx
+++ /dev/null
@@ -1,62 +0,0 @@
-import { extractErrorMessage } from '@/ai/utils/extractErrorMessage';
-import { useTheme } from '@emotion/react';
-import styled from '@emotion/styled';
-import { t } from '@lingui/core/macro';
-import { IconAlertCircle } from 'twenty-ui/display';
-
-const StyledContainer = styled.div`
-  align-items: flex-start;
-  background-color: ${({ theme }) => theme.color.red10};
-  border: 1px solid ${({ theme }) => theme.color.red20};
-  border-radius: ${({ theme }) => theme.border.radius.md};
-  display: flex;
-  gap: ${({ theme }) => theme.spacing(2)};
-  margin-block: ${({ theme }) => theme.spacing(2)};
-  padding: ${({ theme }) => theme.spacing(3)};
-`;
-
-const StyledIconContainer = styled.div`
-  align-items: center;
-  color: ${({ theme }) => theme.color.red60};
-  display: flex;
-  flex-shrink: 0;
-  justify-content: center;
-`;
-
-const StyledContent = styled.div`
-  flex: 1;
-`;
-
-const StyledTitle = styled.div`
-  font-weight: ${({ theme }) => theme.font.weight.medium};
-  color: ${({ theme }) => theme.color.red80};
-  margin-bottom: ${({ theme }) => theme.spacing(1)};
-`;
-
-const StyledMessage = styled.div`
-  color: ${({ theme }) => theme.color.red70};
-  line-height: ${({ theme }) => theme.text.lineHeight.lg};
-`;
-
-export const ErrorStepRenderer = ({
-  message,
-  error,
-}: {
-  message: string;
-  error?: unknown;
-}) => {
-  const theme = useTheme();
-  const errorMessage = error ? extractErrorMessage(error) : message;
-
-  return (
-    <StyledContainer>
-      <StyledIconContainer>
-        <IconAlertCircle size={theme.icon.size.md} />
-      </StyledIconContainer>
-      <StyledContent>
-        <StyledTitle>{t`Error`}</StyledTitle>
-        <StyledMessage>{errorMessage}</StyledMessage>
-      </StyledContent>
-    </StyledContainer>
-  );
-};
diff --git a/packages/twenty-front/src/modules/ai/components/internal/AIChatErrorMessageWithRecordsContext.tsx b/packages/twenty-front/src/modules/ai/components/internal/AIChatErrorMessageWithRecordsContext.tsx
new file mode 100644
index 00000000000..08a6bc3aa43
--- /dev/null
+++ b/packages/twenty-front/src/modules/ai/components/internal/AIChatErrorMessageWithRecordsContext.tsx
@@ -0,0 +1,14 @@
+import { AIChatErrorMessage } from '@/ai/components/AIChatErrorMessage';
+import { useFindManyRecordsSelectedInContextStore } from '@/context-store/hooks/useFindManyRecordsSelectedInContextStore';
+
+export const AIChatErrorMessageWithRecordsContext = ({
+  error,
+}: {
+  error: Error;
+}) => {
+  const { records } = useFindManyRecordsSelectedInContextStore({
+    limit: 10,
+  });
+
+  return <AIChatErrorMessage error={error} records={records} />;
+};
diff --git a/packages/twenty-front/src/modules/ai/components/internal/SendMessageButton.tsx b/packages/twenty-front/src/modules/ai/components/internal/SendMessageButton.tsx
index 392c4aa617c..cff2d61880e 100644
--- a/packages/twenty-front/src/modules/ai/components/internal/SendMessageButton.tsx
+++ b/packages/twenty-front/src/modules/ai/components/internal/SendMessageButton.tsx
@@ -1,24 +1,57 @@
 import { AI_CHAT_INPUT_ID } from '@/ai/constants/AiChatInputId';
+import { useAgentChat } from '@/ai/hooks/useAgentChat';
 import { useAgentChatContextOrThrow } from '@/ai/hooks/useAgentChatContextOrThrow';
+import { useAgentChatRequestBody } from '@/ai/hooks/useAgentChatRequestBody';
+import { agentChatUploadedFilesState } from '@/ai/states/agentChatUploadedFilesState';
 import { type ObjectRecord } from '@/object-record/types/ObjectRecord';
 import { useHotkeysOnFocusedElement } from '@/ui/utilities/hotkey/hooks/useHotkeysOnFocusedElement';
+import { useChat } from '@ai-sdk/react';
 import { t } from '@lingui/core/macro';
+import { useRecoilState } from 'recoil';
 import { Key } from 'ts-key-enum';
 import { Button } from 'twenty-ui/input';
 
 export const SendMessageButton = ({
+  agentId,
   records,
 }: {
+  agentId: string;
   records?: ObjectRecord[];
 }) => {
-  const { handleSendMessage, isLoading, input } = useAgentChatContextOrThrow();
+  const { input, isLoading, handleInputChange } = useAgentChat(agentId);
+  const { chat } = useAgentChatContextOrThrow();
+  const { buildRequestBody } = useAgentChatRequestBody();
+  const { sendMessage } = useChat({ chat });
+
+  const [agentChatUploadedFiles, setAgentChatUploadedFiles] = useRecoilState(
+    agentChatUploadedFilesState,
+  );
+
+  const handleSendMessage = () => {
+    if (input.trim() === '' || isLoading) {
+      return;
+    }
+
+    sendMessage(
+      {
+        text: input,
+        files: agentChatUploadedFiles,
+      },
+      {
+        body: buildRequestBody(records),
+      },
+    );
+
+    handleInputChange('');
+    setAgentChatUploadedFiles([]);
+  };
 
   useHotkeysOnFocusedElement({
     keys: [Key.Enter],
     callback: (event: KeyboardEvent) => {
       if (!event.ctrlKey && !event.metaKey) {
         event.preventDefault();
-        handleSendMessage(records);
+        handleSendMessage();
       }
     },
     focusId: AI_CHAT_INPUT_ID,
@@ -31,7 +64,7 @@ export const SendMessageButton = ({
   return (
     <Button
       hotkeys={input && !isLoading ? ['⏎'] : undefined}
-      onClick={() => handleSendMessage(records)}
+      onClick={handleSendMessage}
       disabled={!input || isLoading}
       variant="primary"
       accent="blue"
diff --git a/packages/twenty-front/src/modules/ai/components/internal/SendMessageWithRecordsContextButton.tsx b/packages/twenty-front/src/modules/ai/components/internal/SendMessageWithRecordsContextButton.tsx
index 26e7f777323..06319f824d1 100644
--- a/packages/twenty-front/src/modules/ai/components/internal/SendMessageWithRecordsContextButton.tsx
+++ b/packages/twenty-front/src/modules/ai/components/internal/SendMessageWithRecordsContextButton.tsx
@@ -1,10 +1,14 @@
 import { SendMessageButton } from '@/ai/components/internal/SendMessageButton';
 import { useFindManyRecordsSelectedInContextStore } from '@/context-store/hooks/useFindManyRecordsSelectedInContextStore';
 
-export const SendMessageWithRecordsContextButton = () => {
+export const SendMessageWithRecordsContextButton = ({
+  agentId,
+}: {
+  agentId: string;
+}) => {
   const { records } = useFindManyRecordsSelectedInContextStore({
     limit: 10,
   });
 
-  return <SendMessageButton records={records} />;
+  return <SendMessageButton agentId={agentId} records={records} />;
 };
diff --git a/packages/twenty-front/src/modules/ai/contexts/AgentChatContext.ts b/packages/twenty-front/src/modules/ai/contexts/AgentChatContext.ts
index fe93db2c68d..d25a6ee6d93 100644
--- a/packages/twenty-front/src/modules/ai/contexts/AgentChatContext.ts
+++ b/packages/twenty-front/src/modules/ai/contexts/AgentChatContext.ts
@@ -1,24 +1,10 @@
-import { type AIChatObjectMetadataAndRecordContext } from '@/ai/states/agentChatContextState';
 import { type UIMessageWithMetadata } from '@/ai/types/UIMessageWithMetadata';
+import { type Chat } from '@ai-sdk/react';
 import { createContext } from 'react';
-import { type ObjectRecord } from '../../object-record/types/ObjectRecord';
 
 export type AgentChatContextValue = {
-  messages: UIMessageWithMetadata[];
-  isStreaming: boolean;
-  isLoading: boolean;
-  error?: Error;
-
-  input: string;
-  handleInputChange: (value: string) => void;
-
-  handleSendMessage: (records?: ObjectRecord[]) => Promise<void>;
-  handleSetContext: (
-    items: Array<AIChatObjectMetadataAndRecordContext>,
-  ) => Promise<void>;
-
-  scrollWrapperId: string;
-  context: Array<AIChatObjectMetadataAndRecordContext>;
+  chat: Chat<UIMessageWithMetadata>;
+  isLoadingData: boolean;
 };
 
 export const AgentChatContext = createContext<AgentChatContextValue | null>(
diff --git a/packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts b/packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts
index c63d4649b93..67740632622 100644
--- a/packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts
+++ b/packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts
@@ -1,142 +1,39 @@
-import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
 import { useRecoilState, useRecoilValue } from 'recoil';
 
-import {
-  type AIChatObjectMetadataAndRecordContext,
-  agentChatContextState,
-} from '@/ai/states/agentChatContextState';
+import { useAgentChatContextOrThrow } from '@/ai/hooks/useAgentChatContextOrThrow';
 import { agentChatSelectedFilesState } from '@/ai/states/agentChatSelectedFilesState';
-import { agentChatUploadedFilesState } from '@/ai/states/agentChatUploadedFilesState';
 import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
-import { isAgentChatCurrentContextActiveState } from '@/ai/states/isAgentChatCurrentContextActiveState';
-import { type UIMessageWithMetadata } from '@/ai/types/UIMessageWithMetadata';
-import { getTokenPair } from '@/apollo/utils/getTokenPair';
-import { contextStoreCurrentObjectMetadataItemIdComponentState } from '@/context-store/states/contextStoreCurrentObjectMetadataItemIdComponentState';
-import { useGetObjectMetadataItemById } from '@/object-metadata/hooks/useGetObjectMetadataItemById';
-import { type ObjectRecord } from '@/object-record/types/ObjectRecord';
-import { useSnackBar } from '@/ui/feedback/snack-bar-manager/hooks/useSnackBar';
-import { useScrollWrapperHTMLElement } from '@/ui/utilities/scroll/hooks/useScrollWrapperHTMLElement';
 import { useChat } from '@ai-sdk/react';
-import { DefaultChatTransport } from 'ai';
-import { isDefined } from 'twenty-shared/utils';
-import { REST_API_BASE_URL } from '../../apollo/constant/rest-api-base-url';
 import { agentChatInputState } from '../states/agentChatInputState';
 
-export const useAgentChat = (
-  agentId: string,
-  uiMessages: UIMessageWithMetadata[],
-) => {
-  const { getObjectMetadataItemById } = useGetObjectMetadataItemById();
-
-  const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
-    contextStoreCurrentObjectMetadataItemIdComponentState,
-  );
-
-  const isAgentChatCurrentContextActive = useRecoilValue(
-    isAgentChatCurrentContextActiveState,
-  );
+export const useAgentChat = (agentId: string) => {
+  const { chat, isLoadingData } = useAgentChatContextOrThrow();
 
   const agentChatSelectedFiles = useRecoilValue(agentChatSelectedFilesState);
 
-  const [agentChatContext, setAgentChatContext] = useRecoilState(
-    agentChatContextState,
-  );
-
   const currentAIChatThread = useRecoilValue(currentAIChatThreadState);
 
-  const [agentChatUploadedFiles, setAgentChatUploadedFiles] = useRecoilState(
-    agentChatUploadedFilesState,
-  );
-
   const [agentChatInput, setAgentChatInput] =
     useRecoilState(agentChatInputState);
 
   const scrollWrapperId = `scroll-wrapper-ai-chat-${agentId}`;
 
-  const { scrollWrapperHTMLElement } =
-    useScrollWrapperHTMLElement(scrollWrapperId);
-
-  const { enqueueErrorSnackBar } = useSnackBar();
-
-  const { sendMessage, messages, status, error } = useChat({
-    transport: new DefaultChatTransport({
-      api: `${REST_API_BASE_URL}/agent-chat/stream`,
-      headers: () => ({
-        Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
-      }),
-    }),
-    messages: uiMessages,
-    id: `${currentAIChatThread}-${uiMessages.length}`,
-    onError: (error) => {
-      enqueueErrorSnackBar({ message: error.message });
-    },
+  const { messages, status, error } = useChat({
+    chat,
   });
 
   const isStreaming = status === 'streaming';
 
-  const scrollToBottom = () => {
-    scrollWrapperHTMLElement?.scroll({
-      top: scrollWrapperHTMLElement.scrollHeight,
-      behavior: 'smooth',
-    });
-  };
-
   const isLoading =
-    !currentAIChatThread || isStreaming || agentChatSelectedFiles.length > 0;
-
-  const handleSendMessage = async (records?: ObjectRecord[]) => {
-    if (agentChatInput.trim() === '' || isLoading === true) {
-      return;
-    }
-
-    const content = agentChatInput.trim();
-    setAgentChatInput('');
-
-    const recordIdsByObjectMetadataNameSingular = [];
-
-    if (
-      isAgentChatCurrentContextActive === true &&
-      isDefined(records) &&
-      isDefined(contextStoreCurrentObjectMetadataItemId)
-    ) {
-      recordIdsByObjectMetadataNameSingular.push({
-        objectMetadataNameSingular: getObjectMetadataItemById(
-          contextStoreCurrentObjectMetadataItemId,
-        ).nameSingular,
-        recordIds: records.map(({ id }) => id),
-      });
-    }
-
-    sendMessage(
-      {
-        text: content,
-        files: agentChatUploadedFiles,
-      },
-      {
-        body: {
-          threadId: currentAIChatThread,
-          recordIdsByObjectMetadataNameSingular,
-        },
-      },
-    );
-
-    setAgentChatUploadedFiles([]);
-    setTimeout(scrollToBottom, 100);
-  };
-
-  const handleSetContext = async (
-    items: Array<AIChatObjectMetadataAndRecordContext>,
-  ) => {
-    setAgentChatContext(items);
-  };
+    isLoadingData ||
+    !currentAIChatThread ||
+    isStreaming ||
+    agentChatSelectedFiles.length > 0;
 
   return {
     handleInputChange: (value: string) => setAgentChatInput(value),
     messages,
     input: agentChatInput,
-    context: agentChatContext,
-    handleSetContext,
-    handleSendMessage,
     isLoading,
     scrollWrapperId,
     isStreaming,
diff --git a/packages/twenty-front/src/modules/ai/hooks/useAgentChatRequestBody.ts b/packages/twenty-front/src/modules/ai/hooks/useAgentChatRequestBody.ts
new file mode 100644
index 00000000000..8bd0d8ca623
--- /dev/null
+++ b/packages/twenty-front/src/modules/ai/hooks/useAgentChatRequestBody.ts
@@ -0,0 +1,45 @@
+import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
+import { isAgentChatCurrentContextActiveState } from '@/ai/states/isAgentChatCurrentContextActiveState';
+import { contextStoreCurrentObjectMetadataItemIdComponentState } from '@/context-store/states/contextStoreCurrentObjectMetadataItemIdComponentState';
+import { useGetObjectMetadataItemById } from '@/object-metadata/hooks/useGetObjectMetadataItemById';
+import { type ObjectRecord } from '@/object-record/types/ObjectRecord';
+import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
+import { useRecoilValue } from 'recoil';
+import { isDefined } from 'twenty-shared/utils';
+
+export const useAgentChatRequestBody = () => {
+  const currentAIChatThread = useRecoilValue(currentAIChatThreadState);
+  const { getObjectMetadataItemById } = useGetObjectMetadataItemById();
+
+  const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
+    contextStoreCurrentObjectMetadataItemIdComponentState,
+  );
+
+  const isAgentChatCurrentContextActive = useRecoilValue(
+    isAgentChatCurrentContextActiveState,
+  );
+
+  const buildRequestBody = (records?: ObjectRecord[]) => {
+    const recordIdsByObjectMetadataNameSingular = [];
+
+    if (
+      isAgentChatCurrentContextActive === true &&
+      isDefined(records) &&
+      isDefined(contextStoreCurrentObjectMetadataItemId)
+    ) {
+      recordIdsByObjectMetadataNameSingular.push({
+        objectMetadataNameSingular: getObjectMetadataItemById(
+          contextStoreCurrentObjectMetadataItemId,
+        ).nameSingular,
+        recordIds: records.map(({ id }) => id),
+      });
+    }
+
+    return {
+      threadId: currentAIChatThread,
+      recordIdsByObjectMetadataNameSingular,
+    };
+  };
+
+  return { buildRequestBody };
+};
diff --git a/packages/twenty-front/src/modules/ai/states/agentChatContextState.ts b/packages/twenty-front/src/modules/ai/states/agentChatContextState.ts
deleted file mode 100644
index 23e1442f053..00000000000
--- a/packages/twenty-front/src/modules/ai/states/agentChatContextState.ts
+++ /dev/null
@@ -1,13 +0,0 @@
-import { atom } from 'recoil';
-
-export type AIChatObjectMetadataAndRecordContext = {
-  type: 'objectMetadataId' | 'recordId';
-  id: string;
-};
-
-export const agentChatContextState = atom<
-  Array<AIChatObjectMetadataAndRecordContext>
->({
-  key: 'ai/agentChatContextState',
-  default: [],
-});
diff --git a/packages/twenty-server/src/engine/metadata-modules/agent/utils/mapUIMessagePartsToDBParts.ts b/packages/twenty-server/src/engine/metadata-modules/agent/utils/mapUIMessagePartsToDBParts.ts
index bc9e8ced726..8b08df5f5da 100644
--- a/packages/twenty-server/src/engine/metadata-modules/agent/utils/mapUIMessagePartsToDBParts.ts
+++ b/packages/twenty-server/src/engine/metadata-modules/agent/utils/mapUIMessagePartsToDBParts.ts
@@ -64,7 +64,7 @@ export const mapUIMessagePartsToDBParts = (
       default:
         {
           if (isToolPart(part)) {
-            const { toolCallId, input, output, errorText } = part;
+            const { toolCallId, input, output, errorText, state } = part;
 
             return {
               ...basePart,
@@ -72,6 +72,7 @@ export const mapUIMessagePartsToDBParts = (
               toolInput: input,
               toolOutput: output,
               errorMessage: errorText,
+              state,
             };
           }
         }