perf(core): cache loadApiKey to reduce redundant keychain access (#21520)

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();
+    const apiKey1 = await loadApiKey();
-    expect(apiKey).toBe('test-key');
+    expect(apiKey1).toBe('test-key');
-    expect(getCredentialsMock).toHaveBeenCalledWith('default-api-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();
+    const apiKey1 = await loadApiKey();
-    expect(apiKey).toBeNull();
+    expect(apiKey1).toBeNull();
-    expect(getCredentialsMock).toHaveBeenCalledWith('default-api-key');
+    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 () => {
+  it('should not throw when clearing an API key fails during saveApiKey', async () => {
-    await clearApiKey();
-    expect(deleteCredentialsMock).toHaveBeenCalledWith('default-api-key');
-  });
-
-  it('should not throw when clearing an API key fails', 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');
+  });
 });

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<string, Promise<string | null>>({
+  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<string | null> {
-  try {
+  return apiKeyCache.getOrCreate(DEFAULT_API_KEY_ENTRY, async () => {
-    const credentials = await storage.getCredentials(DEFAULT_API_KEY_ENTRY);
+    try {
+      const credentials = await storage.getCredentials(DEFAULT_API_KEY_ENTRY);
 
-    if (credentials?.token?.accessToken) {
+      if (credentials?.token?.accessToken) {
-      return 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<string | null> {
 export async function saveApiKey(
   apiKey: string | null | undefined,
 ): Promise<void> {
+  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<void> {
+  apiKeyCache.delete(DEFAULT_API_KEY_ENTRY);
   try {
     await storage.deleteCredentials(DEFAULT_API_KEY_ENTRY);
   } catch (error: unknown) {