Add token renewal to agent chat transport (#15727)

Author: abdulrahmancodes

packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts

@@ -1,12 +1,14 @@
 import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
-import { useRecoilState, useRecoilValue } from 'recoil';
+import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil';
 
 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 { renewToken } from '@/auth/services/AuthService';
+import { tokenPairState } from '@/auth/states/tokenPairState';
 import { contextStoreCurrentObjectMetadataItemIdComponentState } from '@/context-store/states/contextStoreCurrentObjectMetadataItemIdComponentState';
 import { useGetObjectMetadataItemById } from '@/object-metadata/hooks/useGetObjectMetadataItemById';
 import { type ObjectRecord } from '@/object-record/types/ObjectRecord';
@@ -15,10 +17,14 @@ import { useChat } from '@ai-sdk/react';
 import { DefaultChatTransport } from 'ai';
 import { type ExtendedUIMessage } from 'twenty-shared/ai';
 import { isDefined } from 'twenty-shared/utils';
+import { REACT_APP_SERVER_BASE_URL } from '~/config';
+import { cookieStorage } from '~/utils/cookie-storage';
 import { REST_API_BASE_URL } from '../../apollo/constant/rest-api-base-url';
 import { agentChatInputState } from '../states/agentChatInputState';
 
 export const useAgentChat = (uiMessages: ExtendedUIMessage[]) => {
+  const setTokenPair = useSetRecoilState(tokenPairState);
+
   const { getObjectMetadataItemById } = useGetObjectMetadataItemById();
 
   const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
@@ -45,12 +51,68 @@ export const useAgentChat = (uiMessages: ExtendedUIMessage[]) => {
   const { scrollWrapperHTMLElement } =
     useScrollWrapperHTMLElement(scrollWrapperId);
 
+  const retryFetchWithRenewedToken = async (
+    input: RequestInfo | URL,
+    init?: RequestInit,
+  ) => {
+    const tokenPair = getTokenPair();
+
+    if (!isDefined(tokenPair)) {
+      return null;
+    }
+
+    try {
+      const renewedTokens = await renewToken(
+        `${REACT_APP_SERVER_BASE_URL}/graphql`,
+        tokenPair,
+      );
+
+      if (!isDefined(renewedTokens)) {
+        setTokenPair(null);
+        return null;
+      }
+
+      const renewedAccessToken =
+        renewedTokens.accessOrWorkspaceAgnosticToken?.token;
+
+      if (!isDefined(renewedAccessToken)) {
+        setTokenPair(null);
+        return null;
+      }
+
+      cookieStorage.setItem('tokenPair', JSON.stringify(renewedTokens));
+      setTokenPair(renewedTokens);
+
+      const updatedHeaders = new Headers(init?.headers ?? {});
+      updatedHeaders.set('Authorization', `Bearer ${renewedAccessToken}`);
+
+      return fetch(input, {
+        ...init,
+        headers: updatedHeaders,
+      });
+    } catch {
+      setTokenPair(null);
+      return null;
+    }
+  };
+
   const { sendMessage, messages, status, error, regenerate } = useChat({
     transport: new DefaultChatTransport({
       api: `${REST_API_BASE_URL}/agent-chat/stream`,
       headers: () => ({
         Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
       }),
+      fetch: async (input, init) => {
+        const response = await fetch(input, init);
+
+        if (response.status !== 401) {
+          return response;
+        }
+
+        const retriedResponse = await retryFetchWithRenewedToken(input, init);
+
+        return retriedResponse ?? response;
+      },
     }),
     messages: uiMessages,
     id: `${currentAIChatThread}-${uiMessages.length}`,