feat(auth): improve API key authentication flow (#11760)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

packages/core/src/config/config.test.ts

@@ -226,7 +226,7 @@ describe('Server Config (config.ts)', () => {
         apiKey: 'test-key',
       };
 
-      vi.mocked(createContentGeneratorConfig).mockReturnValue(
+      vi.mocked(createContentGeneratorConfig).mockResolvedValue(
         mockContentConfig,
       );
 
@@ -251,7 +251,7 @@ describe('Server Config (config.ts)', () => {
       const config = new Config(baseParams);
 
       vi.mocked(createContentGeneratorConfig).mockImplementation(
-        (_: Config, authType: AuthType | undefined) =>
+        async (_: Config, authType: AuthType | undefined) =>
           ({ authType }) as unknown as ContentGeneratorConfig,
       );
 
@@ -268,7 +268,7 @@ describe('Server Config (config.ts)', () => {
       const config = new Config(baseParams);
 
       vi.mocked(createContentGeneratorConfig).mockImplementation(
-        (_: Config, authType: AuthType | undefined) =>
+        async (_: Config, authType: AuthType | undefined) =>
           ({ authType }) as unknown as ContentGeneratorConfig,
       );
 
@@ -1105,7 +1105,9 @@ describe('BaseLlmClient Lifecycle', () => {
     const authType = AuthType.USE_GEMINI;
     const mockContentConfig = { model: 'gemini-flash', apiKey: 'test-key' };
 
-    vi.mocked(createContentGeneratorConfig).mockReturnValue(mockContentConfig);
+    vi.mocked(createContentGeneratorConfig).mockResolvedValue(
+      mockContentConfig,
+    );
 
     await config.refreshAuth(authType);
 

packages/core/src/config/config.ts

@@ -579,7 +579,7 @@ export class Config {
       this.geminiClient.stripThoughtsFromHistory();
     }
 
-    const newContentGeneratorConfig = createContentGeneratorConfig(
+    const newContentGeneratorConfig = await createContentGeneratorConfig(
       this,
       authMethod,
     );

packages/core/src/core/apiKeyCredentialStorage.test.ts

@@ -0,0 +1,82 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import {
+  loadApiKey,
+  saveApiKey,
+  clearApiKey,
+} from './apiKeyCredentialStorage.js';
+
+const getCredentialsMock = vi.hoisted(() => vi.fn());
+const setCredentialsMock = vi.hoisted(() => vi.fn());
+const deleteCredentialsMock = vi.hoisted(() => vi.fn());
+
+vi.mock('../mcp/token-storage/hybrid-token-storage.js', () => ({
+  HybridTokenStorage: vi.fn().mockImplementation(() => ({
+    getCredentials: getCredentialsMock,
+    setCredentials: setCredentialsMock,
+    deleteCredentials: deleteCredentialsMock,
+  })),
+}));
+
+describe('ApiKeyCredentialStorage', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('should load an API key', async () => {
+    getCredentialsMock.mockResolvedValue({
+      serverName: 'default-api-key',
+      token: {
+        accessToken: 'test-key',
+        tokenType: 'ApiKey',
+      },
+      updatedAt: Date.now(),
+    });
+
+    const apiKey = await loadApiKey();
+    expect(apiKey).toBe('test-key');
+    expect(getCredentialsMock).toHaveBeenCalledWith('default-api-key');
+  });
+
+  it('should return null if no API key is stored', async () => {
+    getCredentialsMock.mockResolvedValue(null);
+    const apiKey = await loadApiKey();
+    expect(apiKey).toBeNull();
+    expect(getCredentialsMock).toHaveBeenCalledWith('default-api-key');
+  });
+
+  it('should save an API key', async () => {
+    await saveApiKey('new-key');
+    expect(setCredentialsMock).toHaveBeenCalledWith(
+      expect.objectContaining({
+        serverName: 'default-api-key',
+        token: expect.objectContaining({
+          accessToken: 'new-key',
+          tokenType: 'ApiKey',
+        }),
+      }),
+    );
+  });
+
+  it('should clear an API key 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 () => {
+    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');
+  });
+});

packages/core/src/core/apiKeyCredentialStorage.ts

@@ -0,0 +1,77 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+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';
+
+const KEYCHAIN_SERVICE_NAME = 'gemini-cli-api-key';
+const DEFAULT_API_KEY_ENTRY = 'default-api-key';
+
+const storage = new HybridTokenStorage(KEYCHAIN_SERVICE_NAME);
+
+/**
+ * Load cached API key
+ */
+export async function loadApiKey(): Promise<string | null> {
+  try {
+    const credentials = await storage.getCredentials(DEFAULT_API_KEY_ENTRY);
+
+    if (credentials?.token?.accessToken) {
+      return credentials.token.accessToken;
+    }
+
+    return null;
+  } catch (error: unknown) {
+    // Ignore "file not found" error from FileTokenStorage, it just means no key is saved yet.
+    // This is common in fresh environments like e2e tests.
+    if (
+      error instanceof Error &&
+      error.message === 'Token file does not exist'
+    ) {
+      return null;
+    }
+
+    // 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;
+  }
+}
+
+/**
+ * Save API key
+ */
+export async function saveApiKey(
+  apiKey: string | null | undefined,
+): Promise<void> {
+  if (!apiKey || apiKey.trim() === '') {
+    await storage.deleteCredentials(DEFAULT_API_KEY_ENTRY);
+    return;
+  }
+
+  // Wrap API key in OAuthCredentials format as required by HybridTokenStorage
+  const credentials: OAuthCredentials = {
+    serverName: DEFAULT_API_KEY_ENTRY,
+    token: {
+      accessToken: apiKey,
+      tokenType: 'ApiKey',
+    },
+    updatedAt: Date.now(),
+  };
+
+  await storage.setCredentials(credentials);
+}
+
+/**
+ * Clear cached API key
+ */
+export async function clearApiKey(): Promise<void> {
+  try {
+    await storage.deleteCredentials(DEFAULT_API_KEY_ENTRY);
+  } catch (error: unknown) {
+    debugLogger.error('Failed to clear API key from storage:', error);
+  }
+}

packages/core/src/core/contentGenerator.test.ts

@@ -15,11 +15,16 @@ import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
 import { GoogleGenAI } from '@google/genai';
 import type { Config } from '../config/config.js';
 import { LoggingContentGenerator } from './loggingContentGenerator.js';
+import { loadApiKey } from './apiKeyCredentialStorage.js';
 import { FakeContentGenerator } from './fakeContentGenerator.js';
 import { RecordingContentGenerator } from './recordingContentGenerator.js';
 
 vi.mock('../code_assist/codeAssist.js');
 vi.mock('@google/genai');
+vi.mock('./apiKeyCredentialStorage.js', () => ({
+  loadApiKey: vi.fn(),
+}));
+
 vi.mock('./fakeContentGenerator.js');
 
 const mockConfig = {} as unknown as Config;
@@ -184,6 +189,17 @@ describe('createContentGeneratorConfig', () => {
     expect(config.vertexai).toBeUndefined();
   });
 
+  it('should not configure for Gemini if GEMINI_API_KEY is not set and storage is empty', async () => {
+    vi.stubEnv('GEMINI_API_KEY', '');
+    vi.mocked(loadApiKey).mockResolvedValue(null);
+    const config = await createContentGeneratorConfig(
+      mockConfig,
+      AuthType.USE_GEMINI,
+    );
+    expect(config.apiKey).toBeUndefined();
+    expect(config.vertexai).toBeUndefined();
+  });
+
   it('should configure for Vertex AI using GOOGLE_API_KEY when set', async () => {
     vi.stubEnv('GOOGLE_API_KEY', 'env-google-key');
     const config = await createContentGeneratorConfig(

packages/core/src/core/contentGenerator.ts

@@ -15,6 +15,7 @@ import type {
 import { GoogleGenAI } from '@google/genai';
 import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
 import type { Config } from '../config/config.js';
+import { loadApiKey } from './apiKeyCredentialStorage.js';
 
 import type { UserTierId } from '../code_assist/types.js';
 import { LoggingContentGenerator } from './loggingContentGenerator.js';
@@ -57,11 +58,12 @@ export type ContentGeneratorConfig = {
   proxy?: string;
 };
 
-export function createContentGeneratorConfig(
+export async function createContentGeneratorConfig(
   config: Config,
   authType: AuthType | undefined,
-): ContentGeneratorConfig {
-  const geminiApiKey = process.env['GEMINI_API_KEY'] || undefined;
+): Promise<ContentGeneratorConfig> {
+  const geminiApiKey =
+    (await loadApiKey()) || process.env['GEMINI_API_KEY'] || undefined;
   const googleApiKey = process.env['GOOGLE_API_KEY'] || undefined;
   const googleCloudProject =
     process.env['GOOGLE_CLOUD_PROJECT'] ||

packages/core/src/index.ts

@@ -37,6 +37,7 @@ export * from './code_assist/codeAssist.js';
 export * from './code_assist/oauth2.js';
 export * from './code_assist/server.js';
 export * from './code_assist/types.js';
+export * from './core/apiKeyCredentialStorage.js';
 
 // Export utilities
 export * from './utils/paths.js';