Refactor: Improve auth flow and service initialization

workspace-mcp-server/WORKSPACE-Context.md

@@ -154,6 +154,11 @@ Choose output format based on use case:
 
 ## 🔍 Error Handling Patterns
 
+### Authentication Errors
+- If any tool returns `{"error":"invalid_request"}`, it likely indicates an expired or invalid session.
+- **Action:** Call `auth.clear` to reset credentials and force a re-login.
+- Inform the user that you are resetting authentication due to an error.
+
 ### Graceful Degradation
 - If a folder doesn't exist, offer to create it
 - If search returns no results, suggest alternatives

workspace-mcp-server/src/__tests__/auth/AuthManager.test.ts

@@ -0,0 +1,79 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { AuthManager } from '../../auth/AuthManager';
+import { OAuthCredentialStorage } from '../../auth/token-storage/oauth-credential-storage';
+import { google } from 'googleapis';
+
+// Mock dependencies
+jest.mock('../../auth/token-storage/oauth-credential-storage');
+jest.mock('googleapis');
+jest.mock('../../utils/logger');
+jest.mock('../../utils/secure-browser-launcher');
+
+describe('AuthManager', () => {
+  let authManager: AuthManager;
+  let mockOAuth2Client: any;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+
+    // Setup mock OAuth2 client
+    mockOAuth2Client = {
+      setCredentials: jest.fn(),
+      generateAuthUrl: jest.fn(),
+      on: jest.fn(),
+      credentials: {}
+    };
+
+    (google.auth.OAuth2 as unknown as jest.Mock).mockReturnValue(mockOAuth2Client);
+
+    authManager = new AuthManager(['scope1']);
+  });
+
+  it('should set up tokens event listener on client creation', async () => {
+    (OAuthCredentialStorage.loadCredentials as jest.Mock).mockResolvedValue({
+        access_token: 'old_token',
+        refresh_token: 'old_refresh',
+        scope: 'scope1'
+    });
+
+    await authManager.getAuthenticatedClient();
+
+    // Verify 'on' was called for 'tokens'
+    expect(mockOAuth2Client.on).toHaveBeenCalledWith('tokens', expect.any(Function));
+  });
+
+  it('should save credentials when tokens event is emitted', async () => {
+    (OAuthCredentialStorage.loadCredentials as jest.Mock).mockResolvedValue({
+        access_token: 'old_token',
+        refresh_token: 'old_refresh',
+        scope: 'scope1'
+    });
+
+    await authManager.getAuthenticatedClient();
+
+    // Get the registered callback
+    const tokensCallback = mockOAuth2Client.on.mock.calls.find((call: any[]) => call[0] === 'tokens')[1];
+    expect(tokensCallback).toBeDefined();
+
+    // Simulate tokens event
+    const newTokens = {
+        access_token: 'new_token',
+        expiry_date: 123456789
+    };
+
+    await tokensCallback(newTokens);
+
+    // Verify saveCredentials was called with merged tokens
+    expect(OAuthCredentialStorage.saveCredentials).toHaveBeenCalledWith({
+        access_token: 'new_token',
+        refresh_token: 'old_refresh', // Should be preserved
+        expiry_date: 123456789,
+        scope: 'scope1'
+    });
+  });
+});

workspace-mcp-server/src/__tests__/services/CalendarService.test.ts

@@ -49,35 +49,16 @@ describe('CalendarService', () => {
 
     // Create CalendarService instance
     calendarService = new CalendarService(mockAuthManager);
+
+    const mockAuthClient = { access_token: 'test-token' };
+    mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient);
   });
 
   afterEach(() => {
     jest.restoreAllMocks();
   });
 
