MCP and OAuth

/**
 * OAuth provider for MCP server authentication
 */
class MCPOAuthProvider {
  constructor(tokenStorage?: MCPOAuthTokenStorage);
  
  /**
   * Perform OAuth authentication flow
   * @param serverName - Name of the MCP server
   * @param config - OAuth configuration
   * @param mcpServerUrl - Optional MCP server URL
   * @returns OAuth token
   */
  authenticate(
    serverName: string,
    config: MCPOAuthConfig,
    mcpServerUrl?: string
  ): Promise<OAuthToken>;
  
  /**
   * Get a valid token, refreshing if necessary
   * @param serverName - Name of the MCP server
   * @param config - OAuth configuration
   * @returns Valid access token or null if authentication required
   */
  getValidToken(serverName: string, config: MCPOAuthConfig): Promise<string | null>;
  
  /**
   * Refresh an access token using refresh token
   * @param serverName - Name of the MCP server
   * @param config - OAuth configuration
   * @param refreshToken - Refresh token
   * @returns New OAuth token
   */
  refreshAccessToken(
    serverName: string,
    config: MCPOAuthConfig,
    refreshToken: string
  ): Promise<OAuthToken>;
  
  /**
   * Revoke tokens for a server
   * @param serverName - Name of the MCP server
   * @param config - OAuth configuration
   */
  revokeTokens(serverName: string, config: MCPOAuthConfig): Promise<void>;
  
  /**
   * Check if authentication is required
   * @param serverName - Name of the MCP server
   * @returns True if authentication is required
   */
  isAuthenticationRequired(serverName: string): Promise<boolean>;
}

/**
 * OAuth configuration for MCP servers
 */
interface MCPOAuthConfig {
  /** OAuth authorization endpoint URL */
  authorizationUrl: string;
  /** OAuth token endpoint URL */
  tokenUrl: string;
  /** OAuth client identifier */
  clientId: string;
  /** OAuth client secret (optional for public clients) */
  clientSecret?: string;
  /** OAuth scopes to request */
  scopes?: string[];
  /** OAuth redirect URI */
  redirectUri?: string;
  /** Additional OAuth parameters */
  additionalParams?: Record<string, string>;
  /** PKCE support */
  usePKCE?: boolean;
}

/**
 * OAuth authorization response
 */
interface OAuthAuthorizationResponse {
  /** Authorization code */
  code: string;
  /** State parameter for CSRF protection */
  state?: string;
  /** Error code if authorization failed */
  error?: string;
  /** Error description */
  errorDescription?: string;
}

/**
 * OAuth token response
 */
interface OAuthTokenResponse {
  /** Access token */
  accessToken: string;
  /** Token type (usually 'Bearer') */
  tokenType: string;
  /** Token expiration time in seconds */
  expiresIn?: number;
  /** Refresh token */
  refreshToken?: string;
  /** Token scope */
  scope?: string;
}

/**
 * OAuth client registration request
 */
interface OAuthClientRegistrationRequest {
  /** Client name */
  clientName: string;
  /** Redirect URIs */
  redirectUris: string[];
  /** Grant types */
  grantTypes: string[];
  /** Response types */
  responseTypes: string[];
  /** Token endpoint auth method */
  tokenEndpointAuthMethod?: string;
  /** Application type */
  applicationType?: 'web' | 'native';
}

/**
 * OAuth client registration response
 */
interface OAuthClientRegistrationResponse {
  /** Client identifier */
  clientId: string;
  /** Client secret */
  clientSecret?: string;
  /** Client ID issued at timestamp */
  clientIdIssuedAt?: number;
  /** Client secret expiration */
  clientSecretExpiresAt?: number;
  /** Registration access token */
  registrationAccessToken?: string;
  /** Registration client URI */
  registrationClientUri?: string;
}

/**
 * OAuth token storage management
 */
class MCPOAuthTokenStorage {
  constructor(storagePath?: string);
  
  /**
   * Store OAuth token for a server
   * @param serverName - Name of the MCP server
   * @param token - OAuth token to store
   */
  storeToken(serverName: string, token: OAuthToken): Promise<void>;
  
