diff --git a/packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts b/packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts index da15d88fb3..b6c92555e2 100644 --- a/packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts +++ b/packages/twenty-front/src/modules/ai/hooks/useAgentChat.ts @@ -1,5 +1,5 @@ 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'; @@ -7,6 +7,8 @@ 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}`,