-  describe('initialize', () => {
-    it('should initialize the Calendar API client', async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient);
-
-      await calendarService.initialize();
-
-      expect(mockAuthManager.getAuthenticatedClient).toHaveBeenCalledTimes(1);
-      expect(google.calendar).toHaveBeenCalledWith(
-        expect.objectContaining({
-          version: 'v3',
-          auth: mockAuthClient,
-        })
-      );
-    });
-  });
-
   describe('listCalendars', () => {
-    beforeEach(async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient);
-    });
-
     it('should list all calendars', async () => {
       const mockCalendars = [
         { id: 'primary', summary: 'Primary Calendar' },
@@ -137,8 +118,6 @@ describe('CalendarService', () => {
 
   describe('createEvent', () => {
     beforeEach(async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient);
       mockCalendarAPI.calendarList.list.mockResolvedValue({
         data: {
           items: [{ id: 'primary-calendar-id', primary: true }],
@@ -232,8 +211,6 @@ describe('CalendarService', () => {
 
   describe('listEvents', () => {
     beforeEach(async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient);
       mockCalendarAPI.calendarList.list.mockResolvedValue({
         data: {
           items: [{ id: 'primary-calendar-id', primary: true }],

workspace-mcp-server/src/__tests__/services/ChatService.test.ts

@@ -26,9 +26,6 @@ describe('ChatService', () => {
     // Create mock AuthManager
     mockAuthManager = {
       getAuthenticatedClient: jest.fn(),
-      loadSavedCredentialsIfExist: jest.fn(),
-      saveCredentials: jest.fn(),
-      authorize: jest.fn(),
     } as any;
 
     // Create mock Chat API
@@ -60,42 +57,16 @@ describe('ChatService', () => {
 
     // Create ChatService instance
     chatService = new ChatService(mockAuthManager);
+
+    const mockAuthClient = { access_token: 'test-token' };
+    mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient as any);
   });
 
   afterEach(() => {
     jest.restoreAllMocks();
   });
 
-  describe('initialize', () => {
-    it('should initialize Chat and People API clients', async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient as any);
-
-      await chatService.initialize();
-
-      expect(mockAuthManager.getAuthenticatedClient).toHaveBeenCalledTimes(1);
-      expect(google.chat).toHaveBeenCalledWith(
-        expect.objectContaining({
-          version: 'v1',
-          auth: mockAuthClient,
-        })
-      );
-      expect(google.people).toHaveBeenCalledWith(
-        expect.objectContaining({
-          version: 'v1',
-          auth: mockAuthClient,
-        })
-      );
-    });
-  });
-
   describe('listSpaces', () => {
-    beforeEach(async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient as any);
-      await chatService.initialize();
-    });
-
     it('should list all chat spaces', async () => {
       const mockSpaces = [
         { name: 'spaces/space1', displayName: 'Team Chat' },
@@ -139,12 +110,6 @@ describe('ChatService', () => {
   });
 
   describe('sendMessage', () => {
-    beforeEach(async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient as any);
-      await chatService.initialize();
-    });
-
     it('should send a message to a space', async () => {
       const mockResponse = {
         name: 'spaces/space1/messages/msg1',
@@ -186,12 +151,6 @@ describe('ChatService', () => {
   });
 
   describe('findSpaceByName', () => {
-    beforeEach(async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient as any);
-      await chatService.initialize();
-    });
-
     it('should find spaces by display name', async () => {
       const mockSpaces = [
         { name: 'spaces/space1', displayName: 'Team Chat' },
@@ -263,12 +222,6 @@ describe('ChatService', () => {
   });
 
   describe('getMessages', () => {
-    beforeEach(async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient as any);
-      await chatService.initialize();
-    });
-
     it('should list messages from a space', async () => {
       const mockMessages = [
         { name: 'spaces/space1/messages/msg1', text: 'Hello' },
@@ -404,12 +357,6 @@ describe('ChatService', () => {
   });
 
   describe('sendDm', () => {
-    beforeEach(async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient as any);
-      await chatService.initialize();
-    });
-
     it('should send a direct message to a user', async () => {
       const mockSpace = {
         name: 'spaces/dm123',
@@ -477,12 +424,6 @@ describe('ChatService', () => {
   });
 
   describe('findDmByEmail', () => {
-    beforeEach(async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient as any);
-      await chatService.initialize();
-    });
-
     it('should find a DM space by user email using spaces.setup', async () => {
       const mockSpace = {
         name: 'spaces/dm123',
@@ -528,12 +469,6 @@ describe('ChatService', () => {
   });
 
   describe('createSpace', () => {
-    beforeEach(async () => {
-      const mockAuthClient = { access_token: 'test-token' };
-      mockAuthManager.getAuthenticatedClient.mockResolvedValue(mockAuthClient as any);
-      await chatService.initialize();
-    });
-
     it('should create a space and return the space data', async () => {
       const mockResponse = {
         name: 'spaces/space1',