fix: Improve token schedule to avoid session close

.eslintignore

@@ -31,4 +31,6 @@ npm-debug.log.*
 src/workers/backups/process.ts
 src/workers/backups/*
 src/workers/filesystems/*
-src/test/*
+src/test/*
+
+tests/vitest/*

package.json

@@ -258,7 +258,9 @@
     "webpack-bundle-analyzer": "^4.5.0",
     "webpack-cli": "^4.9.1",
     "webpack-dev-server": "^4.7.1",
-    "webpack-merge": "^5.8.0"
+    "webpack-merge": "^5.8.0",
+    "vitest-mock-extended": "^3.1.0",
+    "ts-essentials": "^10.1.1"
   },
   "dependencies": {
     "@headlessui/react": "^1.4.2",

src/apps/main/auth/deeplink/initialize_current_user.ts

@@ -1,24 +1,22 @@
 import { logger } from '@internxt/drive-desktop-core/build/backend';
-import { updateCredentials } from '../service';
 import { User } from '../../types';
 import { driveServerModule } from '../../../../infra/drive-server/drive-server.module';
 import ConfigStore from '../../config';
+import { updateCredentials } from '../service';
 
 export async function initializeCurrentUser() {
   try {
-    logger.debug({ msg: 'Initializing current user...' });
+    logger.debug({ tag: 'AUTH', msg: 'Initializing current user...' });
+
     const refreshResult = await driveServerModule.auth.refresh();
     if (refreshResult.isLeft()) {
       throw new Error(`Failed to refresh user data: ${refreshResult.getLeft().message}`);
     }
 
     const refreshData = refreshResult.getRight();
-
-    logger.debug({ msg: 'User data refreshed successfully', userId: refreshData.user?.userId });
     updateCredentials(refreshData.token, refreshData.newToken);
 
     const currentUser = ConfigStore.get('userData') as User;
-
     const updatedUser: User = {
       ...currentUser,
       ...refreshData.user,
@@ -27,8 +25,8 @@ export async function initializeCurrentUser() {
 
     ConfigStore.set('userData', updatedUser);
 
-    logger.debug({ msg: 'Current user initialized successfully', userEmail: updatedUser.email });
+    logger.debug({ tag: 'AUTH', msg: 'Current user initialized successfully' });
   } catch (error) {
-    throw logger.error({ msg: 'Failed to initialize current user', error });
+    throw logger.error({ tag: 'AUTH', msg: 'Failed to initialize current user', error });
   }
 }

src/apps/main/auth/handlers.ts

@@ -1,23 +1,11 @@
 import { ipcMain } from 'electron';
 import { logger } from '@internxt/drive-desktop-core/build/backend';
 
-import { AccessResponse } from '../../renderer/pages/Login/service';
 import { applicationOpened } from '../analytics/service';
 import eventBus from '../event-bus';
-import { setupRootFolder } from '../virtual-root-folder/service';
 import { getWidget } from '../windows/widget';
-import { createTokenSchedule } from './refresh-token/refresh-token';
-import {
-  canHisConfigBeRestored,
-  encryptToken,
-  getHeaders,
-  getNewApiHeaders,
-  getUser,
-  logout,
-  obtainToken,
-  setCredentials,
-  tokensArePresent,
-} from './service';
+import { createTokenScheduleWithRetry } from './refresh-token/create-token-schedule-with-retry';
+import { encryptToken, getHeaders, getNewApiHeaders, getUser, logout, obtainToken, tokensArePresent } from './service';
 
 let isLoggedIn = false;
 
@@ -73,7 +61,7 @@ eventBus.on('APP_IS_READY', async (): Promise<void> => {
   if (isLoggedIn) {
     encryptToken();
     applicationOpened();
-    await createTokenSchedule();
+    await createTokenScheduleWithRetry();
     eventBus.emit('USER_LOGGED_IN');
   }
 });

src/apps/main/auth/refresh-token/create-token-schedule-with-retry.test.ts

@@ -0,0 +1,58 @@
+import { createTokenScheduleWithRetry } from './create-token-schedule-with-retry';
+import * as authServiceModule from '../service';
+import { call, calls, partialSpyOn } from 'tests/vitest/utils.helper';
+import { loggerMock } from 'tests/vitest/mocks.helper';
+import { TokenScheduler } from '../../token-scheduler/TokenScheduler';
+import { Job } from 'node-schedule';
+
+describe('createTokenScheduleWithRetry', () => {
+  const obtainTokensMock = partialSpyOn(authServiceModule, 'obtainTokens');
+  const scheduleMock = partialSpyOn(TokenScheduler.prototype, 'schedule');
+
+  const jobMock: Partial<Job> = {
+    cancel: vi.fn(),
+  };
+  const validTokens = ['token-1', 'token-2'];
+
+  beforeEach(() => {
+    scheduleMock.mockReturnValue(jobMock as Job);
+  });
+
+  it('should create token schedule with provided refreshedTokens parameter', async () => {
+    await createTokenScheduleWithRetry({ refreshedTokens: validTokens });
+
+    calls(scheduleMock).toHaveLength(1);
+    calls(obtainTokensMock).toHaveLength(0);
+  });
+
+  it('should create token schedule using obtainStoredTokens when no parameter provided', async () => {
+    obtainTokensMock.mockReturnValue(validTokens);
+
+    await createTokenScheduleWithRetry();
+
+    calls(obtainTokensMock).toHaveLength(1);
+    calls(scheduleMock).toHaveLength(1);
+  });
+
+  it('should attempt to schedule only once when schedule() succeeds immediately', async () => {
+    scheduleMock.mockReturnValue(jobMock);
+
+    await createTokenScheduleWithRetry({ refreshedTokens: validTokens });
+
+    calls(scheduleMock).toHaveLength(1);
+    calls(loggerMock.debug).toHaveLength(0);
+  });
+
+  it('should retry when schedule() fails and succeed on second attempt', async () => {
+    scheduleMock.mockReturnValueOnce(undefined).mockReturnValueOnce(jobMock);
+
+    await createTokenScheduleWithRetry({ refreshedTokens: validTokens });
+
+    calls(scheduleMock).toHaveLength(2);
+    calls(loggerMock.debug).toHaveLength(1);
+    call(loggerMock.debug).toMatchObject({
+      msg: '[TOKEN] Failed to create token schedule, retrying...',
+      tag: 'AUTH',
+    });
+  });
+});

src/apps/main/auth/refresh-token/create-token-schedule-with-retry.ts

@@ -0,0 +1,28 @@
+import { TokenScheduler } from '../../token-scheduler/TokenScheduler';
+import { onUserUnauthorized } from '../handlers';
+import { obtainTokens as obtainStoredTokens } from '../service';
+import { logger } from '@internxt/drive-desktop-core/build/backend';
+import { refreshToken } from './refresh-token';
+
+const CREATE_SCHEDULE_RETRY_LIMIT = 3;
+
+type Props = {
+  refreshedTokens?: Array<string | undefined>;
+};
+
+export async function createTokenScheduleWithRetry({ refreshedTokens }: Props = {}) {
+  const tokens = refreshedTokens || obtainStoredTokens();
+  const tokenScheduler = new TokenScheduler(5, tokens, onUserUnauthorized);
+
+  let attempt = 0;
+  while (attempt < CREATE_SCHEDULE_RETRY_LIMIT) {
+    const schedule = tokenScheduler.schedule(() => refreshToken());
+    if (schedule) break;
+
+    attempt++;
+    logger.debug({
+      tag: 'AUTH',
+      msg: '[TOKEN] Failed to create token schedule, retrying...',
+    });
+  }
+}