  /**
   * Retrieve OAuth token for a server
   * @param serverName - Name of the MCP server
   * @returns Stored OAuth token or null if not found
   */
  getToken(serverName: string): Promise<OAuthToken | null>;
  
  /**
   * Delete OAuth token for a server
   * @param serverName - Name of the MCP server
   */
  deleteToken(serverName: string): Promise<void>;
  
  /**
   * List all stored server names
   * @returns Array of server names with stored tokens
   */
  listServers(): Promise<string[]>;
  
  /**
   * Clear all stored tokens
   */
  clearAllTokens(): Promise<void>;
  
  /**
   * Check if token exists for server
   * @param serverName - Name of the MCP server
   * @returns True if token exists
   */
  hasToken(serverName: string): Promise<boolean>;
}

/**
 * OAuth token structure
 */
interface OAuthToken {
  /** Access token */
  accessToken: string;
  /** Token type */
  tokenType: string;
  /** Token expiration timestamp (Unix timestamp) */
  expiresAt?: number;
  /** Refresh token */
  refreshToken?: string;
  /** Token scope */
  scope?: string;
  /** Additional token metadata */
  metadata?: Record<string, any>;
}

/**
 * OAuth credentials structure
 */
interface OAuthCredentials {
  /** Client identifier */
  clientId: string;
  /** Client secret */
  clientSecret?: string;
  /** Authorization server metadata */
  authorizationServerMetadata?: OAuthAuthorizationServerMetadata;
  /** Protected resource metadata */
  protectedResourceMetadata?: OAuthProtectedResourceMetadata;
}

/**
 * OAuth utility functions and helpers
 */
class OAuthUtils {
  /**
   * Generate PKCE code verifier and challenge
   * @returns PKCE parameters
   */
  static generatePKCE(): {
    codeVerifier: string;
    codeChallenge: string;
    codeChallengeMethod: string;
  };
  
  /**
   * Build authorization URL with parameters
   * @param config - OAuth configuration
   * @param state - State parameter
   * @param codeChallenge - PKCE code challenge
   * @returns Authorization URL
   */
  static buildAuthorizationUrl(
    config: MCPOAuthConfig,
    state?: string,
    codeChallenge?: string
  ): string;
  
  /**
   * Parse authorization response from URL
   * @param url - Authorization response URL
   * @returns Parsed authorization response
   */
  static parseAuthorizationResponse(url: string): OAuthAuthorizationResponse;
  
  /**
   * Exchange authorization code for tokens
   * @param config - OAuth configuration
   * @param authorizationCode - Authorization code
   * @param codeVerifier - PKCE code verifier
   * @returns OAuth token response
   */
  static exchangeCodeForTokens(
    config: MCPOAuthConfig,
    authorizationCode: string,
    codeVerifier?: string
  ): Promise<OAuthTokenResponse>;
  
  /**
   * Validate token expiration
   * @param token - OAuth token to validate
   * @param bufferSeconds - Buffer time before expiration
   * @returns True if token is valid
   */
  static isTokenValid(token: OAuthToken, bufferSeconds?: number): boolean;
  
  /**
   * Discover OAuth server metadata
   * @param issuerUrl - OAuth issuer URL
   * @returns Authorization server metadata
   */
  static discoverServerMetadata(issuerUrl: string): Promise<OAuthAuthorizationServerMetadata>;
}

/**
 * OAuth authorization server metadata
 */
interface OAuthAuthorizationServerMetadata {
  /** Issuer identifier */
  issuer: string;
  /** Authorization endpoint */
  authorizationEndpoint: string;
  /** Token endpoint */
  tokenEndpoint: string;
  /** Supported response types */
  responseTypesSupported: string[];
  /** Supported subject types */
  subjectTypesSupported: string[];
  /** Supported ID token signing algorithms */
  idTokenSigningAlgValuesSupported: string[];
  /** Revocation endpoint */
  revocationEndpoint?: string;
  /** Introspection endpoint */
  introspectionEndpoint?: string;
  /** JWKS URI */
  jwksUri?: string;
  /** Supported scopes */
  scopesSupported?: string[];
  /** Supported grant types */
  grantTypesSupported?: string[];
  /** Supported token endpoint auth methods */
  tokenEndpointAuthMethodsSupported?: string[];
  /** PKCE support */
  codeChallengeMethodsSupported?: string[];
}

