fix(core): preserve OAuth refresh tokens during rotation and retrieval

Author: cocosheng-g

packages/core/src/code_assist/oauth-credential-storage.test.ts

@@ -242,6 +242,39 @@ describe('OAuthCredentialStorage', () => {
       );
     });
 
+    it('should merge existing refresh token when new payload lacks one', async () => {
+      const oldCredentials: OAuthCredentials = {
+        serverName: 'main-account',
+        token: {
+          accessToken: 'old-access-token',
+          refreshToken: 'persistent-refresh-token',
+          tokenType: 'Bearer',
+          expiresAt: Date.now() + 3600000,
+          scope: 'email',
+        },
+        updatedAt: Date.now(),
+      };
+      vi.spyOn(mockHybridTokenStorage, 'getCredentials').mockResolvedValue(
+        oldCredentials,
+      );
+
+      const newTokens: Credentials = {
+        access_token: 'new-access-token',
+        expiry_date: Date.now() + 3600000,
+      };
+
+      await OAuthCredentialStorage.saveCredentials(newTokens);
+
+      expect(mockHybridTokenStorage.setCredentials).toHaveBeenCalledWith(
+        expect.objectContaining({
+          token: expect.objectContaining({
+            accessToken: 'new-access-token',
+            refreshToken: 'persistent-refresh-token', // correctly merged
+          }),
+        }),
+      );
+    });
+
     it('should throw an error if access_token is missing', async () => {
       const invalidCredentials: Credentials = {
         ...mockCredentials,

packages/core/src/code_assist/oauth-credential-storage.ts

@@ -66,12 +66,16 @@ export class OAuthCredentialStorage {
       throw new Error('Attempted to save credentials without an access token.');
     }
 
+    const existing = await this.storage.getCredentials(MAIN_ACCOUNT_KEY);
+    const mergedRefreshToken =
+      credentials.refresh_token || existing?.token.refreshToken;
+
     // Convert Google Credentials to OAuthCredentials format
     const mcpCredentials: OAuthCredentials = {
       serverName: MAIN_ACCOUNT_KEY,
       token: {
         accessToken: credentials.access_token,
-        refreshToken: credentials.refresh_token || undefined,
+        refreshToken: mergedRefreshToken || undefined,
         tokenType: credentials.token_type || 'Bearer',
         scope: credentials.scope || undefined,
         expiresAt: credentials.expiry_date || undefined,

packages/core/src/mcp/oauth-token-storage.ts

@@ -143,9 +143,18 @@ export class MCPOAuthTokenStorage implements TokenStorage {
   ): Promise<void> {
     await this.ensureConfigDir();
 
+    const existing = await this.getCredentials(serverName);
+    const mergedRefreshToken =
+      token.refreshToken || existing?.token.refreshToken;
+
+    const mergedToken = {
+      ...token,
+      refreshToken: mergedRefreshToken,
+    };
+
     const credential: OAuthCredentials = {
       serverName,
-      token,
+      token: mergedToken,
       clientId,
       tokenUrl,
       mcpServerUrl,

packages/core/src/mcp/token-storage/keychain-token-storage.test.ts

@@ -72,7 +72,7 @@ describe('KeychainTokenStorage', () => {
       expect(retrieved?.serverName).toBe('test-server');
     });
 
-    it('should return null if no credentials are found or they are expired', async () => {
+    it('should return null if no credentials are found or they are expired and unrefreshable', async () => {
       expect(await storage.getCredentials('missing')).toBeNull();
 
       const expiredCreds = {
@@ -81,6 +81,20 @@ describe('KeychainTokenStorage', () => {
       };
       await storage.setCredentials(expiredCreds);
       expect(await storage.getCredentials('test-server')).toBeNull();
+
+      // Ensure that if it has a refresh token, it is NOT returned as null
+      const expiredWithRefresh = {
+        ...validCredentials,
+        token: {
+          ...validCredentials.token,
+          expiresAt: Date.now() - 1000,
+          refreshToken: 'some-refresh-token',
+        },
+      };
+      await storage.setCredentials(expiredWithRefresh);
+      const retrieved = await storage.getCredentials('test-server');
+      expect(retrieved).not.toBeNull();
+      expect(retrieved?.token.refreshToken).toBe('some-refresh-token');
     });
 
     it('should throw if stored data is corrupted JSON', async () => {

packages/core/src/mcp/token-storage/keychain-token-storage.ts

@@ -36,7 +36,7 @@ export class KeychainTokenStorage
       // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
       const credentials = JSON.parse(data) as OAuthCredentials;
 
-      if (this.isTokenExpired(credentials)) {
+      if (this.isTokenExpired(credentials) && !credentials.token.refreshToken) {
         return null;
       }
 
@@ -104,7 +104,7 @@ export class KeychainTokenStorage
         try {
           // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
           const data = JSON.parse(cred.password) as OAuthCredentials;
-          if (!this.isTokenExpired(data)) {
+          if (!this.isTokenExpired(data) || data.token.refreshToken) {
             result.set(cred.account, data);
           }
         } catch (error) {