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
This commit is contained in:
Abdul Rahman
2025-11-05 15:20:50 +05:30
committed by GitHub
parent 003b04e9ae
commit cc7343a8f2
7 changed files with 120 additions and 159 deletions
@@ -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>
@@ -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,
@@ -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>
);
@@ -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"
@@ -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>(
@@ -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,
};
};
@@ -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 };
};