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

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,

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 };
+};