From 0f1258305a06ce1fd03b1579a48fcc98dbe35e99 Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Mon, 9 Mar 2026 12:53:24 -0400 Subject: [PATCH] perf(core): cache loadApiKey to reduce redundant keychain access (#21520) --- .../src/core/apiKeyCredentialStorage.test.ts | 90 +++++++++++++++---- .../core/src/core/apiKeyCredentialStorage.ts | 41 ++++++--- 2 files changed, 103 insertions(+), 28 deletions(-) diff --git a/packages/core/src/core/apiKeyCredentialStorage.test.ts b/packages/core/src/core/apiKeyCredentialStorage.test.ts index b0b0551f4b..b1eb9b21b7 100644 --- a/packages/core/src/core/apiKeyCredentialStorage.test.ts +++ b/packages/core/src/core/apiKeyCredentialStorage.test.ts @@ -9,6 +9,7 @@ import { loadApiKey, saveApiKey, clearApiKey, + resetApiKeyCacheForTesting, } from './apiKeyCredentialStorage.js'; const getCredentialsMock = vi.hoisted(() => vi.fn()); @@ -26,9 +27,10 @@ vi.mock('../mcp/token-storage/hybrid-token-storage.js', () => ({ describe('ApiKeyCredentialStorage', () => { beforeEach(() => { vi.clearAllMocks(); + resetApiKeyCacheForTesting(); }); - it('should load an API key', async () => { + it('should load an API key and cache it', async () => { getCredentialsMock.mockResolvedValue({ serverName: 'default-api-key', token: { @@ -38,19 +40,39 @@ describe('ApiKeyCredentialStorage', () => { updatedAt: Date.now(), }); - const apiKey = await loadApiKey(); - expect(apiKey).toBe('test-key'); - expect(getCredentialsMock).toHaveBeenCalledWith('default-api-key'); + const apiKey1 = await loadApiKey(); + expect(apiKey1).toBe('test-key'); + expect(getCredentialsMock).toHaveBeenCalledTimes(1); + + const apiKey2 = await loadApiKey(); + expect(apiKey2).toBe('test-key'); + expect(getCredentialsMock).toHaveBeenCalledTimes(1); // Should be cached }); - it('should return null if no API key is stored', async () => { + it('should return null if no API key is stored and cache it', async () => { getCredentialsMock.mockResolvedValue(null); - const apiKey = await loadApiKey(); - expect(apiKey).toBeNull(); - expect(getCredentialsMock).toHaveBeenCalledWith('default-api-key'); + const apiKey1 = await loadApiKey(); + expect(apiKey1).toBeNull(); + expect(getCredentialsMock).toHaveBeenCalledTimes(1); + + const apiKey2 = await loadApiKey(); + expect(apiKey2).toBeNull(); + expect(getCredentialsMock).toHaveBeenCalledTimes(1); // Should be cached }); - it('should save an API key', async () => { + it('should save an API key and clear cache', async () => { + getCredentialsMock.mockResolvedValue({ + serverName: 'default-api-key', + token: { + accessToken: 'old-key', + tokenType: 'ApiKey', + }, + updatedAt: Date.now(), + }); + + await loadApiKey(); + expect(getCredentialsMock).toHaveBeenCalledTimes(1); + await saveApiKey('new-key'); expect(setCredentialsMock).toHaveBeenCalledWith( expect.objectContaining({ @@ -61,28 +83,62 @@ describe('ApiKeyCredentialStorage', () => { }), }), ); + + getCredentialsMock.mockResolvedValue({ + serverName: 'default-api-key', + token: { + accessToken: 'new-key', + tokenType: 'ApiKey', + }, + updatedAt: Date.now(), + }); + + await loadApiKey(); + expect(getCredentialsMock).toHaveBeenCalledTimes(2); // Should have fetched again }); - it('should clear an API key when saving empty key', async () => { + it('should clear an API key and clear cache', async () => { + getCredentialsMock.mockResolvedValue({ + serverName: 'default-api-key', + token: { + accessToken: 'old-key', + tokenType: 'ApiKey', + }, + updatedAt: Date.now(), + }); + + await loadApiKey(); + expect(getCredentialsMock).toHaveBeenCalledTimes(1); + + await clearApiKey(); + expect(deleteCredentialsMock).toHaveBeenCalledWith('default-api-key'); + + getCredentialsMock.mockResolvedValue(null); + await loadApiKey(); + expect(getCredentialsMock).toHaveBeenCalledTimes(2); // Should have fetched again + }); + + it('should clear an API key and cache when saving empty key', async () => { await saveApiKey(''); expect(deleteCredentialsMock).toHaveBeenCalledWith('default-api-key'); expect(setCredentialsMock).not.toHaveBeenCalled(); }); - it('should clear an API key when saving null key', async () => { + it('should clear an API key and cache when saving null key', async () => { await saveApiKey(null); expect(deleteCredentialsMock).toHaveBeenCalledWith('default-api-key'); expect(setCredentialsMock).not.toHaveBeenCalled(); }); - it('should clear an API key', async () => { - await clearApiKey(); - expect(deleteCredentialsMock).toHaveBeenCalledWith('default-api-key'); - }); - - it('should not throw when clearing an API key fails', async () => { + it('should not throw when clearing an API key fails during saveApiKey', async () => { deleteCredentialsMock.mockRejectedValueOnce(new Error('Failed to delete')); await expect(saveApiKey('')).resolves.not.toThrow(); expect(deleteCredentialsMock).toHaveBeenCalledWith('default-api-key'); }); + + it('should not throw when clearing an API key fails during clearApiKey', async () => { + deleteCredentialsMock.mockRejectedValueOnce(new Error('Failed to delete')); + await expect(clearApiKey()).resolves.not.toThrow(); + expect(deleteCredentialsMock).toHaveBeenCalledWith('default-api-key'); + }); }); diff --git a/packages/core/src/core/apiKeyCredentialStorage.ts b/packages/core/src/core/apiKeyCredentialStorage.ts index 4836ba075b..41b3a0276a 100644 --- a/packages/core/src/core/apiKeyCredentialStorage.ts +++ b/packages/core/src/core/apiKeyCredentialStorage.ts @@ -7,29 +7,46 @@ import { HybridTokenStorage } from '../mcp/token-storage/hybrid-token-storage.js'; import type { OAuthCredentials } from '../mcp/token-storage/types.js'; import { debugLogger } from '../utils/debugLogger.js'; +import { createCache } from '../utils/cache.js'; const KEYCHAIN_SERVICE_NAME = 'gemini-cli-api-key'; const DEFAULT_API_KEY_ENTRY = 'default-api-key'; const storage = new HybridTokenStorage(KEYCHAIN_SERVICE_NAME); +// Cache to store the results of loadApiKey to avoid redundant keychain access. +const apiKeyCache = createCache>({ + storage: 'map', + defaultTtl: 30000, // 30 seconds +}); + +/** + * Resets the API key cache. Used exclusively for test isolation. + * @internal + */ +export function resetApiKeyCacheForTesting() { + apiKeyCache.clear(); +} + /** * Load cached API key */ export async function loadApiKey(): Promise { - try { - const credentials = await storage.getCredentials(DEFAULT_API_KEY_ENTRY); + return apiKeyCache.getOrCreate(DEFAULT_API_KEY_ENTRY, async () => { + try { + const credentials = await storage.getCredentials(DEFAULT_API_KEY_ENTRY); - if (credentials?.token?.accessToken) { - return credentials.token.accessToken; + if (credentials?.token?.accessToken) { + return credentials.token.accessToken; + } + + return null; + } catch (error: unknown) { + // Log other errors but don't crash, just return null so user can re-enter key + debugLogger.error('Failed to load API key from storage:', error); + return null; } - - return null; - } catch (error: unknown) { - // Log other errors but don't crash, just return null so user can re-enter key - debugLogger.error('Failed to load API key from storage:', error); - return null; - } + }); } /** @@ -38,6 +55,7 @@ export async function loadApiKey(): Promise { export async function saveApiKey( apiKey: string | null | undefined, ): Promise { + apiKeyCache.delete(DEFAULT_API_KEY_ENTRY); if (!apiKey || apiKey.trim() === '') { try { await storage.deleteCredentials(DEFAULT_API_KEY_ENTRY); @@ -65,6 +83,7 @@ export async function saveApiKey( * Clear cached API key */ export async function clearApiKey(): Promise { + apiKeyCache.delete(DEFAULT_API_KEY_ENTRY); try { await storage.deleteCredentials(DEFAULT_API_KEY_ENTRY); } catch (error: unknown) {