/**
 * OAuth protected resource metadata
 */
interface OAuthProtectedResourceMetadata {
  /** Resource identifier */
  resource: string;
  /** Authorization servers */
  authorizationServers?: string[];
  /** Bearer token usage */
  bearerTokenUsage?: string[];
  /** Resource documentation */
  resourceDocumentation?: string;
}

/**
 * MCP-specific prompt logic and templates
 */
namespace MCPPrompts {
  /**
   * Generate authentication prompt for MCP server
   * @param serverName - Name of the MCP server
   * @param config - OAuth configuration
   * @returns Authentication prompt
   */
  function getAuthenticationPrompt(serverName: string, config: MCPOAuthConfig): string;
  
  /**
   * Generate server connection prompt
   * @param serverName - Name of the MCP server
   * @returns Connection prompt
   */
  function getConnectionPrompt(serverName: string): string;
  
  /**
   * Generate tool invocation prompt
   * @param serverName - Name of the MCP server
   * @param toolName - Name of the tool
   * @param parameters - Tool parameters
   * @returns Tool invocation prompt
   */
  function getToolInvocationPrompt(
    serverName: string,
    toolName: string,
    parameters: any
  ): string;
}

import { 
  MCPOAuthProvider,
  MCPOAuthTokenStorage,
  OAuthUtils,
  MCPOAuthConfig,
  OAuthToken
} from '@google/gemini-cli-core';

// Setup OAuth provider with token storage
const tokenStorage = new MCPOAuthTokenStorage();
const oauthProvider = new MCPOAuthProvider(tokenStorage);

// Configure OAuth for an MCP server
const mcpConfig: MCPOAuthConfig = {
  authorizationUrl: 'https://auth.example.com/oauth/authorize',
  tokenUrl: 'https://auth.example.com/oauth/token',
  clientId: 'mcp-client-123',
  scopes: ['read', 'write'],
  usePKCE: true
};

// Check if authentication is required
const serverName = 'example-mcp-server';
const authRequired = await oauthProvider.isAuthenticationRequired(serverName);

if (authRequired) {
  console.log('Authentication required, starting OAuth flow...');
  
  // Perform authentication
  const token = await oauthProvider.authenticate(serverName, mcpConfig);
  console.log('Authentication successful!');
  console.log(`Token expires at: ${new Date(token.expiresAt! * 1000)}`);
} else {
  console.log('Using existing authentication');
}

// Get a valid token (will refresh if needed)
const validToken = await oauthProvider.getValidToken(serverName, mcpConfig);

if (validToken) {
  console.log('Got valid token:', validToken.substring(0, 10) + '...');
  
  // Use the token for MCP server requests
  // ... MCP client code here ...
} else {
  console.log('Failed to obtain valid token');
}

// OAuth utilities usage
const pkceParams = OAuthUtils.generatePKCE();
console.log('Generated PKCE verifier:', pkceParams.codeVerifier);

const authUrl = OAuthUtils.buildAuthorizationUrl(mcpConfig, 'random-state', pkceParams.codeChallenge);
console.log('Authorization URL:', authUrl);

// Token validation
const storedToken = await tokenStorage.getToken(serverName);
if (storedToken && OAuthUtils.isTokenValid(storedToken, 300)) { // 5 minute buffer
  console.log('Stored token is still valid');
} else {
  console.log('Stored token expired or invalid');
}

// Server metadata discovery
try {
  const metadata = await OAuthUtils.discoverServerMetadata('https://auth.example.com');
  console.log('Discovered server metadata:', metadata.issuer);
  console.log('Supported scopes:', metadata.scopesSupported);
} catch (error) {
  console.error('Failed to discover server metadata:', error);
}

// Token storage management
const allServers = await tokenStorage.listServers();
console.log('Servers with stored tokens:', allServers);

// Clear expired or unwanted tokens
for (const server of allServers) {
  const token = await tokenStorage.getToken(server);
  if (token && !OAuthUtils.isTokenValid(token)) {
    await tokenStorage.deleteToken(server);
    console.log(`Deleted expired token for ${server}`);
  }
}