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/refresh-token';
+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 './refresh-token';
+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(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(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(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/obtain-tokens.test.ts

@@ -0,0 +1,42 @@
+import { Either, left, right } from '../../../../context/shared/domain/Either';
+import { obtainTokens } from './refresh-token';
+import { RefreshTokenResponse } from '../../../../infra/drive-server/services/auth/auth.types';
+import { driveServerModule } from '../../../../infra/drive-server/drive-server.module';
+import { calls, partialSpyOn } from 'tests/vitest/utils.helper';
+import * as handlers from '../handlers';
+
+describe('obtainTokens', () => {
+  const authRefreshMock = partialSpyOn(driveServerModule.auth, 'refresh');
+  const onUserUnauthorizedMock = partialSpyOn(handlers, 'onUserUnauthorized');
+
+  const refreshResult: Either<Error, RefreshTokenResponse> = right({
+    token: 'abc',
+    newToken: 'xyz',
+    user: {
+      email: '[email protected]',
+      userId: 'user-id',
+      uuid: 'user-uuid',
+    } as RefreshTokenResponse['user'],
+  });
+
+  it('should properly return the user refresh token and the old token if the refresh was successful', async () => {
+    authRefreshMock.mockResolvedValue(refreshResult);
+    const result = await obtainTokens();
+
+    expect(result.isRight()).toBe(true);
+    expect(result.getRight()).toMatchObject({ token: 'abc', newToken: 'xyz' });
+    expect(onUserUnauthorizedMock).not.toHaveBeenCalled();
+  });
+
+  it('should return an error if the token is not obtained', async () => {
+    const mockError = new Error('Refresh request was not successful');
+    const leftResult: Either<Error, RefreshTokenResponse> = left(mockError);
+    authRefreshMock.mockResolvedValue(leftResult);
+
+    const result = await obtainTokens();
+
+    expect(result.isLeft()).toBe(true);
+    expect(result.getLeft()).toStrictEqual(mockError);
+    calls(onUserUnauthorizedMock).toHaveLength(1);
+  });
+});