Make token update synchronous on FE (#11486)

1. Removing tokenPair internal variable of ApolloFactory. We will relay
on cookieStorage
2. setting the cookie explicitely instead of only relaying on recoil
cookieEffect which is too late

packages/twenty-front/src/modules/apollo/hooks/useApolloFactory.ts

@@ -16,8 +16,8 @@ import { useUpdateEffect } from '~/hooks/useUpdateEffect';
 
 import { currentUserWorkspaceState } from '@/auth/states/currentUserWorkspaceState';
 import { AppPath } from '@/types/AppPath';
-import { ApolloFactory, Options } from '../services/apollo.factory';
 import { isDefined } from 'twenty-shared/utils';
+import { ApolloFactory, Options } from '../services/apollo.factory';
 
 export const useApolloFactory = (options: Partial<Options<any>> = {}) => {
   // eslint-disable-next-line @nx/workspace-no-state-useref
@@ -26,7 +26,7 @@ export const useApolloFactory = (options: Partial<Options<any>> = {}) => {
 
   const navigate = useNavigate();
   const { isMatchingLocation } = useIsMatchingLocation();
-  const [tokenPair, setTokenPair] = useRecoilState(tokenPairState);
+  const setTokenPair = useSetRecoilState(tokenPairState);
   const [currentWorkspace, setCurrentWorkspace] = useRecoilState(
     currentWorkspaceState,
   );
@@ -62,8 +62,6 @@ export const useApolloFactory = (options: Partial<Options<any>> = {}) => {
         },
       },
       connectToDevTools: isDebugMode,
-      // We don't want to re-create the client on token change or it will cause infinite loop
-      initialTokenPair: tokenPair,
       currentWorkspaceMember: currentWorkspaceMember,
       onTokenPairChange: (tokenPair) => {
         setTokenPair(tokenPair);
@@ -104,12 +102,6 @@ export const useApolloFactory = (options: Partial<Options<any>> = {}) => {
     setPreviousUrl,
   ]);
 
-  useUpdateEffect(() => {
-    if (isDefined(apolloRef.current)) {
-      apolloRef.current.updateTokenPair(tokenPair);
-    }
-  }, [tokenPair]);
-
   useUpdateEffect(() => {
     if (isDefined(apolloRef.current)) {
       apolloRef.current.updateWorkspaceMember(currentWorkspaceMember);

packages/twenty-front/src/modules/apollo/services/__tests__/apollo.factory.test.ts

@@ -33,10 +33,6 @@ const mockWorkspaceMember = {
 
 const createMockOptions = (): Options<any> => ({
   uri: 'http://localhost:3000',
-  initialTokenPair: {
-    accessToken: { token: 'mockAccessToken', expiresAt: '' },
-    refreshToken: { token: 'mockRefreshToken', expiresAt: '' },
-  },
   currentWorkspaceMember: mockWorkspaceMember,
   cache: new InMemoryCache(),
   isDebugMode: true,

packages/twenty-front/src/modules/apollo/services/apollo.factory.ts

@@ -18,9 +18,11 @@ import { logDebug } from '~/utils/logDebug';
 
 import { i18n } from '@lingui/core';
 import { GraphQLFormattedError } from 'graphql';
+import { isDefined } from 'twenty-shared/utils';
+import { cookieStorage } from '~/utils/cookie-storage';
+import { isUndefinedOrNull } from '~/utils/isUndefinedOrNull';
 import { ApolloManager } from '../types/apolloManager.interface';
 import { loggerLink } from '../utils/loggerLink';
-import { isDefined } from 'twenty-shared/utils';
 
 const logger = loggerLink(() => 'Twenty');
 
@@ -29,7 +31,6 @@ export interface Options<TCacheShape> extends ApolloClientOptions<TCacheShape> {
   onNetworkError?: (err: Error | ServerParseError | ServerError) => void;
   onTokenPairChange?: (tokenPair: AuthTokenPair) => void;
   onUnauthenticatedError?: () => void;
-  initialTokenPair: AuthTokenPair | null;
   currentWorkspaceMember: CurrentWorkspaceMember | null;
   extraLinks?: ApolloLink[];
   isDebugMode?: boolean;
@@ -37,7 +38,6 @@ export interface Options<TCacheShape> extends ApolloClientOptions<TCacheShape> {
 
 export class ApolloFactory<TCacheShape> implements ApolloManager<TCacheShape> {
   private client: ApolloClient<TCacheShape>;
-  private tokenPair: AuthTokenPair | null = null;
   private currentWorkspaceMember: CurrentWorkspaceMember | null = null;
 
   constructor(opts: Options<TCacheShape>) {
@@ -47,28 +47,45 @@ export class ApolloFactory<TCacheShape> implements ApolloManager<TCacheShape> {
       onNetworkError,
       onTokenPairChange,
       onUnauthenticatedError,
-      initialTokenPair,
       currentWorkspaceMember,
       extraLinks,
       isDebugMode,
       ...options
     } = opts;
 
-    this.tokenPair = initialTokenPair;
     this.currentWorkspaceMember = currentWorkspaceMember;
 
+    const getTokenPair = () => {
+      const stringTokenPair = cookieStorage.getItem('tokenPair');
+      const tokenPair = isDefined(stringTokenPair)
+        ? (JSON.parse(stringTokenPair) as AuthTokenPair)
+        : undefined;
+      return tokenPair;
+    };
+
     const buildApolloLink = (): ApolloLink => {
       const httpLink = createUploadLink({
         uri,
       });
 
       const authLink = setContext(async (_, { headers }) => {
+        const tokenPair = getTokenPair();
+
+        if (isUndefinedOrNull(tokenPair)) {
+          return {
+            headers: {
+              ...headers,
+              ...options.headers,
+            },
+          };
+        }
+
         return {
           headers: {
             ...headers,
             ...options.headers,
-            authorization: this.tokenPair?.accessToken.token
-              ? `Bearer ${this.tokenPair?.accessToken.token}`
+            authorization: tokenPair.accessToken.token
+              ? `Bearer ${tokenPair.accessToken.token}`
               : '',
             ...(this.currentWorkspaceMember?.locale
               ? { 'x-locale': this.currentWorkspaceMember.locale }
@@ -93,7 +110,7 @@ export class ApolloFactory<TCacheShape> implements ApolloManager<TCacheShape> {
             for (const graphQLError of graphQLErrors) {
               if (graphQLError.message === 'Unauthorized') {
                 return fromPromise(
-                  renewToken(uri, this.tokenPair)
+                  renewToken(uri, getTokenPair())
                     .then((tokens) => {
                       if (isDefined(tokens)) {
                         onTokenPairChange?.(tokens);
@@ -108,10 +125,14 @@ export class ApolloFactory<TCacheShape> implements ApolloManager<TCacheShape> {
               switch (graphQLError?.extensions?.code) {
                 case 'UNAUTHENTICATED': {
                   return fromPromise(
-                    renewToken(uri, this.tokenPair)
+                    renewToken(uri, getTokenPair())
                       .then((tokens) => {
                         if (isDefined(tokens)) {
                           onTokenPairChange?.(tokens);
+                          cookieStorage.setItem(
+                            'tokenPair',
+                            JSON.stringify(tokens),
+                          );
                         }
                       })
                       .catch(() => {
@@ -162,10 +183,6 @@ export class ApolloFactory<TCacheShape> implements ApolloManager<TCacheShape> {
     });
   }
 
-  updateTokenPair(tokenPair: AuthTokenPair | null) {
-    this.tokenPair = tokenPair;
-  }
-
   updateWorkspaceMember(workspaceMember: CurrentWorkspaceMember | null) {
     this.currentWorkspaceMember = workspaceMember;
   }

packages/twenty-front/src/modules/apollo/types/apolloManager.interface.ts

@@ -1,10 +1,8 @@
 import { ApolloClient } from '@apollo/client';
 
 import { CurrentWorkspaceMember } from '@/auth/states/currentWorkspaceMemberState';
-import { AuthTokenPair } from '~/generated/graphql';
 
 export interface ApolloManager<TCacheShape> {
   getClient(): ApolloClient<TCacheShape>;
-  updateTokenPair(tokenPair: AuthTokenPair | null): void;
   updateWorkspaceMember(workspaceMember: CurrentWorkspaceMember | null): void;
 }