diff --git a/packages/core/src/mcp/oauth-token-storage.test.ts b/packages/core/src/mcp/oauth-token-storage.test.ts index 6b7f9c8c3f..cd8841aaee 100644 --- a/packages/core/src/mcp/oauth-token-storage.test.ts +++ b/packages/core/src/mcp/oauth-token-storage.test.ts @@ -4,13 +4,15 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { promises as fs } from 'node:fs'; import * as path from 'node:path'; import { MCPOAuthTokenStorage } from './oauth-token-storage.js'; -import type { OAuthToken, OAuthCredentials } from './token-storage/types.js'; +import { FORCE_ENCRYPTED_FILE_ENV_VAR } from './token-storage/index.js'; +import type { OAuthCredentials, OAuthToken } from './token-storage/types.js'; +import { GEMINI_DIR } from '../utils/paths.js'; -// Mock file system operations +// Mock dependencies vi.mock('node:fs', () => ({ promises: { readFile: vi.fn(), @@ -18,13 +20,33 @@ vi.mock('node:fs', () => ({ mkdir: vi.fn(), unlink: vi.fn(), }, - mkdirSync: vi.fn(), })); -vi.mock('node:os', () => ({ - homedir: vi.fn(() => '/mock/home'), +vi.mock('node:path', () => ({ + dirname: vi.fn(), + join: vi.fn(), })); +vi.mock('../config/storage.js', () => ({ + Storage: { + getMcpOAuthTokensPath: vi.fn(), + }, +})); + +const mockHybridTokenStorage = { + listServers: vi.fn(), + setCredentials: vi.fn(), + getCredentials: vi.fn(), + deleteCredentials: vi.fn(), + clearAll: vi.fn(), + getAllCredentials: vi.fn(), +}; +vi.mock('./token-storage/hybrid-token-storage.js', () => ({ + HybridTokenStorage: vi.fn(() => mockHybridTokenStorage), +})); + +const ONE_HR_MS = 3600000; + describe('MCPOAuthTokenStorage', () => { let tokenStorage: MCPOAuthTokenStorage; @@ -33,7 +55,7 @@ describe('MCPOAuthTokenStorage', () => { refreshToken: 'refresh_token_456', tokenType: 'Bearer', scope: 'read write', - expiresAt: Date.now() + 3600000, // 1 hour from now + expiresAt: Date.now() + ONE_HR_MS, }; const mockCredentials: OAuthCredentials = { @@ -44,283 +66,385 @@ describe('MCPOAuthTokenStorage', () => { updatedAt: Date.now(), }; - beforeEach(() => { - tokenStorage = new MCPOAuthTokenStorage(); - vi.clearAllMocks(); - vi.spyOn(console, 'error').mockImplementation(() => {}); - }); + describe('with encrypted flag false', () => { + beforeEach(() => { + vi.stubEnv(FORCE_ENCRYPTED_FILE_ENV_VAR, 'false'); + tokenStorage = new MCPOAuthTokenStorage(); - afterEach(() => { - vi.restoreAllMocks(); - }); - - describe('getAllCredentials', () => { - it('should return empty map when token file does not exist', async () => { - vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' }); - - const tokens = await tokenStorage.getAllCredentials(); - - expect(tokens.size).toBe(0); - expect(console.error).not.toHaveBeenCalled(); + vi.clearAllMocks(); + vi.spyOn(console, 'error'); }); - it('should load tokens from file successfully', async () => { - const tokensArray = [mockCredentials]; - vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(tokensArray)); - - const tokens = await tokenStorage.getAllCredentials(); - - expect(tokens.size).toBe(1); - expect(tokens.get('test-server')).toEqual(mockCredentials); - expect(fs.readFile).toHaveBeenCalledWith( - path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'), - 'utf-8', - ); + afterEach(() => { + vi.unstubAllEnvs(); + vi.restoreAllMocks(); }); - it('should handle corrupted token file gracefully', async () => { - vi.mocked(fs.readFile).mockResolvedValue('invalid json'); + describe('getAllCredentials', () => { + it('should return empty map when token file does not exist', async () => { + vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' }); - const tokens = await tokenStorage.getAllCredentials(); + const tokens = await tokenStorage.getAllCredentials(); - expect(tokens.size).toBe(0); - expect(console.error).toHaveBeenCalledWith( - expect.stringContaining('Failed to load MCP OAuth tokens'), - ); + expect(tokens.size).toBe(0); + expect(console.error).not.toHaveBeenCalled(); + }); + + it('should load tokens from file successfully', async () => { + const tokensArray = [mockCredentials]; + vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(tokensArray)); + + const tokens = await tokenStorage.getAllCredentials(); + + expect(tokens.size).toBe(1); + expect(tokens.get('test-server')).toEqual(mockCredentials); + expect(fs.readFile).toHaveBeenCalledWith( + path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'), + 'utf-8', + ); + }); + + it('should handle corrupted token file gracefully', async () => { + vi.mocked(fs.readFile).mockResolvedValue('invalid json'); + + const tokens = await tokenStorage.getAllCredentials(); + + expect(tokens.size).toBe(0); + expect(console.error).toHaveBeenCalledWith( + expect.stringContaining('Failed to load MCP OAuth tokens'), + ); + }); + + it('should handle file read errors other than ENOENT', async () => { + const error = new Error('Permission denied'); + vi.mocked(fs.readFile).mockRejectedValue(error); + + const tokens = await tokenStorage.getAllCredentials(); + + expect(tokens.size).toBe(0); + expect(console.error).toHaveBeenCalledWith( + expect.stringContaining('Failed to load MCP OAuth tokens'), + ); + }); }); - it('should handle file read errors other than ENOENT', async () => { - const error = new Error('Permission denied'); - vi.mocked(fs.readFile).mockRejectedValue(error); + describe('saveToken', () => { + it('should save token with restricted permissions', async () => { + vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' }); + vi.mocked(fs.mkdir).mockResolvedValue(undefined); + vi.mocked(fs.writeFile).mockResolvedValue(undefined); - const tokens = await tokenStorage.getAllCredentials(); + await tokenStorage.saveToken( + 'test-server', + mockToken, + 'client-id', + 'https://token.url', + ); - expect(tokens.size).toBe(0); - expect(console.error).toHaveBeenCalledWith( - expect.stringContaining('Failed to load MCP OAuth tokens'), - ); + expect(fs.mkdir).toHaveBeenCalledWith( + path.join('/mock/home', GEMINI_DIR), + { recursive: true }, + ); + expect(fs.writeFile).toHaveBeenCalledWith( + path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'), + expect.stringContaining('test-server'), + { mode: 0o600 }, + ); + }); + + it('should update existing token for same server', async () => { + const existingCredentials: OAuthCredentials = { + ...mockCredentials, + serverName: 'existing-server', + }; + vi.mocked(fs.readFile).mockResolvedValue( + JSON.stringify([existingCredentials]), + ); + vi.mocked(fs.writeFile).mockResolvedValue(undefined); + + const newToken: OAuthToken = { + ...mockToken, + accessToken: 'new_access_token', + }; + await tokenStorage.saveToken('existing-server', newToken); + + const writeCall = vi.mocked(fs.writeFile).mock.calls[0]; + const savedData = JSON.parse( + writeCall[1] as string, + ) as OAuthCredentials[]; + + expect(savedData).toHaveLength(1); + expect(savedData[0].token.accessToken).toBe('new_access_token'); + expect(savedData[0].serverName).toBe('existing-server'); + }); + + it('should handle write errors gracefully', async () => { + vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' }); + vi.mocked(fs.mkdir).mockResolvedValue(undefined); + const writeError = new Error('Disk full'); + vi.mocked(fs.writeFile).mockRejectedValue(writeError); + + await expect( + tokenStorage.saveToken('test-server', mockToken), + ).rejects.toThrow('Disk full'); + + expect(console.error).toHaveBeenCalledWith( + expect.stringContaining('Failed to save MCP OAuth token'), + ); + }); + }); + + describe('getCredentials', () => { + it('should return token for existing server', async () => { + vi.mocked(fs.readFile).mockResolvedValue( + JSON.stringify([mockCredentials]), + ); + + const result = await tokenStorage.getCredentials('test-server'); + + expect(result).toEqual(mockCredentials); + }); + + it('should return null for non-existent server', async () => { + vi.mocked(fs.readFile).mockResolvedValue( + JSON.stringify([mockCredentials]), + ); + + const result = await tokenStorage.getCredentials('non-existent'); + + expect(result).toBeNull(); + }); + + it('should return null when no tokens file exists', async () => { + vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' }); + + const result = await tokenStorage.getCredentials('test-server'); + + expect(result).toBeNull(); + }); + }); + + describe('deleteCredentials', () => { + it('should remove token for specific server', async () => { + const credentials1: OAuthCredentials = { + ...mockCredentials, + serverName: 'server1', + }; + const credentials2: OAuthCredentials = { + ...mockCredentials, + serverName: 'server2', + }; + vi.mocked(fs.readFile).mockResolvedValue( + JSON.stringify([credentials1, credentials2]), + ); + vi.mocked(fs.writeFile).mockResolvedValue(undefined); + + await tokenStorage.deleteCredentials('server1'); + + const writeCall = vi.mocked(fs.writeFile).mock.calls[0]; + const savedData = JSON.parse(writeCall[1] as string); + + expect(savedData).toHaveLength(1); + expect(savedData[0].serverName).toBe('server2'); + }); + + it('should remove token file when no tokens remain', async () => { + vi.mocked(fs.readFile).mockResolvedValue( + JSON.stringify([mockCredentials]), + ); + vi.mocked(fs.unlink).mockResolvedValue(undefined); + + await tokenStorage.deleteCredentials('test-server'); + + expect(fs.unlink).toHaveBeenCalledWith( + path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'), + ); + expect(fs.writeFile).not.toHaveBeenCalled(); + }); + + it('should handle removal of non-existent token gracefully', async () => { + vi.mocked(fs.readFile).mockResolvedValue( + JSON.stringify([mockCredentials]), + ); + + await tokenStorage.deleteCredentials('non-existent'); + + expect(fs.writeFile).not.toHaveBeenCalled(); + expect(fs.unlink).not.toHaveBeenCalled(); + }); + + it('should handle file operation errors gracefully', async () => { + vi.mocked(fs.readFile).mockResolvedValue( + JSON.stringify([mockCredentials]), + ); + vi.mocked(fs.unlink).mockRejectedValue(new Error('Permission denied')); + + await tokenStorage.deleteCredentials('test-server'); + + expect(console.error).toHaveBeenCalledWith( + expect.stringContaining('Failed to remove MCP OAuth token'), + ); + }); + }); + + describe('isTokenExpired', () => { + it('should return false for token without expiry', () => { + const tokenWithoutExpiry: OAuthToken = { ...mockToken }; + delete tokenWithoutExpiry.expiresAt; + + const result = tokenStorage.isTokenExpired(tokenWithoutExpiry); + + expect(result).toBe(false); + }); + + it('should return false for valid token', () => { + const futureToken: OAuthToken = { + ...mockToken, + expiresAt: Date.now() + ONE_HR_MS, + }; + + const result = tokenStorage.isTokenExpired(futureToken); + + expect(result).toBe(false); + }); + + it('should return true for expired token', () => { + const expiredToken: OAuthToken = { + ...mockToken, + expiresAt: Date.now() - ONE_HR_MS, + }; + + const result = tokenStorage.isTokenExpired(expiredToken); + + expect(result).toBe(true); + }); + + it('should return true for token expiring within buffer time', () => { + const soonToExpireToken: OAuthToken = { + ...mockToken, + expiresAt: Date.now() + 60000, // 1 minute from now (within 5-minute buffer) + }; + + const result = tokenStorage.isTokenExpired(soonToExpireToken); + + expect(result).toBe(true); + }); + }); + + describe('clearAll', () => { + it('should remove token file successfully', async () => { + vi.mocked(fs.unlink).mockResolvedValue(undefined); + + await tokenStorage.clearAll(); + + expect(fs.unlink).toHaveBeenCalledWith( + path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'), + ); + }); + + it('should handle non-existent file gracefully', async () => { + vi.mocked(fs.unlink).mockRejectedValue({ code: 'ENOENT' }); + + await tokenStorage.clearAll(); + + expect(console.error).not.toHaveBeenCalled(); + }); + + it('should handle other file errors gracefully', async () => { + vi.mocked(fs.unlink).mockRejectedValue(new Error('Permission denied')); + + await tokenStorage.clearAll(); + + expect(console.error).toHaveBeenCalledWith( + expect.stringContaining('Failed to clear MCP OAuth tokens'), + ); + }); }); }); - describe('saveToken', () => { - it('should save token with restricted permissions', async () => { - vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' }); - vi.mocked(fs.mkdir).mockResolvedValue(undefined); - vi.mocked(fs.writeFile).mockResolvedValue(undefined); + describe('with encrypted flag true', () => { + beforeEach(() => { + vi.stubEnv(FORCE_ENCRYPTED_FILE_ENV_VAR, 'true'); + tokenStorage = new MCPOAuthTokenStorage(); + + vi.clearAllMocks(); + vi.spyOn(console, 'error'); + }); + + afterEach(() => { + vi.unstubAllEnvs(); + vi.restoreAllMocks(); + }); + + it('should use HybridTokenStorage to list all credentials', async () => { + mockHybridTokenStorage.getAllCredentials.mockResolvedValue(new Map()); + const servers = await tokenStorage.getAllCredentials(); + expect(mockHybridTokenStorage.getAllCredentials).toHaveBeenCalled(); + expect(servers).toEqual(new Map()); + }); + + it('should use HybridTokenStorage to list servers', async () => { + mockHybridTokenStorage.listServers.mockResolvedValue(['server1']); + const servers = await tokenStorage.listServers(); + expect(mockHybridTokenStorage.listServers).toHaveBeenCalled(); + expect(servers).toEqual(['server1']); + }); + + it('should use HybridTokenStorage to set credentials', async () => { + await tokenStorage.setCredentials(mockCredentials); + expect(mockHybridTokenStorage.setCredentials).toHaveBeenCalledWith( + mockCredentials, + ); + }); + + it('should use HybridTokenStorage to save a token', async () => { + const serverName = 'server1'; + const now = Date.now(); + vi.spyOn(Date, 'now').mockReturnValue(now); await tokenStorage.saveToken( - 'test-server', + serverName, mockToken, - 'client-id', - 'https://token.url', + 'clientId', + 'tokenUrl', + 'mcpUrl', ); - expect(fs.mkdir).toHaveBeenCalledWith( - path.join('/mock/home', '.gemini'), - { recursive: true }, - ); - expect(fs.writeFile).toHaveBeenCalledWith( - path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'), - expect.stringContaining('test-server'), - { mode: 0o600 }, - ); - }); - - it('should update existing token for same server', async () => { - const existingCredentials = { - ...mockCredentials, - serverName: 'existing-server', + const expectedCredential: OAuthCredentials = { + serverName, + token: mockToken, + clientId: 'clientId', + tokenUrl: 'tokenUrl', + mcpServerUrl: 'mcpUrl', + updatedAt: now, }; - vi.mocked(fs.readFile).mockResolvedValue( - JSON.stringify([existingCredentials]), + + expect(mockHybridTokenStorage.setCredentials).toHaveBeenCalledWith( + expectedCredential, ); - vi.mocked(fs.writeFile).mockResolvedValue(undefined); - - const newToken = { ...mockToken, accessToken: 'new_access_token' }; - await tokenStorage.saveToken('existing-server', newToken); - - const writeCall = vi.mocked(fs.writeFile).mock.calls[0]; - const savedData = JSON.parse(writeCall[1] as string); - - expect(savedData).toHaveLength(1); - expect(savedData[0].token.accessToken).toBe('new_access_token'); - expect(savedData[0].serverName).toBe('existing-server'); + expect(path.dirname).toHaveBeenCalled(); + expect(fs.mkdir).toHaveBeenCalled(); }); - it('should handle write errors gracefully', async () => { - vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' }); - vi.mocked(fs.mkdir).mockResolvedValue(undefined); - const writeError = new Error('Disk full'); - vi.mocked(fs.writeFile).mockRejectedValue(writeError); - - await expect( - tokenStorage.saveToken('test-server', mockToken), - ).rejects.toThrow('Disk full'); - - expect(console.error).toHaveBeenCalledWith( - expect.stringContaining('Failed to save MCP OAuth token'), + it('should use HybridTokenStorage to get credentials', async () => { + mockHybridTokenStorage.getCredentials.mockResolvedValue(mockCredentials); + const result = await tokenStorage.getCredentials('server1'); + expect(mockHybridTokenStorage.getCredentials).toHaveBeenCalledWith( + 'server1', ); - }); - }); - - describe('getCredentials', () => { - it('should return token for existing server', async () => { - vi.mocked(fs.readFile).mockResolvedValue( - JSON.stringify([mockCredentials]), - ); - - const result = await tokenStorage.getCredentials('test-server'); - - expect(result).toEqual(mockCredentials); + expect(result).toBe(mockCredentials); }); - it('should return null for non-existent server', async () => { - vi.mocked(fs.readFile).mockResolvedValue( - JSON.stringify([mockCredentials]), - ); - - const result = await tokenStorage.getCredentials('non-existent'); - - expect(result).toBeNull(); - }); - - it('should return null when no tokens file exists', async () => { - vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' }); - - const result = await tokenStorage.getCredentials('test-server'); - - expect(result).toBeNull(); - }); - }); - - describe('deleteCredentials', () => { - it('should remove token for specific server', async () => { - const credentials1 = { ...mockCredentials, serverName: 'server1' }; - const credentials2 = { ...mockCredentials, serverName: 'server2' }; - vi.mocked(fs.readFile).mockResolvedValue( - JSON.stringify([credentials1, credentials2]), - ); - vi.mocked(fs.writeFile).mockResolvedValue(undefined); - + it('should use HybridTokenStorage to delete credentials', async () => { await tokenStorage.deleteCredentials('server1'); - - const writeCall = vi.mocked(fs.writeFile).mock.calls[0]; - const savedData = JSON.parse(writeCall[1] as string); - - expect(savedData).toHaveLength(1); - expect(savedData[0].serverName).toBe('server2'); - }); - - it('should remove token file when no tokens remain', async () => { - vi.mocked(fs.readFile).mockResolvedValue( - JSON.stringify([mockCredentials]), - ); - vi.mocked(fs.unlink).mockResolvedValue(undefined); - - await tokenStorage.deleteCredentials('test-server'); - - expect(fs.unlink).toHaveBeenCalledWith( - path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'), - ); - expect(fs.writeFile).not.toHaveBeenCalled(); - }); - - it('should handle removal of non-existent token gracefully', async () => { - vi.mocked(fs.readFile).mockResolvedValue( - JSON.stringify([mockCredentials]), - ); - - await tokenStorage.deleteCredentials('non-existent'); - - expect(fs.writeFile).not.toHaveBeenCalled(); - expect(fs.unlink).not.toHaveBeenCalled(); - }); - - it('should handle file operation errors gracefully', async () => { - vi.mocked(fs.readFile).mockResolvedValue( - JSON.stringify([mockCredentials]), - ); - vi.mocked(fs.unlink).mockRejectedValue(new Error('Permission denied')); - - await tokenStorage.deleteCredentials('test-server'); - - expect(console.error).toHaveBeenCalledWith( - expect.stringContaining('Failed to remove MCP OAuth token'), + expect(mockHybridTokenStorage.deleteCredentials).toHaveBeenCalledWith( + 'server1', ); }); - }); - - describe('isTokenExpired', () => { - it('should return false for token without expiry', () => { - const tokenWithoutExpiry = { ...mockToken }; - delete tokenWithoutExpiry.expiresAt; - - const result = tokenStorage.isTokenExpired(tokenWithoutExpiry); - - expect(result).toBe(false); - }); - - it('should return false for valid token', () => { - const futureToken = { - ...mockToken, - expiresAt: Date.now() + 3600000, // 1 hour from now - }; - - const result = tokenStorage.isTokenExpired(futureToken); - - expect(result).toBe(false); - }); - - it('should return true for expired token', () => { - const expiredToken = { - ...mockToken, - expiresAt: Date.now() - 3600000, // 1 hour ago - }; - - const result = tokenStorage.isTokenExpired(expiredToken); - - expect(result).toBe(true); - }); - - it('should return true for token expiring within buffer time', () => { - const soonToExpireToken = { - ...mockToken, - expiresAt: Date.now() + 60000, // 1 minute from now (within 5-minute buffer) - }; - - const result = tokenStorage.isTokenExpired(soonToExpireToken); - - expect(result).toBe(true); - }); - }); - - describe('clearAll', () => { - it('should remove token file successfully', async () => { - vi.mocked(fs.unlink).mockResolvedValue(undefined); + it('should use HybridTokenStorage to clear all tokens', async () => { await tokenStorage.clearAll(); - - expect(fs.unlink).toHaveBeenCalledWith( - path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'), - ); - }); - - it('should handle non-existent file gracefully', async () => { - vi.mocked(fs.unlink).mockRejectedValue({ code: 'ENOENT' }); - - await tokenStorage.clearAll(); - - expect(console.error).not.toHaveBeenCalled(); - }); - - it('should handle other file errors gracefully', async () => { - vi.mocked(fs.unlink).mockRejectedValue(new Error('Permission denied')); - - await tokenStorage.clearAll(); - - expect(console.error).toHaveBeenCalledWith( - expect.stringContaining('Failed to clear MCP OAuth tokens'), - ); + expect(mockHybridTokenStorage.clearAll).toHaveBeenCalled(); }); }); }); diff --git a/packages/core/src/mcp/oauth-token-storage.ts b/packages/core/src/mcp/oauth-token-storage.ts index ba8af40229..d9d98ff417 100644 --- a/packages/core/src/mcp/oauth-token-storage.ts +++ b/packages/core/src/mcp/oauth-token-storage.ts @@ -13,11 +13,22 @@ import type { OAuthCredentials, TokenStorage, } from './token-storage/types.js'; +import { HybridTokenStorage } from './token-storage/hybrid-token-storage.js'; +import { + DEFAULT_SERVICE_NAME, + FORCE_ENCRYPTED_FILE_ENV_VAR, +} from './token-storage/index.js'; /** * Class for managing MCP OAuth token storage and retrieval. */ export class MCPOAuthTokenStorage implements TokenStorage { + private readonly hybridTokenStorage = new HybridTokenStorage( + DEFAULT_SERVICE_NAME, + ); + private readonly useEncryptedFile = + process.env[FORCE_ENCRYPTED_FILE_ENV_VAR] === 'true'; + /** * Get the path to the token storage file. * @@ -41,6 +52,9 @@ export class MCPOAuthTokenStorage implements TokenStorage { * @returns A map of server names to credentials */ async getAllCredentials(): Promise> { + if (this.useEncryptedFile) { + return this.hybridTokenStorage.getAllCredentials(); + } const tokenMap = new Map(); try { @@ -64,11 +78,17 @@ export class MCPOAuthTokenStorage implements TokenStorage { } async listServers(): Promise { + if (this.useEncryptedFile) { + return this.hybridTokenStorage.listServers(); + } const tokens = await this.getAllCredentials(); return Array.from(tokens.keys()); } async setCredentials(credentials: OAuthCredentials): Promise { + if (this.useEncryptedFile) { + return this.hybridTokenStorage.setCredentials(credentials); + } const tokens = await this.getAllCredentials(); tokens.set(credentials.serverName, credentials); @@ -116,6 +136,9 @@ export class MCPOAuthTokenStorage implements TokenStorage { updatedAt: Date.now(), }; + if (this.useEncryptedFile) { + return this.hybridTokenStorage.setCredentials(credential); + } await this.setCredentials(credential); } @@ -126,6 +149,9 @@ export class MCPOAuthTokenStorage implements TokenStorage { * @returns The stored credentials or null if not found */ async getCredentials(serverName: string): Promise { + if (this.useEncryptedFile) { + return this.hybridTokenStorage.getCredentials(serverName); + } const tokens = await this.getAllCredentials(); return tokens.get(serverName) || null; } @@ -136,6 +162,9 @@ export class MCPOAuthTokenStorage implements TokenStorage { * @param serverName The name of the MCP server */ async deleteCredentials(serverName: string): Promise { + if (this.useEncryptedFile) { + return this.hybridTokenStorage.deleteCredentials(serverName); + } const tokens = await this.getAllCredentials(); if (tokens.delete(serverName)) { @@ -179,6 +208,9 @@ export class MCPOAuthTokenStorage implements TokenStorage { * Clear all stored MCP OAuth tokens. */ async clearAll(): Promise { + if (this.useEncryptedFile) { + return this.hybridTokenStorage.clearAll(); + } try { const tokenFile = this.getTokenFilePath(); await fs.unlink(tokenFile); diff --git a/packages/core/src/mcp/token-storage/index.ts b/packages/core/src/mcp/token-storage/index.ts new file mode 100644 index 0000000000..0b48a933a9 --- /dev/null +++ b/packages/core/src/mcp/token-storage/index.ts @@ -0,0 +1,14 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export * from './types.js'; +export * from './base-token-storage.js'; +export * from './file-token-storage.js'; +export * from './hybrid-token-storage.js'; + +export const DEFAULT_SERVICE_NAME = 'gemini-cli-oauth'; +export const FORCE_ENCRYPTED_FILE_ENV_VAR = + 'GEMINI_FORCE_ENCRYPTED_FILE_STORAGE';