All files / packages/gemini-core/src/mcp/token-storage hybrid-token-storage.ts

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
 
 
 
 
 
 
6x
6x
 
6x
 
6x
 
6x
6x
6x
6x
 
 
6x
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */
 
import { BaseTokenStorage } from './base-token-storage';
import { FileTokenStorage } from './file-token-storage';
import type { TokenStorage, OAuthCredentials } from './types';
import { TokenStorageType } from './types';
 
const FORCE_FILE_STORAGE_ENV_VAR = 'GEMINI_FORCE_FILE_STORAGE';
 
export class HybridTokenStorage extends BaseTokenStorage {
  private storage: TokenStorage | null = null;
  private storageType: TokenStorageType | null = null;
  private storageInitPromise: Promise<TokenStorage> | null = null;
 
  constructor(serviceName: string) {
    super(serviceName);
  }
 
  private async initializeStorage(): Promise<TokenStorage> {
    const forceFileStorage = process.env[FORCE_FILE_STORAGE_ENV_VAR] === 'true';
 
    Iif (!forceFileStorage) {
      try {
        const { KeychainTokenStorage } = await import(
          './keychain-token-storage'
        );
        const keychainStorage = new KeychainTokenStorage(this.serviceName);
 
        const isAvailable = await keychainStorage.isAvailable();
        Iif (isAvailable) {
          this.storage = keychainStorage;
          this.storageType = TokenStorageType.KEYCHAIN;
          return this.storage;
        }
      } catch (_e) {
        // Fallback to file storage if keychain fails to initialize
      }
    }
 
    this.storage = new FileTokenStorage(this.serviceName);
    this.storageType = TokenStorageType.ENCRYPTED_FILE;
    return this.storage;
  }
 
  private async getStorage(): Promise<TokenStorage> {
    Iif (this.storage !== null) {
      return this.storage;
    }
 
    // Use a single initialization promise to avoid race conditions
    Iif (!this.storageInitPromise) {
      this.storageInitPromise = this.initializeStorage();
    }
 
    // Wait for initialization to complete
    return await this.storageInitPromise;
  }
 
  async getCredentials(serverName: string): Promise<OAuthCredentials | null> {
    const storage = await this.getStorage();
    return storage.getCredentials(serverName);
  }
 
  async setCredentials(credentials: OAuthCredentials): Promise<void> {
    const storage = await this.getStorage();
    await storage.setCredentials(credentials);
  }
 
  async deleteCredentials(serverName: string): Promise<void> {
    const storage = await this.getStorage();
    await storage.deleteCredentials(serverName);
  }
 
  async listServers(): Promise<string[]> {
    const storage = await this.getStorage();
    return storage.listServers();
  }
 
  async getAllCredentials(): Promise<Map<string, OAuthCredentials>> {
    const storage = await this.getStorage();
    return storage.getAllCredentials();
  }
 
  async clearAll(): Promise<void> {
    const storage = await this.getStorage();
    await storage.clearAll();
  }
 
  async getStorageType(): Promise<TokenStorageType> {
    await this.getStorage();
    return this.storageType!;
  }
}