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
2025-10-15 10:31:18 +02:00
parent 7b179806a9
commit dfb796ec5e
15 changed files with 310 additions and 304 deletions
@@ -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>
  );
 };
@@ -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>
+  );
+};
@@ -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>
@@ -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>
@@ -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>
@@ -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>
-  );
-};
@@ -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>
-  );
-};
@@ -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} />;
+};
@@ -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"
@@ -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} />;
 };
@@ -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>(
@@ -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,
@@ -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 };
+};
@@ -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: [],
-});
@@ -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,
            };
          }
        }