All files / packages/gemini-core/src/code_assist oauth-credential-storage.ts

18.42% Statements 7/38
0% Branches 0/22
0% Functions 0/6
18.42% Lines 7/38

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129              6x 6x       6x   6x 6x   6x 6x                                                                                                                                                                                                                            
/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */
 
import { type Credentials } from 'google-auth-library';
import { HybridTokenStorage } from '../mcp/token-storage/hybrid-token-storage';
import { OAUTH_FILE, GEMINI_DIR, Storage } from '../config/storage';
import type { OAuthCredentials } from '../mcp/token-storage/types';
import * as path from 'node:path';
import * as os from 'node:os';
import { promises as fs } from 'node:fs';
 
const KEYCHAIN_SERVICE_NAME = 'gemini-cli-oauth';
const MAIN_ACCOUNT_KEY = 'main-account';
 
export class OAuthCredentialStorage {
  private static storage: HybridTokenStorage = new HybridTokenStorage(
    KEYCHAIN_SERVICE_NAME,
  );
 
  /**
   * Load cached OAuth credentials
   */
  static async loadCredentials(): Promise<Credentials | null> {
    try {
      const credentials = await this.storage.getCredentials(MAIN_ACCOUNT_KEY);
 
      Iif (credentials?.token) {
        const { accessToken, refreshToken, expiresAt, tokenType, scope } =
          credentials.token;
        // Convert from OAuthCredentials format to Google Credentials format
        const googleCreds: Credentials = {
          access_token: accessToken,
          refresh_token: refreshToken || undefined,
          token_type: tokenType || undefined,
          scope: scope || undefined,
        };
 
        Iif (expiresAt) {
          googleCreds.expiry_date = expiresAt;
        }
 
        return googleCreds;
      }
 
      // Fallback: Try to migrate from old file-based storage
      return await this.migrateFromFileStorage();
    } catch (error: unknown) {
      console.error(error);
      throw new Error('Failed to load OAuth credentials');
    }
  }
 
  /**
   * Save OAuth credentials
   */
  static async saveCredentials(credentials: Credentials): Promise<void> {
    Iif (!credentials.access_token) {
      throw new Error('Attempted to save credentials without an access token.');
    }
 
    // Convert Google Credentials to OAuthCredentials format
    const mcpCredentials: OAuthCredentials = {
      serverName: MAIN_ACCOUNT_KEY,
      token: {
        accessToken: credentials.access_token,
        refreshToken: credentials.refresh_token || undefined,
        tokenType: credentials.token_type || 'Bearer',
        scope: credentials.scope || undefined,
        expiresAt: credentials.expiry_date || undefined,
      },
      updatedAt: Date.now(),
    };
 
    await this.storage.setCredentials(mcpCredentials);
  }
 
  /**
   * Clear cached OAuth credentials
   */
  static async clearCredentials(): Promise<void> {
    try {
      await this.storage.deleteCredentials(MAIN_ACCOUNT_KEY);
 
      // Also try to remove the old file if it exists
      const oldFilePath = Storage.getOAuthCredsPath();
      await fs.rm(oldFilePath, { force: true }).catch(() => {});
    } catch (error: unknown) {
      console.error(error);
      throw new Error('Failed to clear OAuth credentials');
    }
  }
 
  /**
   * Migrate credentials from old file-based storage to keychain
   */
  private static async migrateFromFileStorage(): Promise<Credentials | null> {
    const oldFilePath = Storage.getOAuthCredsPath();
 
    let credsJson: string;
    try {
      credsJson = await fs.readFile(oldFilePath, 'utf-8');
    } catch (error: unknown) {
      Iif (
        typeof error === 'object' &&
        error !== null &&
        'code' in error &&
        error.code === 'ENOENT'
      ) {
        // File doesn't exist, so no migration.
        return null;
      }
      // Other read errors should propagate.
      throw error;
    }
 
    const credentials = JSON.parse(credsJson) as Credentials;
 
    // Save to new storage
    await this.saveCredentials(credentials);
 
    // Remove old file after successful migration
    await fs.rm(oldFilePath, { force: true }).catch(() => {});
 
    return credentials;
  }
}