@office-ai/aioncli-core 0.2.2 → 0.8.1

package/dist/src/mcp/oauth-token-storage.js

@@ -5,27 +5,28 @@
  */
 import { promises as fs } from 'node:fs';
 import * as path from 'node:path';
-import * as os from 'node:os';
+import { Storage } from '../config/storage.js';
 import { getErrorMessage } from '../utils/errors.js';
+import { HybridTokenStorage } from './token-storage/hybrid-token-storage.js';
+import { DEFAULT_SERVICE_NAME, FORCE_ENCRYPTED_FILE_ENV_VAR, } from './token-storage/index.js';
 /**
  * Class for managing MCP OAuth token storage and retrieval.
  */
 export class MCPOAuthTokenStorage {
-    static TOKEN_FILE = 'mcp-oauth-tokens.json';
-    static CONFIG_DIR = '.gemini';
+    hybridTokenStorage = new HybridTokenStorage(DEFAULT_SERVICE_NAME);
+    useEncryptedFile = process.env[FORCE_ENCRYPTED_FILE_ENV_VAR] === 'true';
     /**
      * Get the path to the token storage file.
      *
      * @returns The full path to the token storage file
      */
-    static getTokenFilePath() {
-        const homeDir = os.homedir();
-        return path.join(homeDir, this.CONFIG_DIR, this.TOKEN_FILE);
+    getTokenFilePath() {
+        return Storage.getMcpOAuthTokensPath();
     }
     /**
      * Ensure the config directory exists.
      */
-    static async ensureConfigDir() {
+    async ensureConfigDir() {
         const configDir = path.dirname(this.getTokenFilePath());
         await fs.mkdir(configDir, { recursive: true });
     }
@@ -34,7 +35,10 @@ export class MCPOAuthTokenStorage {
      *
      * @returns A map of server names to credentials
      */
-    static async loadTokens() {
+    async getAllCredentials() {
+        if (this.useEncryptedFile) {
+            return this.hybridTokenStorage.getAllCredentials();
+        }
         const tokenMap = new Map();
         try {
             const tokenFile = this.getTokenFilePath();
@@ -52,6 +56,29 @@ export class MCPOAuthTokenStorage {
         }
         return tokenMap;
     }
+    async listServers() {
+        if (this.useEncryptedFile) {
+            return this.hybridTokenStorage.listServers();
+        }
+        const tokens = await this.getAllCredentials();
+        return Array.from(tokens.keys());
+    }
+    async setCredentials(credentials) {
+        if (this.useEncryptedFile) {
+            return this.hybridTokenStorage.setCredentials(credentials);
+        }
+        const tokens = await this.getAllCredentials();
+        tokens.set(credentials.serverName, credentials);
+        const tokenArray = Array.from(tokens.values());
+        const tokenFile = this.getTokenFilePath();
+        try {
+            await fs.writeFile(tokenFile, JSON.stringify(tokenArray, null, 2), { mode: 0o600 });
+        }
+        catch (error) {
+            console.error(`Failed to save MCP OAuth token: ${getErrorMessage(error)}`);
+            throw error;
+        }
+    }
     /**
      * Save a token for a specific MCP server.
      *
@@ -61,9 +88,8 @@ export class MCPOAuthTokenStorage {
      * @param tokenUrl Optional token URL used for this token
      * @param mcpServerUrl Optional MCP server URL
      */
-    static async saveToken(serverName, token, clientId, tokenUrl, mcpServerUrl) {
+    async saveToken(serverName, token, clientId, tokenUrl, mcpServerUrl) {
         await this.ensureConfigDir();
-        const tokens = await this.loadTokens();
         const credential = {
             serverName,
             token,
@@ -72,16 +98,10 @@ export class MCPOAuthTokenStorage {
             mcpServerUrl,
             updatedAt: Date.now(),
         };
-        tokens.set(serverName, credential);
-        const tokenArray = Array.from(tokens.values());
-        const tokenFile = this.getTokenFilePath();
-        try {
-            await fs.writeFile(tokenFile, JSON.stringify(tokenArray, null, 2), { mode: 0o600 });
-        }
-        catch (error) {
-            console.error(`Failed to save MCP OAuth token: ${getErrorMessage(error)}`);
-            throw error;
+        if (this.useEncryptedFile) {
+            return this.hybridTokenStorage.setCredentials(credential);
         }
+        await this.setCredentials(credential);
     }
     /**
      * Get a token for a specific MCP server.
@@ -89,8 +109,11 @@ export class MCPOAuthTokenStorage {
      * @param serverName The name of the MCP server
      * @returns The stored credentials or null if not found
      */
-    static async getToken(serverName) {
-        const tokens = await this.loadTokens();
+    async getCredentials(serverName) {
+        if (this.useEncryptedFile) {
+            return this.hybridTokenStorage.getCredentials(serverName);
+        }
+        const tokens = await this.getAllCredentials();
         return tokens.get(serverName) || null;
     }
     /**
@@ -98,8 +121,11 @@ export class MCPOAuthTokenStorage {
      *
      * @param serverName The name of the MCP server
      */
-    static async removeToken(serverName) {
-        const tokens = await this.loadTokens();
+    async deleteCredentials(serverName) {
+        if (this.useEncryptedFile) {
+            return this.hybridTokenStorage.deleteCredentials(serverName);
+        }
+        const tokens = await this.getAllCredentials();
         if (tokens.delete(serverName)) {
             const tokenArray = Array.from(tokens.values());
             const tokenFile = this.getTokenFilePath();
@@ -125,7 +151,7 @@ export class MCPOAuthTokenStorage {
      * @param token The token to check
      * @returns True if the token is expired
      */
-    static isTokenExpired(token) {
+    isTokenExpired(token) {
         if (!token.expiresAt) {
             return false; // No expiry, assume valid
         }
@@ -136,7 +162,10 @@ export class MCPOAuthTokenStorage {
     /**
      * Clear all stored MCP OAuth tokens.
      */
-    static async clearAllTokens() {
+    async clearAll() {
+        if (this.useEncryptedFile) {
+            return this.hybridTokenStorage.clearAll();
+        }
         try {
             const tokenFile = this.getTokenFilePath();
             await fs.unlink(tokenFile);

package/dist/src/mcp/oauth-token-storage.js.map

@@ -1 +1 @@
-{"version":3,"file":"oauth-token-storage.js","sourceRoot":"","sources":["../../../src/mcp/oauth-token-storage.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAEH,OAAO,EAAE,QAAQ,IAAI,EAAE,EAAE,MAAM,SAAS,CAAC;AACzC,OAAO,KAAK,IAAI,MAAM,WAAW,CAAC;AAClC,OAAO,KAAK,EAAE,MAAM,SAAS,CAAC;AAC9B,OAAO,EAAE,eAAe,EAAE,MAAM,oBAAoB,CAAC;AAyBrD;;GAEG;AACH,MAAM,OAAO,oBAAoB;IACvB,MAAM,CAAU,UAAU,GAAG,uBAAuB,CAAC;IACrD,MAAM,CAAU,UAAU,GAAG,SAAS,CAAC;IAE/C;;;;OAIG;IACK,MAAM,CAAC,gBAAgB;QAC7B,MAAM,OAAO,GAAG,EAAE,CAAC,OAAO,EAAE,CAAC;QAC7B,OAAO,IAAI,CAAC,IAAI,CAAC,OAAO,EAAE,IAAI,CAAC,UAAU,EAAE,IAAI,CAAC,UAAU,CAAC,CAAC;IAC9D,CAAC;IAED;;OAEG;IACK,MAAM,CAAC,KAAK,CAAC,eAAe;QAClC,MAAM,SAAS,GAAG,IAAI,CAAC,OAAO,CAAC,IAAI,CAAC,gBAAgB,EAAE,CAAC,CAAC;QACxD,MAAM,EAAE,CAAC,KAAK,CAAC,SAAS,EAAE,EAAE,SAAS,EAAE,IAAI,EAAE,CAAC,CAAC;IACjD,CAAC;IAED;;;;OAIG;IACH,MAAM,CAAC,KAAK,CAAC,UAAU;QACrB,MAAM,QAAQ,GAAG,IAAI,GAAG,EAA+B,CAAC;QAExD,IAAI,CAAC;YACH,MAAM,SAAS,GAAG,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1C,MAAM,IAAI,GAAG,MAAM,EAAE,CAAC,QAAQ,CAAC,SAAS,EAAE,OAAO,CAAC,CAAC;YACnD,MAAM,MAAM,GAAG,IAAI,CAAC,KAAK,CAAC,IAAI,CAA0B,CAAC;YAEzD,KAAK,MAAM,UAAU,IAAI,MAAM,EAAE,CAAC;gBAChC,QAAQ,CAAC,GAAG,CAAC,UAAU,CAAC,UAAU,EAAE,UAAU,CAAC,CAAC;YAClD,CAAC;QACH,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,qDAAqD;YACrD,IAAK,KAA+B,CAAC,IAAI,KAAK,QAAQ,EAAE,CAAC;gBACvD,OAAO,CAAC,KAAK,CACX,oCAAoC,eAAe,CAAC,KAAK,CAAC,EAAE,CAC7D,CAAC;YACJ,CAAC;QACH,CAAC;QAED,OAAO,QAAQ,CAAC;IAClB,CAAC;IAED;;;;;;;;OAQG;IACH,MAAM,CAAC,KAAK,CAAC,SAAS,CACpB,UAAkB,EAClB,KAAoB,EACpB,QAAiB,EACjB,QAAiB,EACjB,YAAqB;QAErB,MAAM,IAAI,CAAC,eAAe,EAAE,CAAC;QAE7B,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,UAAU,EAAE,CAAC;QAEvC,MAAM,UAAU,GAAwB;YACtC,UAAU;YACV,KAAK;YACL,QAAQ;YACR,QAAQ;YACR,YAAY;YACZ,SAAS,EAAE,IAAI,CAAC,GAAG,EAAE;SACtB,CAAC;QAEF,MAAM,CAAC,GAAG,CAAC,UAAU,EAAE,UAAU,CAAC,CAAC;QAEnC,MAAM,UAAU,GAAG,KAAK,CAAC,IAAI,CAAC,MAAM,CAAC,MAAM,EAAE,CAAC,CAAC;QAC/C,MAAM,SAAS,GAAG,IAAI,CAAC,gBAAgB,EAAE,CAAC;QAE1C,IAAI,CAAC;YACH,MAAM,EAAE,CAAC,SAAS,CAChB,SAAS,EACT,IAAI,CAAC,SAAS,CAAC,UAAU,EAAE,IAAI,EAAE,CAAC,CAAC,EACnC,EAAE,IAAI,EAAE,KAAK,EAAE,CAChB,CAAC;QACJ,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,OAAO,CAAC,KAAK,CACX,mCAAmC,eAAe,CAAC,KAAK,CAAC,EAAE,CAC5D,CAAC;YACF,MAAM,KAAK,CAAC;QACd,CAAC;IACH,CAAC;IAED;;;;;OAKG;IACH,MAAM,CAAC,KAAK,CAAC,QAAQ,CACnB,UAAkB;QAElB,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,UAAU,EAAE,CAAC;QACvC,OAAO,MAAM,CAAC,GAAG,CAAC,UAAU,CAAC,IAAI,IAAI,CAAC;IACxC,CAAC;IAED;;;;OAIG;IACH,MAAM,CAAC,KAAK,CAAC,WAAW,CAAC,UAAkB;QACzC,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,UAAU,EAAE,CAAC;QAEvC,IAAI,MAAM,CAAC,MAAM,CAAC,UAAU,CAAC,EAAE,CAAC;YAC9B,MAAM,UAAU,GAAG,KAAK,CAAC,IAAI,CAAC,MAAM,CAAC,MAAM,EAAE,CAAC,CAAC;YAC/C,MAAM,SAAS,GAAG,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAE1C,IAAI,CAAC;gBACH,IAAI,UAAU,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;oBAC5B,gCAAgC;oBAChC,MAAM,EAAE,CAAC,MAAM,CAAC,SAAS,CAAC,CAAC;gBAC7B,CAAC;qBAAM,CAAC;oBACN,MAAM,EAAE,CAAC,SAAS,CAAC,SAAS,EAAE,IAAI,CAAC,SAAS,CAAC,UAAU,EAAE,IAAI,EAAE,CAAC,CAAC,EAAE;wBACjE,IAAI,EAAE,KAAK;qBACZ,CAAC,CAAC;gBACL,CAAC;YACH,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBACf,OAAO,CAAC,KAAK,CACX,qCAAqC,eAAe,CAAC,KAAK,CAAC,EAAE,CAC9D,CAAC;YACJ,CAAC;QACH,CAAC;IACH,CAAC;IAED;;;;;OAKG;IACH,MAAM,CAAC,cAAc,CAAC,KAAoB;QACxC,IAAI,CAAC,KAAK,CAAC,SAAS,EAAE,CAAC;YACrB,OAAO,KAAK,CAAC,CAAC,0BAA0B;QAC1C,CAAC;QAED,kDAAkD;QAClD,MAAM,QAAQ,GAAG,CAAC,GAAG,EAAE,GAAG,IAAI,CAAC;QAC/B,OAAO,IAAI,CAAC,GAAG,EAAE,GAAG,QAAQ,IAAI,KAAK,CAAC,SAAS,CAAC;IAClD,CAAC;IAED;;OAEG;IACH,MAAM,CAAC,KAAK,CAAC,cAAc;QACzB,IAAI,CAAC;YACH,MAAM,SAAS,GAAG,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1C,MAAM,EAAE,CAAC,MAAM,CAAC,SAAS,CAAC,CAAC;QAC7B,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,IAAK,KAA+B,CAAC,IAAI,KAAK,QAAQ,EAAE,CAAC;gBACvD,OAAO,CAAC,KAAK,CACX,qCAAqC,eAAe,CAAC,KAAK,CAAC,EAAE,CAC9D,CAAC;YACJ,CAAC;QACH,CAAC;IACH,CAAC"}
+{"version":3,"file":"oauth-token-storage.js","sourceRoot":"","sources":["../../../src/mcp/oauth-token-storage.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAEH,OAAO,EAAE,QAAQ,IAAI,EAAE,EAAE,MAAM,SAAS,CAAC;AACzC,OAAO,KAAK,IAAI,MAAM,WAAW,CAAC;AAClC,OAAO,EAAE,OAAO,EAAE,MAAM,sBAAsB,CAAC;AAC/C,OAAO,EAAE,eAAe,EAAE,MAAM,oBAAoB,CAAC;AAMrD,OAAO,EAAE,kBAAkB,EAAE,MAAM,yCAAyC,CAAC;AAC7E,OAAO,EACL,oBAAoB,EACpB,4BAA4B,GAC7B,MAAM,0BAA0B,CAAC;AAElC;;GAEG;AACH,MAAM,OAAO,oBAAoB;IACd,kBAAkB,GAAG,IAAI,kBAAkB,CAC1D,oBAAoB,CACrB,CAAC;IACe,gBAAgB,GAC/B,OAAO,CAAC,GAAG,CAAC,4BAA4B,CAAC,KAAK,MAAM,CAAC;IAEvD;;;;OAIG;IACK,gBAAgB;QACtB,OAAO,OAAO,CAAC,qBAAqB,EAAE,CAAC;IACzC,CAAC;IAED;;OAEG;IACK,KAAK,CAAC,eAAe;QAC3B,MAAM,SAAS,GAAG,IAAI,CAAC,OAAO,CAAC,IAAI,CAAC,gBAAgB,EAAE,CAAC,CAAC;QACxD,MAAM,EAAE,CAAC,KAAK,CAAC,SAAS,EAAE,EAAE,SAAS,EAAE,IAAI,EAAE,CAAC,CAAC;IACjD,CAAC;IAED;;;;OAIG;IACH,KAAK,CAAC,iBAAiB;QACrB,IAAI,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1B,OAAO,IAAI,CAAC,kBAAkB,CAAC,iBAAiB,EAAE,CAAC;QACrD,CAAC;QACD,MAAM,QAAQ,GAAG,IAAI,GAAG,EAA4B,CAAC;QAErD,IAAI,CAAC;YACH,MAAM,SAAS,GAAG,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1C,MAAM,IAAI,GAAG,MAAM,EAAE,CAAC,QAAQ,CAAC,SAAS,EAAE,OAAO,CAAC,CAAC;YACnD,MAAM,MAAM,GAAG,IAAI,CAAC,KAAK,CAAC,IAAI,CAAuB,CAAC;YAEtD,KAAK,MAAM,UAAU,IAAI,MAAM,EAAE,CAAC;gBAChC,QAAQ,CAAC,GAAG,CAAC,UAAU,CAAC,UAAU,EAAE,UAAU,CAAC,CAAC;YAClD,CAAC;QACH,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,qDAAqD;YACrD,IAAK,KAA+B,CAAC,IAAI,KAAK,QAAQ,EAAE,CAAC;gBACvD,OAAO,CAAC,KAAK,CACX,oCAAoC,eAAe,CAAC,KAAK,CAAC,EAAE,CAC7D,CAAC;YACJ,CAAC;QACH,CAAC;QAED,OAAO,QAAQ,CAAC;IAClB,CAAC;IAED,KAAK,CAAC,WAAW;QACf,IAAI,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1B,OAAO,IAAI,CAAC,kBAAkB,CAAC,WAAW,EAAE,CAAC;QAC/C,CAAC;QACD,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,iBAAiB,EAAE,CAAC;QAC9C,OAAO,KAAK,CAAC,IAAI,CAAC,MAAM,CAAC,IAAI,EAAE,CAAC,CAAC;IACnC,CAAC;IAED,KAAK,CAAC,cAAc,CAAC,WAA6B;QAChD,IAAI,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1B,OAAO,IAAI,CAAC,kBAAkB,CAAC,cAAc,CAAC,WAAW,CAAC,CAAC;QAC7D,CAAC;QACD,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,iBAAiB,EAAE,CAAC;QAC9C,MAAM,CAAC,GAAG,CAAC,WAAW,CAAC,UAAU,EAAE,WAAW,CAAC,CAAC;QAEhD,MAAM,UAAU,GAAG,KAAK,CAAC,IAAI,CAAC,MAAM,CAAC,MAAM,EAAE,CAAC,CAAC;QAC/C,MAAM,SAAS,GAAG,IAAI,CAAC,gBAAgB,EAAE,CAAC;QAE1C,IAAI,CAAC;YACH,MAAM,EAAE,CAAC,SAAS,CAChB,SAAS,EACT,IAAI,CAAC,SAAS,CAAC,UAAU,EAAE,IAAI,EAAE,CAAC,CAAC,EACnC,EAAE,IAAI,EAAE,KAAK,EAAE,CAChB,CAAC;QACJ,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,OAAO,CAAC,KAAK,CACX,mCAAmC,eAAe,CAAC,KAAK,CAAC,EAAE,CAC5D,CAAC;YACF,MAAM,KAAK,CAAC;QACd,CAAC;IACH,CAAC;IAED;;;;;;;;OAQG;IACH,KAAK,CAAC,SAAS,CACb,UAAkB,EAClB,KAAiB,EACjB,QAAiB,EACjB,QAAiB,EACjB,YAAqB;QAErB,MAAM,IAAI,CAAC,eAAe,EAAE,CAAC;QAE7B,MAAM,UAAU,GAAqB;YACnC,UAAU;YACV,KAAK;YACL,QAAQ;YACR,QAAQ;YACR,YAAY;YACZ,SAAS,EAAE,IAAI,CAAC,GAAG,EAAE;SACtB,CAAC;QAEF,IAAI,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1B,OAAO,IAAI,CAAC,kBAAkB,CAAC,cAAc,CAAC,UAAU,CAAC,CAAC;QAC5D,CAAC;QACD,MAAM,IAAI,CAAC,cAAc,CAAC,UAAU,CAAC,CAAC;IACxC,CAAC;IAED;;;;;OAKG;IACH,KAAK,CAAC,cAAc,CAAC,UAAkB;QACrC,IAAI,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1B,OAAO,IAAI,CAAC,kBAAkB,CAAC,cAAc,CAAC,UAAU,CAAC,CAAC;QAC5D,CAAC;QACD,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,iBAAiB,EAAE,CAAC;QAC9C,OAAO,MAAM,CAAC,GAAG,CAAC,UAAU,CAAC,IAAI,IAAI,CAAC;IACxC,CAAC;IAED;;;;OAIG;IACH,KAAK,CAAC,iBAAiB,CAAC,UAAkB;QACxC,IAAI,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1B,OAAO,IAAI,CAAC,kBAAkB,CAAC,iBAAiB,CAAC,UAAU,CAAC,CAAC;QAC/D,CAAC;QACD,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,iBAAiB,EAAE,CAAC;QAE9C,IAAI,MAAM,CAAC,MAAM,CAAC,UAAU,CAAC,EAAE,CAAC;YAC9B,MAAM,UAAU,GAAG,KAAK,CAAC,IAAI,CAAC,MAAM,CAAC,MAAM,EAAE,CAAC,CAAC;YAC/C,MAAM,SAAS,GAAG,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAE1C,IAAI,CAAC;gBACH,IAAI,UAAU,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;oBAC5B,gCAAgC;oBAChC,MAAM,EAAE,CAAC,MAAM,CAAC,SAAS,CAAC,CAAC;gBAC7B,CAAC;qBAAM,CAAC;oBACN,MAAM,EAAE,CAAC,SAAS,CAAC,SAAS,EAAE,IAAI,CAAC,SAAS,CAAC,UAAU,EAAE,IAAI,EAAE,CAAC,CAAC,EAAE;wBACjE,IAAI,EAAE,KAAK;qBACZ,CAAC,CAAC;gBACL,CAAC;YACH,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBACf,OAAO,CAAC,KAAK,CACX,qCAAqC,eAAe,CAAC,KAAK,CAAC,EAAE,CAC9D,CAAC;YACJ,CAAC;QACH,CAAC;IACH,CAAC;IAED;;;;;OAKG;IACH,cAAc,CAAC,KAAiB;QAC9B,IAAI,CAAC,KAAK,CAAC,SAAS,EAAE,CAAC;YACrB,OAAO,KAAK,CAAC,CAAC,0BAA0B;QAC1C,CAAC;QAED,kDAAkD;QAClD,MAAM,QAAQ,GAAG,CAAC,GAAG,EAAE,GAAG,IAAI,CAAC;QAC/B,OAAO,IAAI,CAAC,GAAG,EAAE,GAAG,QAAQ,IAAI,KAAK,CAAC,SAAS,CAAC;IAClD,CAAC;IAED;;OAEG;IACH,KAAK,CAAC,QAAQ;QACZ,IAAI,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1B,OAAO,IAAI,CAAC,kBAAkB,CAAC,QAAQ,EAAE,CAAC;QAC5C,CAAC;QACD,IAAI,CAAC;YACH,MAAM,SAAS,GAAG,IAAI,CAAC,gBAAgB,EAAE,CAAC;YAC1C,MAAM,EAAE,CAAC,MAAM,CAAC,SAAS,CAAC,CAAC;QAC7B,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,IAAK,KAA+B,CAAC,IAAI,KAAK,QAAQ,EAAE,CAAC;gBACvD,OAAO,CAAC,KAAK,CACX,qCAAqC,eAAe,CAAC,KAAK,CAAC,EAAE,CAC9D,CAAC;YACJ,CAAC;QACH,CAAC;IACH,CAAC;CACF"}

package/dist/src/mcp/oauth-token-storage.test.js

@@ -3,11 +3,13 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import { promises as fs } from 'node:fs';
 import * as path from 'node:path';
-import { MCPOAuthTokenStorage, } from './oauth-token-storage.js';
-// Mock file system operations
+import { MCPOAuthTokenStorage } from './oauth-token-storage.js';
+import { FORCE_ENCRYPTED_FILE_ENV_VAR } from './token-storage/index.js';
+import { GEMINI_DIR } from '../utils/paths.js';
+// Mock dependencies
 vi.mock('node:fs', () => ({
     promises: {
         readFile: vi.fn(),
@@ -16,16 +18,35 @@ vi.mock('node:fs', () => ({
         unlink: vi.fn(),
     },
 }));
-vi.mock('node:os', () => ({
-    homedir: vi.fn(() => '/mock/home'),
+vi.mock('node:path', () => ({
+    dirname: vi.fn(),
+    join: vi.fn(),
 }));
+vi.mock('../config/storage.js', () => ({
+    Storage: {
+        getMcpOAuthTokensPath: vi.fn(),
+    },
+}));
+const mockHybridTokenStorage = {
+    listServers: vi.fn(),
+    setCredentials: vi.fn(),
+    getCredentials: vi.fn(),
+    deleteCredentials: vi.fn(),
+    clearAll: vi.fn(),
+    getAllCredentials: vi.fn(),
+};
+vi.mock('./token-storage/hybrid-token-storage.js', () => ({
+    HybridTokenStorage: vi.fn(() => mockHybridTokenStorage),
+}));
+const ONE_HR_MS = 3600000;
 describe('MCPOAuthTokenStorage', () => {
+    let tokenStorage;
     const mockToken = {
         accessToken: 'access_token_123',
         refreshToken: 'refresh_token_456',
         tokenType: 'Bearer',
         scope: 'read write',
-        expiresAt: Date.now() + 3600000, // 1 hour from now
+        expiresAt: Date.now() + ONE_HR_MS,
     };
     const mockCredentials = {
         serverName: 'test-server',
@@ -34,171 +55,244 @@ describe('MCPOAuthTokenStorage', () => {
         tokenUrl: 'https://auth.example.com/token',
         updatedAt: Date.now(),
     };
-    beforeEach(() => {
-        vi.clearAllMocks();
-        vi.spyOn(console, 'error').mockImplementation(() => { });
-    });
-    afterEach(() => {
-        vi.restoreAllMocks();
-    });
-    describe('loadTokens', () => {
-        it('should return empty map when token file does not exist', async () => {
-            vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' });
-            const tokens = await MCPOAuthTokenStorage.loadTokens();
-            expect(tokens.size).toBe(0);
-            expect(console.error).not.toHaveBeenCalled();
-        });
-        it('should load tokens from file successfully', async () => {
-            const tokensArray = [mockCredentials];
-            vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(tokensArray));
-            const tokens = await MCPOAuthTokenStorage.loadTokens();
-            expect(tokens.size).toBe(1);
-            expect(tokens.get('test-server')).toEqual(mockCredentials);
-            expect(fs.readFile).toHaveBeenCalledWith(path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'), 'utf-8');
-        });
-        it('should handle corrupted token file gracefully', async () => {
-            vi.mocked(fs.readFile).mockResolvedValue('invalid json');
-            const tokens = await MCPOAuthTokenStorage.loadTokens();
-            expect(tokens.size).toBe(0);
-            expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Failed to load MCP OAuth tokens'));
-        });
-        it('should handle file read errors other than ENOENT', async () => {
-            const error = new Error('Permission denied');
-            vi.mocked(fs.readFile).mockRejectedValue(error);
-            const tokens = await MCPOAuthTokenStorage.loadTokens();
-            expect(tokens.size).toBe(0);
-            expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Failed to load MCP OAuth tokens'));
+    describe('with encrypted flag false', () => {
+        beforeEach(() => {
+            vi.stubEnv(FORCE_ENCRYPTED_FILE_ENV_VAR, 'false');
+            tokenStorage = new MCPOAuthTokenStorage();
+            vi.clearAllMocks();
+            vi.spyOn(console, 'error');
         });
-    });
-    describe('saveToken', () => {
-        it('should save token with restricted permissions', async () => {
-            vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' });
-            vi.mocked(fs.mkdir).mockResolvedValue(undefined);
-            vi.mocked(fs.writeFile).mockResolvedValue(undefined);
-            await MCPOAuthTokenStorage.saveToken('test-server', mockToken, 'client-id', 'https://token.url');
-            expect(fs.mkdir).toHaveBeenCalledWith(path.join('/mock/home', '.gemini'), { recursive: true });
-            expect(fs.writeFile).toHaveBeenCalledWith(path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'), expect.stringContaining('test-server'), { mode: 0o600 });
-        });
-        it('should update existing token for same server', async () => {
-            const existingCredentials = {
-                ...mockCredentials,
-                serverName: 'existing-server',
-            };
-            vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([existingCredentials]));
-            vi.mocked(fs.writeFile).mockResolvedValue(undefined);
-            const newToken = { ...mockToken, accessToken: 'new_access_token' };
-            await MCPOAuthTokenStorage.saveToken('existing-server', newToken);
-            const writeCall = vi.mocked(fs.writeFile).mock.calls[0];
-            const savedData = JSON.parse(writeCall[1]);
-            expect(savedData).toHaveLength(1);
-            expect(savedData[0].token.accessToken).toBe('new_access_token');
-            expect(savedData[0].serverName).toBe('existing-server');
-        });
-        it('should handle write errors gracefully', async () => {
-            vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' });
-            vi.mocked(fs.mkdir).mockResolvedValue(undefined);
-            const writeError = new Error('Disk full');
-            vi.mocked(fs.writeFile).mockRejectedValue(writeError);
-            await expect(MCPOAuthTokenStorage.saveToken('test-server', mockToken)).rejects.toThrow('Disk full');
-            expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Failed to save MCP OAuth token'));
+        afterEach(() => {
+            vi.unstubAllEnvs();
+            vi.restoreAllMocks();
         });
-    });
-    describe('getToken', () => {
-        it('should return token for existing server', async () => {
-            vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([mockCredentials]));
-            const result = await MCPOAuthTokenStorage.getToken('test-server');
-            expect(result).toEqual(mockCredentials);
-        });
-        it('should return null for non-existent server', async () => {
-            vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([mockCredentials]));
-            const result = await MCPOAuthTokenStorage.getToken('non-existent');
-            expect(result).toBeNull();
-        });
-        it('should return null when no tokens file exists', async () => {
-            vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' });
-            const result = await MCPOAuthTokenStorage.getToken('test-server');
-            expect(result).toBeNull();
+        describe('getAllCredentials', () => {
+            it('should return empty map when token file does not exist', async () => {
+                vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' });
+                const tokens = await tokenStorage.getAllCredentials();
+                expect(tokens.size).toBe(0);
+                expect(console.error).not.toHaveBeenCalled();
+            });
+            it('should load tokens from file successfully', async () => {
+                const tokensArray = [mockCredentials];
+                vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(tokensArray));
+                const tokens = await tokenStorage.getAllCredentials();
+                expect(tokens.size).toBe(1);
+                expect(tokens.get('test-server')).toEqual(mockCredentials);
+                expect(fs.readFile).toHaveBeenCalledWith(path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'), 'utf-8');
+            });
+            it('should handle corrupted token file gracefully', async () => {
+                vi.mocked(fs.readFile).mockResolvedValue('invalid json');
+                const tokens = await tokenStorage.getAllCredentials();
+                expect(tokens.size).toBe(0);
+                expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Failed to load MCP OAuth tokens'));
+            });
+            it('should handle file read errors other than ENOENT', async () => {
+                const error = new Error('Permission denied');
+                vi.mocked(fs.readFile).mockRejectedValue(error);
+                const tokens = await tokenStorage.getAllCredentials();
+                expect(tokens.size).toBe(0);
+                expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Failed to load MCP OAuth tokens'));
+            });
         });
-    });
-    describe('removeToken', () => {
-        it('should remove token for specific server', async () => {
-            const credentials1 = { ...mockCredentials, serverName: 'server1' };
-            const credentials2 = { ...mockCredentials, serverName: 'server2' };
-            vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([credentials1, credentials2]));
-            vi.mocked(fs.writeFile).mockResolvedValue(undefined);
-            await MCPOAuthTokenStorage.removeToken('server1');
-            const writeCall = vi.mocked(fs.writeFile).mock.calls[0];
-            const savedData = JSON.parse(writeCall[1]);
-            expect(savedData).toHaveLength(1);
-            expect(savedData[0].serverName).toBe('server2');
-        });
-        it('should remove token file when no tokens remain', async () => {
-            vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([mockCredentials]));
-            vi.mocked(fs.unlink).mockResolvedValue(undefined);
-            await MCPOAuthTokenStorage.removeToken('test-server');
-            expect(fs.unlink).toHaveBeenCalledWith(path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'));
-            expect(fs.writeFile).not.toHaveBeenCalled();
-        });
-        it('should handle removal of non-existent token gracefully', async () => {
-            vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([mockCredentials]));
-            await MCPOAuthTokenStorage.removeToken('non-existent');
-            expect(fs.writeFile).not.toHaveBeenCalled();
-            expect(fs.unlink).not.toHaveBeenCalled();
-        });
-        it('should handle file operation errors gracefully', async () => {
-            vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([mockCredentials]));
-            vi.mocked(fs.unlink).mockRejectedValue(new Error('Permission denied'));
-            await MCPOAuthTokenStorage.removeToken('test-server');
-            expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Failed to remove MCP OAuth token'));
+        describe('saveToken', () => {
+            it('should save token with restricted permissions', async () => {
+                vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' });
+                vi.mocked(fs.mkdir).mockResolvedValue(undefined);
+                vi.mocked(fs.writeFile).mockResolvedValue(undefined);
+                await tokenStorage.saveToken('test-server', mockToken, 'client-id', 'https://token.url');
+                expect(fs.mkdir).toHaveBeenCalledWith(path.join('/mock/home', GEMINI_DIR), { recursive: true });
+                expect(fs.writeFile).toHaveBeenCalledWith(path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'), expect.stringContaining('test-server'), { mode: 0o600 });
+            });
+            it('should update existing token for same server', async () => {
+                const existingCredentials = {
+                    ...mockCredentials,
+                    serverName: 'existing-server',
+                };
+                vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([existingCredentials]));
+                vi.mocked(fs.writeFile).mockResolvedValue(undefined);
+                const newToken = {
+                    ...mockToken,
+                    accessToken: 'new_access_token',
+                };
+                await tokenStorage.saveToken('existing-server', newToken);
+                const writeCall = vi.mocked(fs.writeFile).mock.calls[0];
+                const savedData = JSON.parse(writeCall[1]);
+                expect(savedData).toHaveLength(1);
+                expect(savedData[0].token.accessToken).toBe('new_access_token');
+                expect(savedData[0].serverName).toBe('existing-server');
+            });
+            it('should handle write errors gracefully', async () => {
+                vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' });
+                vi.mocked(fs.mkdir).mockResolvedValue(undefined);
+                const writeError = new Error('Disk full');
+                vi.mocked(fs.writeFile).mockRejectedValue(writeError);
+                await expect(tokenStorage.saveToken('test-server', mockToken)).rejects.toThrow('Disk full');
+                expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Failed to save MCP OAuth token'));
+            });
+        });
+        describe('getCredentials', () => {
+            it('should return token for existing server', async () => {
+                vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([mockCredentials]));
+                const result = await tokenStorage.getCredentials('test-server');
+                expect(result).toEqual(mockCredentials);
+            });
+            it('should return null for non-existent server', async () => {
+                vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([mockCredentials]));
+                const result = await tokenStorage.getCredentials('non-existent');
+                expect(result).toBeNull();
+            });
+            it('should return null when no tokens file exists', async () => {
+                vi.mocked(fs.readFile).mockRejectedValue({ code: 'ENOENT' });
+                const result = await tokenStorage.getCredentials('test-server');
+                expect(result).toBeNull();
+            });
+        });
+        describe('deleteCredentials', () => {
+            it('should remove token for specific server', async () => {
+                const credentials1 = {
+                    ...mockCredentials,
+                    serverName: 'server1',
+                };
+                const credentials2 = {
+                    ...mockCredentials,
+                    serverName: 'server2',
+                };
+                vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([credentials1, credentials2]));
+                vi.mocked(fs.writeFile).mockResolvedValue(undefined);
+                await tokenStorage.deleteCredentials('server1');
+                const writeCall = vi.mocked(fs.writeFile).mock.calls[0];
+                const savedData = JSON.parse(writeCall[1]);
+                expect(savedData).toHaveLength(1);
+                expect(savedData[0].serverName).toBe('server2');
+            });
+            it('should remove token file when no tokens remain', async () => {
+                vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([mockCredentials]));
+                vi.mocked(fs.unlink).mockResolvedValue(undefined);
+                await tokenStorage.deleteCredentials('test-server');
+                expect(fs.unlink).toHaveBeenCalledWith(path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'));
+                expect(fs.writeFile).not.toHaveBeenCalled();
+            });
+            it('should handle removal of non-existent token gracefully', async () => {
+                vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([mockCredentials]));
+                await tokenStorage.deleteCredentials('non-existent');
+                expect(fs.writeFile).not.toHaveBeenCalled();
+                expect(fs.unlink).not.toHaveBeenCalled();
+            });
+            it('should handle file operation errors gracefully', async () => {
+                vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify([mockCredentials]));
+                vi.mocked(fs.unlink).mockRejectedValue(new Error('Permission denied'));
+                await tokenStorage.deleteCredentials('test-server');
+                expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Failed to remove MCP OAuth token'));
+            });
+        });
+        describe('isTokenExpired', () => {
+            it('should return false for token without expiry', () => {
+                const tokenWithoutExpiry = { ...mockToken };
+                delete tokenWithoutExpiry.expiresAt;
+                const result = tokenStorage.isTokenExpired(tokenWithoutExpiry);
+                expect(result).toBe(false);
+            });
+            it('should return false for valid token', () => {
+                const futureToken = {
+                    ...mockToken,
+                    expiresAt: Date.now() + ONE_HR_MS,
+                };
+                const result = tokenStorage.isTokenExpired(futureToken);
+                expect(result).toBe(false);
+            });
+            it('should return true for expired token', () => {
+                const expiredToken = {
+                    ...mockToken,
+                    expiresAt: Date.now() - ONE_HR_MS,
+                };
+                const result = tokenStorage.isTokenExpired(expiredToken);
+                expect(result).toBe(true);
+            });
+            it('should return true for token expiring within buffer time', () => {
+                const soonToExpireToken = {
+                    ...mockToken,
+                    expiresAt: Date.now() + 60000, // 1 minute from now (within 5-minute buffer)
+                };
+                const result = tokenStorage.isTokenExpired(soonToExpireToken);
+                expect(result).toBe(true);
+            });
+        });
+        describe('clearAll', () => {
+            it('should remove token file successfully', async () => {
+                vi.mocked(fs.unlink).mockResolvedValue(undefined);
+                await tokenStorage.clearAll();
+                expect(fs.unlink).toHaveBeenCalledWith(path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'));
+            });
+            it('should handle non-existent file gracefully', async () => {
+                vi.mocked(fs.unlink).mockRejectedValue({ code: 'ENOENT' });
+                await tokenStorage.clearAll();
+                expect(console.error).not.toHaveBeenCalled();
+            });
+            it('should handle other file errors gracefully', async () => {
+                vi.mocked(fs.unlink).mockRejectedValue(new Error('Permission denied'));
+                await tokenStorage.clearAll();
+                expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Failed to clear MCP OAuth tokens'));
+            });
         });
     });
-    describe('isTokenExpired', () => {
-        it('should return false for token without expiry', () => {
-            const tokenWithoutExpiry = { ...mockToken };
-            delete tokenWithoutExpiry.expiresAt;
-            const result = MCPOAuthTokenStorage.isTokenExpired(tokenWithoutExpiry);
-            expect(result).toBe(false);
-        });
-        it('should return false for valid token', () => {
-            const futureToken = {
-                ...mockToken,
-                expiresAt: Date.now() + 3600000, // 1 hour from now
-            };
-            const result = MCPOAuthTokenStorage.isTokenExpired(futureToken);
-            expect(result).toBe(false);
+    describe('with encrypted flag true', () => {
+        beforeEach(() => {
+            vi.stubEnv(FORCE_ENCRYPTED_FILE_ENV_VAR, 'true');
+            tokenStorage = new MCPOAuthTokenStorage();
+            vi.clearAllMocks();
+            vi.spyOn(console, 'error');
         });
-        it('should return true for expired token', () => {
-            const expiredToken = {
-                ...mockToken,
-                expiresAt: Date.now() - 3600000, // 1 hour ago
-            };
-            const result = MCPOAuthTokenStorage.isTokenExpired(expiredToken);
-            expect(result).toBe(true);
+        afterEach(() => {
+            vi.unstubAllEnvs();
+            vi.restoreAllMocks();
+        });
+        it('should use HybridTokenStorage to list all credentials', async () => {
+            mockHybridTokenStorage.getAllCredentials.mockResolvedValue(new Map());
+            const servers = await tokenStorage.getAllCredentials();
+            expect(mockHybridTokenStorage.getAllCredentials).toHaveBeenCalled();
+            expect(servers).toEqual(new Map());
         });
-        it('should return true for token expiring within buffer time', () => {
-            const soonToExpireToken = {
-                ...mockToken,
-                expiresAt: Date.now() + 60000, // 1 minute from now (within 5-minute buffer)
+        it('should use HybridTokenStorage to list servers', async () => {
+            mockHybridTokenStorage.listServers.mockResolvedValue(['server1']);
+            const servers = await tokenStorage.listServers();
+            expect(mockHybridTokenStorage.listServers).toHaveBeenCalled();
+            expect(servers).toEqual(['server1']);
+        });
+        it('should use HybridTokenStorage to set credentials', async () => {
+            await tokenStorage.setCredentials(mockCredentials);
+            expect(mockHybridTokenStorage.setCredentials).toHaveBeenCalledWith(mockCredentials);
+        });
+        it('should use HybridTokenStorage to save a token', async () => {
+            const serverName = 'server1';
+            const now = Date.now();
+            vi.spyOn(Date, 'now').mockReturnValue(now);
+            await tokenStorage.saveToken(serverName, mockToken, 'clientId', 'tokenUrl', 'mcpUrl');
+            const expectedCredential = {
+                serverName,
+                token: mockToken,
+                clientId: 'clientId',
+                tokenUrl: 'tokenUrl',
+                mcpServerUrl: 'mcpUrl',
+                updatedAt: now,
             };
-            const result = MCPOAuthTokenStorage.isTokenExpired(soonToExpireToken);
-            expect(result).toBe(true);
+            expect(mockHybridTokenStorage.setCredentials).toHaveBeenCalledWith(expectedCredential);
+            expect(path.dirname).toHaveBeenCalled();
+            expect(fs.mkdir).toHaveBeenCalled();
         });
-    });
-    describe('clearAllTokens', () => {
-        it('should remove token file successfully', async () => {
-            vi.mocked(fs.unlink).mockResolvedValue(undefined);
-            await MCPOAuthTokenStorage.clearAllTokens();
-            expect(fs.unlink).toHaveBeenCalledWith(path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'));
-        });
-        it('should handle non-existent file gracefully', async () => {
-            vi.mocked(fs.unlink).mockRejectedValue({ code: 'ENOENT' });
-            await MCPOAuthTokenStorage.clearAllTokens();
-            expect(console.error).not.toHaveBeenCalled();
-        });
-        it('should handle other file errors gracefully', async () => {
-            vi.mocked(fs.unlink).mockRejectedValue(new Error('Permission denied'));
-            await MCPOAuthTokenStorage.clearAllTokens();
-            expect(console.error).toHaveBeenCalledWith(expect.stringContaining('Failed to clear MCP OAuth tokens'));
+        it('should use HybridTokenStorage to get credentials', async () => {
+            mockHybridTokenStorage.getCredentials.mockResolvedValue(mockCredentials);
+            const result = await tokenStorage.getCredentials('server1');
+            expect(mockHybridTokenStorage.getCredentials).toHaveBeenCalledWith('server1');
+            expect(result).toBe(mockCredentials);
+        });
+        it('should use HybridTokenStorage to delete credentials', async () => {
+            await tokenStorage.deleteCredentials('server1');
+            expect(mockHybridTokenStorage.deleteCredentials).toHaveBeenCalledWith('server1');
+        });
+        it('should use HybridTokenStorage to clear all tokens', async () => {
+            await tokenStorage.clearAll();
+            expect(mockHybridTokenStorage.clearAll).toHaveBeenCalled();
         });
     });
 });