feat(security) - Make oauth token storage implement the shared interface (#7802)

Co-authored-by: Shi Shu <shii@google.com>

packages/core/src/mcp/oauth-token-storage.ts

@@ -8,12 +8,16 @@ import { promises as fs } from 'node:fs';
 import * as path from 'node:path';
 import { Storage } from '../config/storage.js';
 import { getErrorMessage } from '../utils/errors.js';
-import type { OAuthToken, OAuthCredentials } from './token-storage/types.js';
+import type {
+  OAuthToken,
+  OAuthCredentials,
+  TokenStorage,
+} from './token-storage/types.js';
 
 /**
  * Class for managing MCP OAuth token storage and retrieval.
  */
-export class MCPOAuthTokenStorage {
+export class MCPOAuthTokenStorage implements TokenStorage {
   /**
    * Get the path to the token storage file.
    *
@@ -36,7 +40,7 @@ export class MCPOAuthTokenStorage {
    *
    * @returns A map of server names to credentials
    */
-  async loadTokens(): Promise<Map<string, OAuthCredentials>> {
+  async getAllCredentials(): Promise<Map<string, OAuthCredentials>> {
     const tokenMap = new Map<string, OAuthCredentials>();
 
     try {
@@ -59,36 +63,14 @@ export class MCPOAuthTokenStorage {
     return tokenMap;
   }
 
-  /**
-   * Save a token for a specific MCP server.
-   *
-   * @param serverName The name of the MCP server
-   * @param token The OAuth token to save
-   * @param clientId Optional client ID used for this token
-   * @param tokenUrl Optional token URL used for this token
-   * @param mcpServerUrl Optional MCP server URL
-   */
-  async saveToken(
-    serverName: string,
-    token: OAuthToken,
-    clientId?: string,
-    tokenUrl?: string,
-    mcpServerUrl?: string,
-  ): Promise<void> {
-    await this.ensureConfigDir();
+  async listServers(): Promise<string[]> {
+    const tokens = await this.getAllCredentials();
+    return Array.from(tokens.keys());
+  }
 
-    const tokens = await this.loadTokens();
-
-    const credential: OAuthCredentials = {
-      serverName,
-      token,
-      clientId,
-      tokenUrl,
-      mcpServerUrl,
-      updatedAt: Date.now(),
-    };
-
-    tokens.set(serverName, credential);
+  async setCredentials(credentials: OAuthCredentials): Promise<void> {
+    const tokens = await this.getAllCredentials();
+    tokens.set(credentials.serverName, credentials);
 
     const tokenArray = Array.from(tokens.values());
     const tokenFile = this.getTokenFilePath();
@@ -107,14 +89,44 @@ export class MCPOAuthTokenStorage {
     }
   }
 
+  /**
+   * Save a token for a specific MCP server.
+   *
+   * @param serverName The name of the MCP server
+   * @param token The OAuth token to save
+   * @param clientId Optional client ID used for this token
+   * @param tokenUrl Optional token URL used for this token
+   * @param mcpServerUrl Optional MCP server URL
+   */
+  async saveToken(
+    serverName: string,
+    token: OAuthToken,
+    clientId?: string,
+    tokenUrl?: string,
+    mcpServerUrl?: string,
+  ): Promise<void> {
+    await this.ensureConfigDir();
+
+    const credential: OAuthCredentials = {
+      serverName,
+      token,
+      clientId,
+      tokenUrl,
+      mcpServerUrl,
+      updatedAt: Date.now(),
+    };
+
+    await this.setCredentials(credential);
+  }
+
   /**
    * Get a token for a specific MCP server.
    *
    * @param serverName The name of the MCP server
    * @returns The stored credentials or null if not found
    */
-  async getToken(serverName: string): Promise<OAuthCredentials | null> {
-    const tokens = await this.loadTokens();
+  async getCredentials(serverName: string): Promise<OAuthCredentials | null> {
+    const tokens = await this.getAllCredentials();
     return tokens.get(serverName) || null;
   }
 
@@ -123,8 +135,8 @@ export class MCPOAuthTokenStorage {
    *
    * @param serverName The name of the MCP server
    */
-  async removeToken(serverName: string): Promise<void> {
-    const tokens = await this.loadTokens();
+  async deleteCredentials(serverName: string): Promise<void> {
+    const tokens = await this.getAllCredentials();
 
     if (tokens.delete(serverName)) {
       const tokenArray = Array.from(tokens.values());
@@ -166,7 +178,7 @@ export class MCPOAuthTokenStorage {
   /**
    * Clear all stored MCP OAuth tokens.
    */
-  async clearAllTokens(): Promise<void> {
+  async clearAll(): Promise<void> {
     try {
       const tokenFile = this.getTokenFilePath();
       await fs.unlink(tokenFile);