From cc7343a8f29387e945309ca7cba7f980441435f3 Mon Sep 17 00:00:00 2001
From: Abdul Rahman <81605929+abdulrahmancodes@users.noreply.github.com>
Date: Wed, 5 Nov 2025 15:20:50 +0530
Subject: [PATCH] Revert: Agent chat umbrella hook refactoring due to streaming
 issue on thread switch (#15621)

## Summary

Reverts commits afc518a, ae22e64, and 800b5b5 that refactored
`useAgentChat` to address umbrella hook pattern feedback.

## Issues Introduced by Refactoring

The refactoring broke several critical functionalities:

1. **Streaming fails on thread switch** - Messages don't stream properly
when switching between threads
2. **Messages lost on tab close** - When the Ask AI tab is closed, the
request is lost instead of continuing in the background
3. **Blank chat requiring force-reload** - Chats often appear blank and
require switching to another chat to force a reload (closes
[#1771](https://github.com/twentyhq/core-team-issues/issues/1771))

## Root Cause

After extensive debugging, it appears **multiple instances of `useChat`
don't work well together**. The refactored architecture inadvertently
created scenarios where multiple `useChat` instances interfere with each
other.

## Resolution

Reverting to restore functionality. The umbrella hook pattern
optimization needs a different architectural approach that doesn't rely
on multiple `useChat` instances.

## Follow-up

While the umbrella hook feedback is valid, we need to rethink the
implementation strategy:
- Find an alternative to multiple `useChat` instances
- Possibly consolidate chat state management differently
---
 .../ai/components/AIChatErrorMessage.tsx      | 20 +---
 .../src/modules/ai/components/AIChatTab.tsx   |  4 +-
 .../ai/components/AgentChatProvider.tsx       | 59 ++---------
 .../components/internal/SendMessageButton.tsx | 37 +------
 .../modules/ai/contexts/AgentChatContext.ts   | 16 ++-
 .../src/modules/ai/hooks/useAgentChat.ts      | 98 +++++++++++++++++--
 .../ai/hooks/useAgentChatRequestBody.ts       | 45 ---------
 7 files changed, 120 insertions(+), 159 deletions(-)
 delete mode 100644 packages/twenty-front/src/modules/ai/hooks/useAgentChatRequestBody.ts

diff --git a/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx b/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx
index b41c15db8f..e10c880fb1 100644
--- a/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx
+++ b/packages/twenty-front/src/modules/ai/components/AIChatErrorMessage.tsx
@@ -1,7 +1,5 @@
 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';
@@ -46,22 +44,12 @@ const StyledErrorMessage = styled.div`
 type AIChatErrorMessageProps = {
   error: Error;
   records?: ObjectRecord[];
+  isRetrying?: boolean;
 };
 
-export const AIChatErrorMessage = ({
-  error,
-  records,
-}: AIChatErrorMessageProps) => {
+export const AIChatErrorMessage = ({ error }: AIChatErrorMessageProps) => {
   const theme = useTheme();
-  const { chat } = useAgentChatContextOrThrow();
-  const { buildRequestBody } = useAgentChatRequestBody();
-  const { regenerate, status } = useChat({ chat });
-
-  const handleRetry = () => {
-    regenerate({
-      body: buildRequestBody(records),
-    });
-  };
+  const { handleRetry, isStreaming } = useAgentChatContextOrThrow();
 
   return (
     <StyledErrorContainer>
@@ -79,7 +67,7 @@ export const AIChatErrorMessage = ({
         size="small"
         Icon={IconRefresh}
         onClick={handleRetry}
-        disabled={status === 'streaming'}
+        disabled={isStreaming}
         title={t`Retry`}
       />
     </StyledErrorContainer>
diff --git a/packages/twenty-front/src/modules/ai/components/AIChatTab.tsx b/packages/twenty-front/src/modules/ai/components/AIChatTab.tsx
index 362692cfb1..a95b922d7b 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 { useAgentChat } from '@/ai/hooks/useAgentChat';
+import { useAgentChatContextOrThrow } from '@/ai/hooks/useAgentChatContextOrThrow';
 import { contextStoreCurrentObjectMetadataItemIdComponentState } from '@/context-store/states/contextStoreCurrentObjectMetadataItemIdComponentState';
 import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
 import { t } from '@lingui/core/macro';
@@ -70,7 +70,7 @@ export const AIChatTab = () => {
     messages,
     isStreaming,
     error,
-  } = useAgentChat();
+  } = useAgentChatContextOrThrow();
 
   const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
     contextStoreCurrentObjectMetadataItemIdComponentState,
diff --git a/packages/twenty-front/src/modules/ai/components/AgentChatProvider.tsx b/packages/twenty-front/src/modules/ai/components/AgentChatProvider.tsx
index f22d5d9b20..66f04ec9d4 100644
--- a/packages/twenty-front/src/modules/ai/components/AgentChatProvider.tsx
+++ b/packages/twenty-front/src/modules/ai/components/AgentChatProvider.tsx
@@ -1,50 +1,25 @@
 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 { REST_API_BASE_URL } from '@/apollo/constant/rest-api-base-url';
-import { getTokenPair } from '@/apollo/utils/getTokenPair';
 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 { type ExtendedUIMessage } from 'twenty-shared/ai';
 import { FeatureFlagKey } from '~/generated/graphql';
 
-const createLoadingChat = () =>
-  new Chat<ExtendedUIMessage>({
-    transport: new DefaultChatTransport({
-      api: `${REST_API_BASE_URL}/agent-chat/stream`,
-      headers: () => ({}),
-    }),
-    messages: [],
-    id: 'loading',
-  });
-
 const AgentChatProviderContent = ({
   children,
 }: {
   children: React.ReactNode;
 }) => {
   const { uiMessages, isLoading } = useAgentChatData();
-  const currentAIChatThread = useRecoilValue(currentAIChatThreadState);
-
-  const chatConfig = isLoading
-    ? createLoadingChat()
-    : new Chat<ExtendedUIMessage>({
-        transport: new DefaultChatTransport({
-          api: `${REST_API_BASE_URL}/agent-chat/stream`,
-          headers: () => ({
-            Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
-          }),
-        }),
-        messages: uiMessages,
-        id: `${currentAIChatThread}-${uiMessages.length}`,
-      });
+  const chatState = useAgentChat(uiMessages);
+  const combinedIsLoading = chatState.isLoading || isLoading;
 
   return (
     <AgentChatContext.Provider
-      value={{ chat: chatConfig, isLoadingData: isLoading }}
+      value={{
+        ...chatState,
+        isLoading: combinedIsLoading,
+      }}
     >
       {children}
     </AgentChatContext.Provider>
@@ -60,30 +35,14 @@ export const AgentChatProvider = ({
 
   if (!isAiEnabled) {
     return (
-      <AgentChatContext.Provider
-        value={{
-          chat: createLoadingChat(),
-          isLoadingData: false,
-        }}
-      >
+      <AgentChatContext.Provider value={null}>
         {children}
       </AgentChatContext.Provider>
     );
   }
 
   return (
-    <Suspense
-      fallback={
-        <AgentChatContext.Provider
-          value={{
-            chat: createLoadingChat(),
-            isLoadingData: true,
-          }}
-        >
-          {children}
-        </AgentChatContext.Provider>
-      }
-    >
+    <Suspense fallback={null}>
       <AgentChatProviderContent>{children}</AgentChatProviderContent>
     </Suspense>
   );
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 2e820012d7..392c4aa617 100644
--- a/packages/twenty-front/src/modules/ai/components/internal/SendMessageButton.tsx
+++ b/packages/twenty-front/src/modules/ai/components/internal/SendMessageButton.tsx
@@ -1,13 +1,8 @@
 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';
 
@@ -16,40 +11,14 @@ export const SendMessageButton = ({
 }: {
   records?: ObjectRecord[];
 }) => {
-  const { input, isLoading, handleInputChange } = useAgentChat();
-  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([]);
-  };
+  const { handleSendMessage, isLoading, input } = useAgentChatContextOrThrow();
 
   useHotkeysOnFocusedElement({
     keys: [Key.Enter],
     callback: (event: KeyboardEvent) => {
       if (!event.ctrlKey && !event.metaKey) {
         event.preventDefault();
-        handleSendMessage();
+        handleSendMessage(records);
       }
     },
     focusId: AI_CHAT_INPUT_ID,
@@ -62,7 +31,7 @@ export const SendMessageButton = ({
   return (
     <Button
       hotkeys={input && !isLoading ? ['⏎'] : undefined}
-      onClick={handleSendMessage}
+      onClick={() => handleSendMessage(records)}
       disabled={!input || isLoading}
       variant="primary"
       accent="blue"
diff --git a/packages/twenty-front/src/modules/ai/contexts/AgentChatContext.ts b/packages/twenty-front/src/modules/ai/contexts/AgentChatContext.ts
index d5151b78da..5abddd422a 100644
--- a/packages/twenty-front/src/modules/ai/contexts/AgentChatContext.ts
+++ b/packages/twenty-front/src/modules/ai/contexts/AgentChatContext.ts
@@ -1,10 +1,20 @@
-import { type Chat } from '@ai-sdk/react';
 import { createContext } from 'react';
 import { type ExtendedUIMessage } from 'twenty-shared/ai';
+import { type ObjectRecord } from '../../object-record/types/ObjectRecord';
 
 export type AgentChatContextValue = {
-  chat: Chat<ExtendedUIMessage>;
-  isLoadingData: boolean;
+  messages: ExtendedUIMessage[];
+  isStreaming: boolean;
+  isLoading: boolean;
+  error?: Error;
+
+  input: string;
+  handleInputChange: (value: string) => void;
+
+  handleSendMessage: (records?: ObjectRecord[]) => Promise<void>;
+
+  scrollWrapperId: string;
+  handleRetry: () => void;
 };
 
 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 93d7410de2..da15d88fb3 100644
--- a/packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts
+++ b/packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts
@@ -1,42 +1,122 @@
+import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
 import { useRecoilState, useRecoilValue } from 'recoil';
 
-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 { 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 { useScrollWrapperHTMLElement } from '@/ui/utilities/scroll/hooks/useScrollWrapperHTMLElement';
 import { useChat } from '@ai-sdk/react';
+import { DefaultChatTransport } from 'ai';
+import { type ExtendedUIMessage } from 'twenty-shared/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 = () => {
-  const { chat, isLoadingData } = useAgentChatContextOrThrow();
+export const useAgentChat = (uiMessages: ExtendedUIMessage[]) => {
+  const { getObjectMetadataItemById } = useGetObjectMetadataItemById();
+
+  const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
+    contextStoreCurrentObjectMetadataItemIdComponentState,
+  );
+
+  const isAgentChatCurrentContextActive = useRecoilValue(
+    isAgentChatCurrentContextActiveState,
+  );
 
   const agentChatSelectedFiles = useRecoilValue(agentChatSelectedFilesState);
 
   const currentAIChatThread = useRecoilValue(currentAIChatThreadState);
 
+  const [agentChatUploadedFiles, setAgentChatUploadedFiles] = useRecoilState(
+    agentChatUploadedFilesState,
+  );
+
   const [agentChatInput, setAgentChatInput] =
     useRecoilState(agentChatInputState);
 
   const scrollWrapperId = `scroll-wrapper-ai-chat-${currentAIChatThread}`;
 
-  const { messages, status, error } = useChat({
-    chat,
+  const { scrollWrapperHTMLElement } =
+    useScrollWrapperHTMLElement(scrollWrapperId);
+
+  const { sendMessage, messages, status, error, regenerate } = useChat({
+    transport: new DefaultChatTransport({
+      api: `${REST_API_BASE_URL}/agent-chat/stream`,
+      headers: () => ({
+        Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
+      }),
+    }),
+    messages: uiMessages,
+    id: `${currentAIChatThread}-${uiMessages.length}`,
   });
 
   const isStreaming = status === 'streaming';
 
+  const scrollToBottom = () => {
+    scrollWrapperHTMLElement?.scroll({
+      top: scrollWrapperHTMLElement.scrollHeight,
+      behavior: 'smooth',
+    });
+  };
+
   const isLoading =
-    isLoadingData ||
-    !currentAIChatThread ||
-    isStreaming ||
-    agentChatSelectedFiles.length > 0;
+    !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);
+  };
 
   return {
     handleInputChange: (value: string) => setAgentChatInput(value),
     messages,
     input: agentChatInput,
+    handleSendMessage,
     isLoading,
     scrollWrapperId,
     isStreaming,
     error,
+    handleRetry: regenerate,
   };
 };
diff --git a/packages/twenty-front/src/modules/ai/hooks/useAgentChatRequestBody.ts b/packages/twenty-front/src/modules/ai/hooks/useAgentChatRequestBody.ts
deleted file mode 100644
index 8bd0d8ca62..0000000000
--- a/packages/twenty-front/src/modules/ai/hooks/useAgentChatRequestBody.ts
+++ /dev/null
@@ -1,45 +0,0 @@
-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 };
-};