From d0372fd8d8c4985e8df7dfc879c81a1f2b8a5198 Mon Sep 17 00:00:00 2001 From: Glen Maddern Date: Tue, 6 May 2025 14:59:38 +1000 Subject: [PATCH] GEMINI TO THE RESCUE --- src/auth/browser-provider.ts | 354 +++++++--------- src/auth/callback.ts | 211 ++++------ src/auth/types.ts | 19 +- src/react/types.ts | 2 +- src/react/useMcp.ts | 774 +++++++++++++++-------------------- 5 files changed, 588 insertions(+), 772 deletions(-) diff --git a/src/auth/browser-provider.ts b/src/auth/browser-provider.ts index 6ab96c2..ae52cb3 100644 --- a/src/auth/browser-provider.ts +++ b/src/auth/browser-provider.ts @@ -1,25 +1,24 @@ -import { OAuthClientInformation, OAuthMetadata, OAuthTokens } from '@modelcontextprotocol/sdk/shared/auth.js'; +// browser-provider.ts +import { OAuthClientInformation, OAuthMetadata, OAuthTokens, OAuthClientMetadata } from '@modelcontextprotocol/sdk/shared/auth.js'; import { OAuthClientProvider } from '@modelcontextprotocol/sdk/client/auth.js'; -import { StoredState } from './types.js'; +// Assuming StoredState is defined in ./types.js and includes fields for provider options +import { StoredState } from './types.js'; // Adjust path if necessary + + /** * Browser-compatible OAuth client provider for MCP using localStorage. - * Handles storing client information, tokens, and managing the popup/redirect flow. */ export class BrowserOAuthClientProvider implements OAuthClientProvider { - private storageKeyPrefix: string; - private serverUrlHash: string; - private clientName: string; - private clientUri: string; - private callbackUrl: string; + readonly serverUrl: string; + readonly storageKeyPrefix: string; + readonly serverUrlHash: string; + readonly clientName: string; + readonly clientUri: string; + readonly callbackUrl: string; - /** - * Creates an instance of BrowserOAuthClientProvider. - * @param serverUrl The base URL of the MCP server requiring authentication. - * @param options Configuration options for the provider. - */ constructor( - readonly serverUrl: string, + serverUrl: string, options: { storageKeyPrefix?: string; clientName?: string; @@ -27,26 +26,21 @@ export class BrowserOAuthClientProvider implements OAuthClientProvider { callbackUrl?: string; } = {}, ) { + this.serverUrl = serverUrl; this.storageKeyPrefix = options.storageKeyPrefix || 'mcp:auth'; - // Hash the server URL to create unique storage keys per server this.serverUrlHash = this.hashString(serverUrl); this.clientName = options.clientName || 'MCP Browser Client'; - this.clientUri = options.clientUri || window.location.origin; - // Default callback URL is /oauth/callback on the current origin - this.callbackUrl = options.callbackUrl || new URL('/oauth/callback', window.location.origin).toString(); + this.clientUri = options.clientUri || (typeof window !== 'undefined' ? window.location.origin : ''); + this.callbackUrl = options.callbackUrl || (typeof window !== 'undefined' ? new URL('/oauth/callback', window.location.origin).toString() : '/oauth/callback'); } - /** - * The redirect URL used for the OAuth flow. - */ + // --- SDK Interface Methods --- + get redirectUrl(): string { return this.callbackUrl; } - /** - * Metadata about this client, sent during dynamic client registration (if supported by the server). - */ - get clientMetadata() { + get clientMetadata(): OAuthClientMetadata { return { redirect_uris: [this.redirectUrl], token_endpoint_auth_method: 'none', // Public client @@ -58,222 +52,180 @@ export class BrowserOAuthClientProvider implements OAuthClientProvider { }; } - /** - * Clears all localStorage items associated with this specific server URL. - * @returns The number of items removed from storage. - */ - clearStorage(): number { - const prefix = `${this.storageKeyPrefix}_${this.serverUrlHash}`; - const keysToRemove: string[] = []; - - // Find keys directly related to this server instance - for (let i = 0; i < localStorage.length; i++) { - const key = localStorage.key(i); - if (key && key.startsWith(prefix)) { - keysToRemove.push(key); - } - } - - // Also check any persisted OAuth state keys - for (let i = 0; i < localStorage.length; i++) { - const key = localStorage.key(i); - if (key && key.startsWith(`${this.storageKeyPrefix}:state_`)) { - try { - const item = localStorage.getItem(key); - if (item) { - const state = JSON.parse(item) as Partial; - // If state belongs to this server, mark for removal - if (state.serverUrlHash === this.serverUrlHash) { - keysToRemove.push(key); - } - } - } catch (e) { - console.warn(`[${this.storageKeyPrefix}] Error parsing state key ${key}:`, e); - // Optionally remove malformed keys - // keysToRemove.push(key); - } - } - } - - // Remove all identified keys - const uniqueKeysToRemove = [...new Set(keysToRemove)]; // Ensure uniqueness - uniqueKeysToRemove.forEach(key => localStorage.removeItem(key)); - - return uniqueKeysToRemove.length; - } - - private hashString(str: string): string { - // Simple, non-cryptographic hash function suitable for creating unique keys - let hash = 0; - for (let i = 0; i < str.length; i++) { - const char = str.charCodeAt(i); - hash = (hash << 5) - hash + char; - hash = hash & hash; // Convert to 32bit integer - } - // Use absolute value and convert to hex for a cleaner key - return Math.abs(hash).toString(16); - } - - private getKey(key: string): string { - return `${this.storageKeyPrefix}_${this.serverUrlHash}_${key}`; - } - - /** - * Retrieves client information (like client_id) from storage. - * @returns The stored client information, or undefined if not found/invalid. - */ async clientInformation(): Promise { const key = this.getKey('client_info'); const data = localStorage.getItem(key); if (!data) return undefined; - try { - // TODO: Add validation using a schema (e.g., Zod) + // TODO: Add validation using a schema return JSON.parse(data) as OAuthClientInformation; } catch (e) { - console.warn(`[${this.storageKeyPrefix}] Failed to parse client information from storage:`, e); - localStorage.removeItem(key); // Clean up invalid data + console.warn(`[${this.storageKeyPrefix}] Failed to parse client information:`, e); + localStorage.removeItem(key); return undefined; } } - /** - * Saves client information to storage. - * @param clientInformation The client information to save. - */ - async saveClientInformation(clientInformation: OAuthClientInformation): Promise { + // NOTE: The SDK's auth() function uses this if dynamic registration is needed. + // Ensure your OAuthClientInformationFull matches the expected structure if DCR is used. + async saveClientInformation(clientInformation: OAuthClientInformation /* | OAuthClientInformationFull */): Promise { const key = this.getKey('client_info'); + // Cast needed if handling OAuthClientInformationFull specifically localStorage.setItem(key, JSON.stringify(clientInformation)); } - /** - * Retrieves OAuth tokens (access token, refresh token) from storage. - * @returns The stored tokens, or undefined if not found/invalid. - */ + async tokens(): Promise { const key = this.getKey('tokens'); const data = localStorage.getItem(key); if (!data) return undefined; - try { - // TODO: Add validation using a schema (e.g., Zod) - const tokens = JSON.parse(data) as OAuthTokens; - // Optional: Check token expiry here if 'expires_at' is stored - return tokens; + // TODO: Add validation + return JSON.parse(data) as OAuthTokens; } catch (e) { - console.warn(`[${this.storageKeyPrefix}] Failed to parse tokens from storage:`, e); - localStorage.removeItem(key); // Clean up invalid data + console.warn(`[${this.storageKeyPrefix}] Failed to parse tokens:`, e); + localStorage.removeItem(key); return undefined; } } - /** - * Saves OAuth tokens to storage. - * @param tokens The tokens to save. - */ async saveTokens(tokens: OAuthTokens): Promise { const key = this.getKey('tokens'); - // Optional: Calculate and store expiry timestamp: - // const storedTokens = { ...tokens, expires_at: Date.now() + (tokens.expires_in || 0) * 1000 }; localStorage.setItem(key, JSON.stringify(tokens)); + // Clean up code verifier and last auth URL after successful token save + localStorage.removeItem(this.getKey('code_verifier')); + localStorage.removeItem(this.getKey('last_auth_url')); } - /** - * Initiates the OAuth authorization flow by opening a popup window. - * Stores necessary state in localStorage for the callback handler. - * @param authorizationUrl The fully constructed authorization URL (including client_id, redirect_uri, scope, code_challenge, etc.). - * @param metadata The discovered OAuth metadata of the server. - * @param options Configuration for the popup window. - * @returns An object indicating success, whether the popup was potentially blocked, and the URL used. - */ - // @ts-ignore - async redirectToAuthorization( - authorizationUrl: URL, - metadata: OAuthMetadata, - options?: { - popupFeatures?: string; - }, - ): Promise<{ success: boolean; popupBlocked?: boolean; url: string }> { - // Generate a unique state parameter for this authorization request - const state = crypto.randomUUID(); // Use crypto.randomUUID for better randomness - const stateKey = `${this.storageKeyPrefix}:state_${state}`; - - // Store context needed by the callback handler, associated with the state param - localStorage.setItem( - stateKey, - JSON.stringify({ - authorizationUrl: authorizationUrl.origin, // Store origin for exchangeAuthorization - metadata, - serverUrlHash: this.serverUrlHash, // Link state back to this server instance - expiry: Date.now() + 1000 * 60 * 10, // State expires in 10 minutes - } as StoredState), - ); - authorizationUrl.searchParams.set('state', state); - - const authUrlString = authorizationUrl.toString(); - const popupFeatures = options?.popupFeatures || 'width=600,height=700,resizable=yes,scrollbars=yes,status=yes'; - - // Persist the exact auth URL in case the popup fails and manual navigation is needed - localStorage.setItem(this.getKey('auth_url'), authUrlString); - - try { - // Attempt to open the authorization URL in a new window - const popup = window.open(authUrlString, `mcp_auth_${this.serverUrlHash}`, popupFeatures); - - // Check if the popup window handle was obtained and if it's not immediately closed - if (!popup || popup.closed || typeof popup.closed === 'undefined') { - // This is the most common indicator of a popup blocker - console.warn(`[${this.storageKeyPrefix}] Popup likely blocked by browser.`); - return { success: false, popupBlocked: true, url: authUrlString }; - } - - // Additional check: try accessing a property that might throw if blocked cross-origin - // Note: This might not always work due to browser security policies evolution - try { - // Accessing location.href will likely throw a cross-origin error immediately - // if the popup is loading the auth server's domain, which is expected. - // What we *really* want to detect is if the popup *failed to open at all*. - // The initial check (!popup || popup.closed) is usually sufficient. - // Focusing the popup can be a sign it opened successfully. - popup.focus(); - } catch (e) { - // If accessing the popup threw an error immediately, it might indicate a deeper issue, - // but often it's just a standard cross-origin security restriction. - // The initial check is more reliable for detecting simple blocking. - console.debug(`[${this.storageKeyPrefix}] Accessing popup properties caused potential cross-origin error (normal):`, e); - } - - // If we reached here, the popup handle was obtained and wasn't immediately closed. - return { success: true, url: authUrlString }; - } catch (e) { - // Catch any unexpected errors during window.open - console.error(`[${this.storageKeyPrefix}] Error opening popup window:`, e); - return { success: false, popupBlocked: true, url: authUrlString }; // Assume blocked on error - } - } - - /** - * Saves the PKCE code verifier to storage. - * @param codeVerifier The code verifier string. - */ async saveCodeVerifier(codeVerifier: string): Promise { const key = this.getKey('code_verifier'); localStorage.setItem(key, codeVerifier); } - /** - * Retrieves the PKCE code verifier from storage. - * @returns The stored code verifier. - * @throws If the code verifier is not found in storage. - */ async codeVerifier(): Promise { const key = this.getKey('code_verifier'); const verifier = localStorage.getItem(key); if (!verifier) { throw new Error(`[${this.storageKeyPrefix}] Code verifier not found in storage for key ${key}. Auth flow likely corrupted or timed out.`); } - // Optionally remove the verifier after retrieving it, as it's single-use - // localStorage.removeItem(key); + // SDK's auth() retrieves this BEFORE exchanging code. Don't remove it here. + // It will be removed in saveTokens on success. return verifier; } -} + + /** + * Redirects the user agent to the authorization URL, storing necessary state. + * This now adheres to the SDK's void return type expectation for the interface. + * @param authorizationUrl The fully constructed authorization URL from the SDK. + */ + async redirectToAuthorization(authorizationUrl: URL): Promise { + // Generate a unique state parameter for this authorization request + const state = crypto.randomUUID(); + const stateKey = `${this.storageKeyPrefix}:state_${state}`; + + // Store context needed by the callback handler, associated with the state param + const stateData: StoredState = { + serverUrlHash: this.serverUrlHash, + expiry: Date.now() + 1000 * 60 * 10, // State expires in 10 minutes + // Store provider options needed to reconstruct on callback + providerOptions: { + serverUrl: this.serverUrl, + storageKeyPrefix: this.storageKeyPrefix, + clientName: this.clientName, + clientUri: this.clientUri, + callbackUrl: this.callbackUrl, + } + }; + localStorage.setItem(stateKey, JSON.stringify(stateData)); + + // Add the state parameter to the URL + authorizationUrl.searchParams.set('state', state); + const authUrlString = authorizationUrl.toString(); + + // Persist the exact auth URL in case the popup fails and manual navigation is needed + localStorage.setItem(this.getKey('last_auth_url'), authUrlString); + + // Attempt to open the popup + const popupFeatures = 'width=600,height=700,resizable=yes,scrollbars=yes,status=yes'; // Make configurable if needed + try { + const popup = window.open(authUrlString, `mcp_auth_${this.serverUrlHash}`, popupFeatures); + + if (!popup || popup.closed || typeof popup.closed === 'undefined') { + console.warn(`[${this.storageKeyPrefix}] Popup likely blocked by browser. Manual navigation might be required using the stored URL.`); + // Cannot signal failure back via SDK auth() directly. + // useMcp will need to rely on timeout or manual trigger if stuck. + } else { + popup.focus(); + console.info(`[${this.storageKeyPrefix}] Redirecting to authorization URL in popup.`); + } + } catch (e) { + console.error(`[${this.storageKeyPrefix}] Error opening popup window:`, e); + // Cannot signal failure back via SDK auth() directly. + } + // Regardless of popup success, the interface expects this method to initiate the redirect. + // If the popup failed, the user journey stops here until manual action or timeout. + } + + // --- Helper Methods --- + + /** + * Retrieves the last URL passed to `redirectToAuthorization`. Useful for manual fallback. + */ + getLastAttemptedAuthUrl(): string | null { + return localStorage.getItem(this.getKey('last_auth_url')); + } + + + clearStorage(): number { + const prefixPattern = `${this.storageKeyPrefix}_${this.serverUrlHash}_`; + const statePattern = `${this.storageKeyPrefix}:state_`; + const keysToRemove: string[] = []; + let count = 0; + + for (let i = 0; i < localStorage.length; i++) { + const key = localStorage.key(i); + if (!key) continue; + + if (key.startsWith(prefixPattern)) { + keysToRemove.push(key); + } else if (key.startsWith(statePattern)) { + try { + const item = localStorage.getItem(key); + if (item) { + // Check if state belongs to this provider instance based on serverUrlHash + // We need to parse cautiously as the structure isn't guaranteed. + const state = JSON.parse(item) as Partial; + if (state.serverUrlHash === this.serverUrlHash) { + keysToRemove.push(key); + } + } + } catch (e) { + console.warn(`[${this.storageKeyPrefix}] Error parsing state key ${key} during clearStorage:`, e); + // Optionally remove malformed keys + // keysToRemove.push(key); + } + } + } + + const uniqueKeysToRemove = [...new Set(keysToRemove)]; + uniqueKeysToRemove.forEach(key => { + localStorage.removeItem(key); + count++; + }); + return count; + } + + private hashString(str: string): string { + let hash = 0; + for (let i = 0; i < str.length; i++) { + const char = str.charCodeAt(i); + hash = (hash << 5) - hash + char; + hash = hash & hash; + } + return Math.abs(hash).toString(16); + } + + getKey(keySuffix: string): string { + return `${this.storageKeyPrefix}_${this.serverUrlHash}_${keySuffix}`; + } +} \ No newline at end of file diff --git a/src/auth/callback.ts b/src/auth/callback.ts index 54fe737..9e9fcae 100644 --- a/src/auth/callback.ts +++ b/src/auth/callback.ts @@ -1,180 +1,137 @@ -import { exchangeAuthorization } from '@modelcontextprotocol/sdk/client/auth.js'; -import { OAuthClientInformation, OAuthMetadata } from '@modelcontextprotocol/sdk/shared/auth.js'; -import { StoredState } from './types.js'; +// callback.ts +import { auth } from '@modelcontextprotocol/sdk/client/auth.js' +import { BrowserOAuthClientProvider } from './browser-provider.js' // Adjust path +import { StoredState } from './types.js' // Adjust path, ensure definition includes providerOptions /** - * Handles the OAuth callback. This function should be invoked on your callback page - * (e.g., /oauth/callback) with the URL query parameters. - * - * It validates the state parameter, retrieves stored context from localStorage, - * exchanges the authorization code for tokens using the stored code verifier, - * saves the tokens, and notifies the opener window. - * - * @param query The URL query parameters (e.g., from `window.location.search`). - * @param options Configuration options. - * @param options.storageKeyPrefix The prefix used for localStorage keys (must match the one used by `BrowserOAuthClientProvider`). Defaults to "mcp:auth". - * @returns An object indicating success or failure with an error message. + * Handles the OAuth callback using the SDK's auth() function. + * Assumes it's running on the page specified as the callbackUrl. */ -export async function onMcpAuthorization( - query: Record, - { - storageKeyPrefix = 'mcp:auth', - }: { - storageKeyPrefix?: string; - } = {}, -): Promise<{ success: boolean; error?: string }> { - const logPrefix = `[${storageKeyPrefix}-callback]`; - console.log(`${logPrefix} Handling OAuth callback with query:`, query); +export async function onMcpAuthorization() { + const queryParams = new URLSearchParams(window.location.search); + const code = queryParams.get('code'); + const state = queryParams.get('state'); + const error = queryParams.get('error'); + const errorDescription = queryParams.get('error_description'); + + const logPrefix = '[mcp-callback]'; // Generic prefix, or derive from stored state later + console.log(`${logPrefix} Handling callback...`, { code, state, error, errorDescription }); + + let provider: BrowserOAuthClientProvider | null = null; + let storedStateData: StoredState | null = null; + const stateKey = state ? `mcp:auth:state_${state}` : null; // Reconstruct state key prefix assumption try { - const code = query.code; - const state = query.state; - const error = query.error; - const errorDescription = query.error_description; - + // --- Basic Error Handling --- if (error) { throw new Error(`OAuth error: ${error} - ${errorDescription || 'No description provided.'}`); } - if (!code) { throw new Error('Authorization code not found in callback query parameters.'); } - if (!state) { - throw new Error('State parameter not found in callback query parameters.'); + if (!state || !stateKey) { + throw new Error('State parameter not found or invalid in callback query parameters.'); } - // Retrieve the stored state using the state parameter from the query - const stateKey = `${storageKeyPrefix}:state_${state}`; + // --- Retrieve Stored State & Provider Options --- const storedStateJSON = localStorage.getItem(stateKey); - if (!storedStateJSON) { throw new Error(`Invalid or expired state parameter "${state}". No matching state found in storage.`); } - - let storedState: StoredState; try { - storedState = JSON.parse(storedStateJSON); + storedStateData = JSON.parse(storedStateJSON) as StoredState; } catch (e) { throw new Error('Failed to parse stored OAuth state.'); } // Validate expiry - if (storedState.expiry < Date.now()) { + if (!storedStateData.expiry || storedStateData.expiry < Date.now()) { + localStorage.removeItem(stateKey); // Clean up expired state throw new Error('OAuth state has expired. Please try initiating authentication again.'); } - const { authorizationUrl, serverUrlHash, metadata } = storedState; - console.log(`${logPrefix} Found valid state for server hash: ${serverUrlHash}`); - - // Construct keys to retrieve client info and code verifier for this server instance - const clientInfoKey = `${storageKeyPrefix}_${serverUrlHash}_client_info`; - const codeVerifierKey = `${storageKeyPrefix}_${serverUrlHash}_code_verifier`; - const tokensKey = `${storageKeyPrefix}_${serverUrlHash}_tokens`; - const authUrlKey = `${storageKeyPrefix}_${serverUrlHash}_auth_url`; // Key for the potentially stored manual auth URL - - const clientInfoStr = localStorage.getItem(clientInfoKey); - const codeVerifier = localStorage.getItem(codeVerifierKey); - - if (!clientInfoStr) { - throw new Error(`Client information not found in storage (key: ${clientInfoKey}).`); - } - if (!codeVerifier) { - throw new Error(`Code verifier not found in storage (key: ${codeVerifierKey}). Auth flow may be incomplete or timed out.`); + // Ensure provider options are present + if (!storedStateData.providerOptions) { + throw new Error('Stored state is missing required provider options.'); } + const { serverUrl, ...providerOptions } = storedStateData.providerOptions; - let clientInfo: OAuthClientInformation; - try { - clientInfo = JSON.parse(clientInfoStr); - } catch (e) { - throw new Error('Failed to parse stored client information.'); - } - console.log(`${logPrefix} Exchanging authorization code for token...`); + // --- Instantiate Provider --- + console.log(`${logPrefix} Re-instantiating provider for server: ${serverUrl}`); + provider = new BrowserOAuthClientProvider(serverUrl, providerOptions); - // Exchange the code for tokens - const tokens = await exchangeAuthorization(new URL(metadata.token_endpoint), { - // Use token_endpoint from metadata - metadata, - clientInformation: clientInfo, - authorizationCode: code, - codeVerifier, - // redirectUrl is typically required by the token endpoint for verification - // @ts-ignore - redirectUrl: new URL('/oauth/callback', window.location.origin).toString(), // TODO: Make this configurable or get from provider instance? - }); + // --- Call SDK Auth Function --- + console.log(`${logPrefix} Calling SDK auth() to exchange code...`); + // The SDK auth() function will internally: + // 1. Use provider.clientInformation() + // 2. Use provider.codeVerifier() + // 3. Call exchangeAuthorization() + // 4. Use provider.saveTokens() on success + const authResult = await auth(provider, { serverUrl, authorizationCode: code }); - console.log(`${logPrefix} Token exchange successful. Saving tokens...`); + if (authResult === 'AUTHORIZED') { + console.log(`${logPrefix} Authorization successful via SDK auth(). Notifying opener...`); + // --- Notify Opener and Close (Success) --- + if (window.opener && !window.opener.closed) { + window.opener.postMessage( + { type: 'mcp_auth_callback', success: true }, + window.location.origin, + ); + window.close(); + } else { + console.warn(`${logPrefix} No opener window detected. Redirecting to root.`); + window.location.href = '/'; // Or a configured post-auth destination + } + // Clean up state ONLY on success and after notifying opener + localStorage.removeItem(stateKey); - // Save the obtained tokens - localStorage.setItem(tokensKey, JSON.stringify(tokens)); - - // Clean up the persisted manual auth URL if it exists - localStorage.removeItem(authUrlKey); - // Remove the state key immediately after retrieval (it's single-use) - localStorage.removeItem(stateKey); - // Remove the code verifier after retrieving it - localStorage.removeItem(codeVerifierKey); - - console.log(`${logPrefix} Tokens saved. Notifying opener window.`); - - // Notify the original window (opener) that authentication succeeded - if (window.opener && !window.opener.closed) { - window.opener.postMessage( - { - type: 'mcp_auth_callback', - success: true, - // Optionally pass serverUrlHash if multiple hooks might be listening - serverUrlHash: serverUrlHash, - }, - window.location.origin, // Target origin must match the opener's origin - ); - // Close the popup window - window.close(); } else { - // If there's no opener, this might be a full page redirect flow. - // Redirect back to the app's main page or a specific post-auth page. - console.warn(`${logPrefix} No opener window detected. Assuming redirect flow. Redirecting to root.`); - window.location.href = '/'; // Adjust as needed + // This case shouldn't happen if `authorizationCode` is provided to `auth()` + console.warn(`${logPrefix} SDK auth() returned unexpected status: ${authResult}`); + throw new Error(`Unexpected result from authentication library: ${authResult}`); } - return { success: true }; } catch (err) { console.error(`${logPrefix} Error during OAuth callback handling:`, err); const errorMessage = err instanceof Error ? err.message : String(err); - // Attempt to notify the opener window about the failure + // --- Notify Opener and Display Error (Failure) --- if (window.opener && !window.opener.closed) { window.opener.postMessage( - { - type: 'mcp_auth_callback', - success: false, - error: errorMessage, - // Optionally pass serverUrlHash - // serverUrlHash: storedState?.serverUrlHash // Might not be available if state parsing failed - }, + { type: 'mcp_auth_callback', success: false, error: errorMessage }, window.location.origin, ); + // Optionally close even on error, depending on UX preference + // window.close(); } - // Display error in the callback window itself for better debugging + // Display error in the callback window try { document.body.innerHTML = ` -
-

Authentication Error

-

- ${errorMessage} -

-

You can close this window.

- 
${
+            

+            
Authentication Error

+            

+                ${errorMessage}
+            

+            
You can close this window or click here to close.

+            
${
         err instanceof Error ? err.stack : ''
       }

-        

-      `;
+
+ `; } catch (displayError) { - // Fallback if document.body is not available or writable console.error(`${logPrefix} Could not display error in callback window:`, displayError); } - - return { success: false, error: errorMessage }; + // Clean up potentially invalid state on error + if (stateKey) { + localStorage.removeItem(stateKey); + } + // Clean up potentially dangling verifier or last_auth_url if auth failed badly + // Note: saveTokens should clean these on success + if (provider) { + localStorage.removeItem(provider.getKey('code_verifier')); + localStorage.removeItem(provider.getKey('last_auth_url')); + } } -} \ No newline at end of file +} diff --git a/src/auth/types.ts b/src/auth/types.ts index 6892739..95159be 100644 --- a/src/auth/types.ts +++ b/src/auth/types.ts @@ -4,9 +4,16 @@ import { OAuthMetadata } from '@modelcontextprotocol/sdk/shared/auth.js'; * Internal type for storing OAuth state in localStorage during the popup flow. * @internal */ -export type StoredState = { - authorizationUrl: string; - metadata: OAuthMetadata; - serverUrlHash: string; // To associate state with a specific server URL - expiry: number; // Timestamp when the state expires -}; \ No newline at end of file +export interface StoredState { + expiry: number; + metadata?: OAuthMetadata; // Optional: might not be needed if auth() rediscovers + serverUrlHash: string; + // Add provider options needed on callback: + providerOptions: { + serverUrl: string; + storageKeyPrefix: string; + clientName: string; + clientUri: string; + callbackUrl: string; + }; +} \ No newline at end of file diff --git a/src/react/types.ts b/src/react/types.ts index 3c6035d..915d339 100644 --- a/src/react/types.ts +++ b/src/react/types.ts @@ -66,7 +66,7 @@ export type UseMcpResult = { * @returns A promise that resolves with the authorization URL opened (or intended to be opened), * or undefined if auth cannot be started. */ - authenticate: () => Promise; + authenticate: () => void; /** Clears all stored authentication data (tokens, client info, etc.) for this server URL from localStorage. */ clearStorage: () => void; }; \ No newline at end of file diff --git a/src/react/useMcp.ts b/src/react/useMcp.ts index 92bd267..e1d8a35 100644 --- a/src/react/useMcp.ts +++ b/src/react/useMcp.ts @@ -1,575 +1,475 @@ -import { CallToolResultSchema, JSONRPCMessage, ListToolsResultSchema, Tool } from '@modelcontextprotocol/sdk/types.js'; -import { useCallback, useEffect, useRef, useState } from 'react'; -import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'; -import { Client } from '@modelcontextprotocol/sdk/client/index.js'; -import { discoverOAuthMetadata, exchangeAuthorization, startAuthorization, UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js'; -import { OAuthClientInformation, OAuthMetadata, OAuthTokens } from '@modelcontextprotocol/sdk/shared/auth.js'; -import { BrowserOAuthClientProvider } from '../auth/browser-provider.js'; -import { assert } from '../utils/assert.js'; -import type { UseMcpOptions, UseMcpResult } from './types.js'; +// useMcp.ts +import { CallToolResultSchema, JSONRPCMessage, ListToolsResultSchema, Tool } from '@modelcontextprotocol/sdk/types.js' +import { useCallback, useEffect, useRef, useState } from 'react' +import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js' +import { Client } from '@modelcontextprotocol/sdk/client/index.js' +// Import the main auth function and error type +import { auth, UnauthorizedError, OAuthClientProvider } from '@modelcontextprotocol/sdk/client/auth.js' +// Keep BrowserOAuthClientProvider +import { BrowserOAuthClientProvider } from '../auth/browser-provider.js' // Adjust path +import { assert } from '../utils/assert.js' // Adjust path +import type { UseMcpOptions, UseMcpResult } from './types.js' // Adjust path -const DEFAULT_RECONNECT_DELAY = 3000; -const DEFAULT_RETRY_DELAY = 5000; -const AUTH_TIMEOUT = 5 * 60 * 1000; // 5 minutes for user to complete auth in popup +const DEFAULT_RECONNECT_DELAY = 3000 +const DEFAULT_RETRY_DELAY = 5000 +const AUTH_TIMEOUT = 5 * 60 * 1000 // 5 minutes for user to complete auth in popup -/** - * React hook to connect to a remote MCP server, handle OAuth authentication, - * and provide access to server-side tools for use with AI SDKs. - * - * @param options Configuration options for connecting to the MCP server. - * @returns State and functions for interacting with the MCP connection. - */ export function useMcp(options: UseMcpOptions): UseMcpResult { const { url, - clientName, // Passed to BrowserOAuthClientProvider - clientUri, // Passed to BrowserOAuthClientProvider + clientName, + clientUri, callbackUrl = typeof window !== 'undefined' ? new URL('/oauth/callback', window.location.origin).toString() : '/oauth/callback', storageKeyPrefix = 'mcp:auth', clientConfig = {}, debug = false, autoRetry = false, autoReconnect = DEFAULT_RECONNECT_DELAY, - popupFeatures = 'width=600,height=700,resizable=yes,scrollbars=yes,status=yes', - } = options; + // popupFeatures is now handled within BrowserOAuthClientProvider + } = options - const [state, setState] = useState('discovering'); - const [tools, setTools] = useState([]); - const [error, setError] = useState(undefined); - const [log, setLog] = useState([]); - const [authUrl, setAuthUrl] = useState(undefined); // For manual auth if popup blocked + const [state, setState] = useState('discovering') + const [tools, setTools] = useState([]) + const [error, setError] = useState(undefined) + const [log, setLog] = useState([]) + // Store the URL for manual fallback if needed + const [authUrl, setAuthUrl] = useState(undefined) - const clientRef = useRef(null); - const transportRef = useRef(null); - const authProviderRef = useRef(null); - const metadataRef = useRef(undefined); - const pendingAuthUrlRef = useRef(undefined); // Stores the URL while waiting for popup - const connectingRef = useRef(false); // Prevent concurrent connection attempts - const isMountedRef = useRef(true); // Track component mount status - const connectAttemptRef = useRef(0); // Track connection attempts for retry logic - const authTimeoutRef = useRef(null); // Timer for auth popup - - // ===== Logging Utility ===== + const clientRef = useRef(null) + const transportRef = useRef(null) + const authProviderRef = useRef(null) + const connectingRef = useRef(false) + const isMountedRef = useRef(true) + const connectAttemptRef = useRef(0) + const authTimeoutRef = useRef(null) // Use number for type safety const addLog = useCallback( (level: UseMcpResult['log'][0]['level'], message: string, ...args: unknown[]) => { - // if (level === 'debug' && !debug) return; - - const fullMessage = args.length > 0 ? `${message} ${args.map(arg => JSON.stringify(arg)).join(' ')}` : message; - console[level](`[useMcp] ${fullMessage}`); + const fullMessage = args.length > 0 ? `${message} ${args.map((arg) => JSON.stringify(arg)).join(' ')}` : message + console[level](`[useMcp] ${fullMessage}`) if (isMountedRef.current) { - setLog((prevLog) => [...prevLog, { level, message: fullMessage, timestamp: Date.now() }]); + setLog((prevLog) => [...prevLog.slice(-100), { level, message: fullMessage, timestamp: Date.now() }]) // Limit log size } }, - [debug], // Remove isMountedRef as it's a ref and shouldn't be a dependency - ); - - // ===== Core Actions ===== + [], // No dependencies needed if `debug` isn't used inside + ) const disconnect = useCallback( async (quiet = false) => { - if (!quiet) addLog('info', 'Disconnecting...'); - connectingRef.current = false; - clearTimeout(authTimeoutRef.current!); - authTimeoutRef.current = null; + if (!quiet) addLog('info', 'Disconnecting...') + connectingRef.current = false + if (authTimeoutRef.current) clearTimeout(authTimeoutRef.current) + authTimeoutRef.current = null - const client = clientRef.current; - const transport = transportRef.current; - clientRef.current = null; - transportRef.current = null; + const client = clientRef.current + const transport = transportRef.current + clientRef.current = null + transportRef.current = null // Prevent reconnection logic on explicit disconnect - // Don't reset state if unmounting if (isMountedRef.current && !quiet) { - setState('discovering'); // Or 'disconnected'? Let's stick to initial flow for retry. - setTools([]); - setError(undefined); - setAuthUrl(undefined); - pendingAuthUrlRef.current = undefined; + setState('discovering') + setTools([]) + setError(undefined) + setAuthUrl(undefined) } - // Close transport first, then client if (transport) { try { - await transport.close(); - if (!quiet) addLog('debug', 'Transport closed'); + await transport.close() + if (!quiet) addLog('debug', 'Transport closed') } catch (err) { - if (!quiet) addLog('warn', 'Error closing transport:', err); + if (!quiet) addLog('warn', 'Error closing transport:', err) } } - // Note: Client doesn't have an explicit async close, it relies on transport. }, [addLog], - ); + ) - const failConnection = useCallback((errorMessage: string, connectionError?: Error) => { - addLog('error', errorMessage, connectionError ?? ''); - if (isMountedRef.current) { - setState('failed'); - setError(errorMessage); - } - connectingRef.current = false; - disconnect(true); // Quiet disconnect - }, [addLog, disconnect]); - - - // ===== Authentication Flow ===== - - const initiateAuthFlow = useCallback(async (): Promise => { - assert(authProviderRef.current, 'Auth provider not initialized'); - assert(metadataRef.current, 'OAuth metadata not discovered'); - addLog('info', 'Initiating OAuth authorization flow...'); - - try { - // Dynamic Client Registration (DCR) - Placeholder - // In a real browser scenario, DCR is complex. Often, a pre-registered - // client_id is used. If DCR is needed, it typically involves a backend proxy. - let clientInfo = await authProviderRef.current.clientInformation(); - if (!clientInfo) { - // Assuming pre-registration or manual setup for browser clients - // If DCR were implemented: - // addLog('info', 'No client info found, attempting dynamic registration...'); - // clientInfo = await registerClient(url, authProviderRef.current.clientMetadata); - // await authProviderRef.current.saveClientInformation(clientInfo); - // addLog('info', 'Client registration successful'); - - // For now, error if no client_id is found (must be pre-configured or manually stored) - addLog('warn', 'Client information not found in storage. Authentication requires a pre-registered client_id.'); - throw new Error('Client information (client_id) not found. Please ensure the client is registered with the auth server and the details are stored, or implement dynamic registration.'); - } - - addLog('debug', 'Using client information:', clientInfo); - - const { authorizationUrl, codeVerifier } = await startAuthorization(url, { - metadata: metadataRef.current, - clientInformation: clientInfo, - redirectUrl: authProviderRef.current.redirectUrl, - // scopes: ['openid', 'profile', 'mcp'], // Optional: Request specific scopes - }); - - addLog('debug', 'Authorization URL created:', authorizationUrl.toString()); - await authProviderRef.current.saveCodeVerifier(codeVerifier); - addLog('debug', 'Code verifier saved.'); - - pendingAuthUrlRef.current = authorizationUrl; // Store for popup opening - setAuthUrl(authorizationUrl.toString()); // Expose for manual fallback - - return authorizationUrl; - - } catch (err) { - failConnection(`Failed to start OAuth authorization: ${err instanceof Error ? err.message : String(err)}`, err instanceof Error ? err : undefined); - return undefined; - } - - }, [url, addLog, failConnection]); - - const openAuthPopup = useCallback(async () => { - assert(authProviderRef.current, 'Auth provider not available'); - assert(metadataRef.current, 'Metadata not available'); - - let authPopupUrl = pendingAuthUrlRef.current; - - // If URL isn't already generated, try generating it now - if (!authPopupUrl) { - addLog('info', 'Auth URL not pre-generated, initiating flow now...'); - authPopupUrl = await initiateAuthFlow(); - if (!authPopupUrl) { - // initiateAuthFlow calls failConnection on error - return undefined; - } - } - - addLog('info', 'Attempting to open authentication popup...'); - setState('authenticating'); // Update state *before* opening popup - - const redirectResult = await authProviderRef.current.redirectToAuthorization( - authPopupUrl, - metadataRef.current, - { popupFeatures }, - ); - - if (redirectResult.success) { - addLog('info', 'Authentication popup opened. Waiting for user action...'); - // Start timeout for auth completion - clearTimeout(authTimeoutRef.current!); - authTimeoutRef.current = setTimeout(() => { - addLog('warn', 'Authentication timed out.'); - if (isMountedRef.current && state === 'authenticating') { - failConnection('Authentication timed out. Please try again.'); + const failConnection = useCallback( + (errorMessage: string, connectionError?: Error) => { + addLog('error', errorMessage, connectionError ?? '') + if (isMountedRef.current) { + setState('failed') + setError(errorMessage) + // Attempt to get manual auth URL if provider exists + const manualUrl = authProviderRef.current?.getLastAttemptedAuthUrl() + if (manualUrl) { + setAuthUrl(manualUrl) + addLog('info', 'Manual authentication URL may be available.', manualUrl) } - }, AUTH_TIMEOUT); - } else { - addLog('warn', 'Authentication popup was blocked or failed to open.'); - setAuthUrl(redirectResult.url); // Ensure the manual URL is available - failConnection( - 'Authentication popup blocked. Please allow popups for this site or use the provided link to authenticate manually.', - new Error('PopupBlocked') // Custom error? - ); - // Do not clear pendingAuthUrlRef here, retry might need it - } - return redirectResult.url; - - }, [addLog, failConnection, initiateAuthFlow, popupFeatures, state]); - - - // ===== Connection Logic ===== + } + connectingRef.current = false + // Don't disconnect here, allow user to potentially retry/authenticate + // disconnect(true); // No longer needed, state is 'failed' + }, + [addLog], + ) const connect = useCallback(async () => { if (connectingRef.current) { - addLog('debug', 'Connection attempt already in progress, skipping.'); - return; + addLog('debug', 'Connection attempt already in progress.') + return } - if (!isMountedRef.current) return; // Don't connect if unmounted + if (!isMountedRef.current) return - connectingRef.current = true; - connectAttemptRef.current += 1; - setError(undefined); - setAuthUrl(undefined); - pendingAuthUrlRef.current = undefined; - setState('discovering'); - addLog('info', `Connecting attempt #${connectAttemptRef.current}...`); + connectingRef.current = true + connectAttemptRef.current += 1 + setError(undefined) + setAuthUrl(undefined) // Clear previous manual URL + setState('discovering') + addLog('info', `Connecting attempt #${connectAttemptRef.current} to ${url}...`) - // 1. Initialize Auth Provider + // 1. Initialize Auth Provider (only once) if (!authProviderRef.current) { authProviderRef.current = new BrowserOAuthClientProvider(url, { storageKeyPrefix, clientName, clientUri, callbackUrl, - }); - addLog('debug', 'BrowserOAuthClientProvider initialized.'); + }) + addLog('debug', 'BrowserOAuthClientProvider initialized.') } - // 2. Initialize MCP Client + // 2. Initialize MCP Client (only once) if (!clientRef.current) { clientRef.current = new Client( { name: clientConfig.name || 'use-mcp-react-client', - version: clientConfig.version || '0.1.0', // TODO: Get version from package.json? + version: clientConfig.version || '0.1.0', }, - { capabilities: { /* Define client caps if any */ } }, - ); - addLog('debug', 'MCP Client initialized.'); + { capabilities: {} }, + ) + addLog('debug', 'MCP Client initialized.') } - // 3. Create SSE Transport - addLog('info', 'Creating SSE transport...'); - setState('connecting'); - const serverUrl = new URL(url); // Ensure it's a URL object - - // Clear previous transport if any + // 3. Create/Recreate SSE Transport + setState('connecting') if (transportRef.current) { - await transportRef.current.close().catch(e => addLog('warn', 'Error closing previous transport:', e)); + await transportRef.current.close().catch((e) => addLog('warn', 'Error closing previous transport:', e)) + transportRef.current = null + } + try { + transportRef.current = new SSEClientTransport(new URL(url), { + authProvider: authProviderRef.current, // Pass provider for automatic token attachment + }) + addLog('debug', 'SSEClientTransport created.') + } catch (err) { + failConnection( + `Failed to create transport: ${err instanceof Error ? err.message : String(err)}`, + err instanceof Error ? err : undefined, + ) + return } - - transportRef.current = new SSEClientTransport(serverUrl, { - // @ts-ignore - authProvider: authProviderRef.current, - // Pass debug flag to transport if it supports it - }); - addLog('debug', 'SSEClientTransport created.'); // 4. Set up Transport Handlers transportRef.current.onmessage = (message: JSONRPCMessage) => { - // @ts-expect-error - SDK type might be union, check message content - addLog('debug', `[Transport] Received message: Method=${message.method || 'N/A'} ID=${message.id || 'N/A'}`); - // Forward notifications/responses to the client instance + addLog('debug', `[Transport] Received: ${JSON.stringify(message)}`) // @ts-ignore - clientRef.current?.handleMessage(message); - }; - + clientRef.current?.handleMessage(message) + } transportRef.current.onerror = (err: Error) => { - // Check if the error indicates an auth issue *before* full connection - if (err instanceof UnauthorizedError || err.message.includes('Unauthorized') || err.message.includes('401')) { - addLog('warn', 'Transport connection requires authentication.'); - // Don't immediately fail, proceed to connect step which handles auth discovery - } else { - // Handle other transport-level errors (network issues, etc.) - failConnection(`Transport error: ${err.message}`, err); - } - }; - + // Transport errors usually mean connection is lost/failed + failConnection(`Transport error: ${err.message}`, err) + } transportRef.current.onclose = () => { - // Only handle close if we weren't explicitly disconnecting or failing - if (connectingRef.current || !isMountedRef.current) return; + if (!isMountedRef.current || connectingRef.current) return // Ignore if unmounting or mid-connect - addLog('info', 'Transport connection closed.'); - // Use the current state directly from the state variable - const currentState = state; - + addLog('info', 'Transport connection closed.') + const currentState = state // Capture state at time of closure if (currentState === 'ready' && autoReconnect) { - const delay = typeof autoReconnect === 'number' ? autoReconnect : DEFAULT_RECONNECT_DELAY; - addLog('info', `Attempting to reconnect in ${delay}ms...`); + const delay = typeof autoReconnect === 'number' ? autoReconnect : DEFAULT_RECONNECT_DELAY + addLog('info', `Attempting to reconnect in ${delay}ms...`) setTimeout(() => { - // Perform a clean disconnect before reconnecting - disconnect().then(() => { - if (isMountedRef.current) connect(); - }); - }, delay); - // Set state to indicate reconnection attempt? Maybe 'connecting'? - if (isMountedRef.current) setState('connecting'); - } else if (currentState !== 'failed') { - // If not ready or autoReconnect is off, treat unexpected close as failure - failConnection('Connection closed unexpectedly.'); - } - }; - - // 5. Connect Client (which starts the transport internally) - try { - addLog('info', 'Connecting client...'); - assert(clientRef.current, "Client not initialized"); - assert(transportRef.current, "Transport not initialized"); - - // The client.connect() method internally calls transport.start() - await clientRef.current.connect(transportRef.current); - - // If connect resolves without error, we are connected (potentially anonymously) - addLog('info', 'Client connected successfully.'); - setState('loading'); // Move to loading state - - // 6. Load Tools - try { - addLog('info', 'Loading tools...'); - const toolsResponse = await clientRef.current.request({ method: 'tools/list' }, ListToolsResultSchema); - if (isMountedRef.current) { - setTools(toolsResponse.tools); - addLog('info', `Loaded ${toolsResponse.tools.length} tools.`); - setState('ready'); // Final success state - connectingRef.current = false; - connectAttemptRef.current = 0; // Reset attempt counter on success - } - } catch (toolErr) { - // If loading tools fails, we are still connected but log the error. - addLog('error', `Failed to load tools: ${toolErr instanceof Error ? toolErr.message : String(toolErr)}`); - // Consider if this should be a 'failed' state or just 'ready' with no tools - if (isMountedRef.current) { - setState('ready'); // Stay ready, but tools array is empty/stale - setTools([]); // Ensure tools are empty on error - } - connectingRef.current = false; // Still finish connection attempt - } - - } catch (connectErr) { - // Handle errors during the client.connect() phase - addLog('debug', 'Client connect error:', connectErr); - - if (connectErr instanceof UnauthorizedError || (connectErr instanceof Error && (connectErr.message.includes('Unauthorized') || connectErr.message.includes('401')))) { - addLog('info', 'Authentication required.'); - - // Discover OAuth metadata - if (!metadataRef.current) { - try { - addLog('info', 'Discovering OAuth metadata...'); - metadataRef.current = await discoverOAuthMetadata(url); - assert(metadataRef.current, 'Server does not provide OAuth metadata.'); - addLog('info', 'OAuth metadata discovered:', metadataRef.current); - } catch (discoverErr) { - failConnection(`Failed to discover OAuth metadata: ${discoverErr instanceof Error ? discoverErr.message : String(discoverErr)}`, discoverErr instanceof Error ? discoverErr : undefined); - return; // Stop connection process + if (isMountedRef.current && clientRef.current /* Check if still relevant */) { + connect() // Re-run the connection logic } - } - - // We have metadata, attempt to open the popup - await openAuthPopup(); // This sets state to 'authenticating' or 'failed' - - } else { - // Handle other connection errors (network, server unavailable, etc.) - failConnection(`Failed to connect: ${connectErr instanceof Error ? connectErr.message : String(connectErr)}`, connectErr instanceof Error ? connectErr : undefined); + }, delay) + setState('connecting') + } else if (currentState !== 'failed' && currentState !== 'authenticating') { + // If closed unexpectedly and not already failed/authenticating + failConnection('Connection closed unexpectedly.') } } - // Note: connectingRef is set to false within success/failure paths above or in openAuthPopup's failConnection. + // 5. Connect Client & Handle Auth + try { + addLog('info', 'Connecting client...') + assert(clientRef.current, 'Client not initialized') + assert(transportRef.current, 'Transport not initialized') + + await clientRef.current.connect(transportRef.current) + + addLog('info', 'Client connected. Loading tools...') + setState('loading') + + const toolsResponse = await clientRef.current.request({ method: 'tools/list' }, ListToolsResultSchema) + if (isMountedRef.current) { + setTools(toolsResponse.tools) + addLog('info', `Loaded ${toolsResponse.tools.length} tools.`) + setState('ready') + connectingRef.current = false + connectAttemptRef.current = 0 // Reset on success + } + } catch (connectErr) { + addLog('debug', 'Client connect error:', connectErr) + const errorInstance = connectErr instanceof Error ? connectErr : new Error(String(connectErr)) + + // Check if it's specifically an UnauthorizedError or smells like one + // NOTE: The SSE Transport might wrap the underlying 401. Check for UnauthorizedError first. + if ( + errorInstance instanceof UnauthorizedError || + errorInstance.message.includes('Unauthorized') || + errorInstance.message.includes('401') + ) { + addLog('info', 'Authentication required. Initiating SDK auth flow...') + setState('authenticating') + // Start auth timeout timer + if (authTimeoutRef.current) clearTimeout(authTimeoutRef.current) + authTimeoutRef.current = setTimeout(() => { + addLog('warn', 'Authentication timed out.') + if (isMountedRef.current && state === 'authenticating') { + // Check state again + failConnection('Authentication timed out. Please try again or use the manual link if available.') + } + }, AUTH_TIMEOUT) + + try { + assert(authProviderRef.current, 'Auth Provider not available for auth flow') + // This call will handle discovery, checks, token refresh attempts, + // and finally trigger provider.redirectToAuthorization if needed. + const authResult = await auth(authProviderRef.current, { serverUrl: url }) + // If auth() resolves with "AUTHORIZED", it means refresh worked or existing token is valid. + if (authResult === 'AUTHORIZED') { + addLog('info', 'Authentication successful via existing token or refresh. Reconnecting transport...') + if (authTimeoutRef.current) clearTimeout(authTimeoutRef.current) // Stop timeout + // We might need to reconnect the transport or the whole client + // Let's try reconnecting the client fully + connectingRef.current = false // Allow reconnect + connect() // Re-initiate connection flow + } + // If authResult is 'REDIRECT', the provider's redirectToAuthorization was called. + // We are now waiting for the popup callback. State is already 'authenticating'. + else if (authResult === 'REDIRECT') { + addLog('info', 'Redirecting for authentication. Waiting for callback...') + // The state is 'authenticating', timeout is running. + } + } catch (sdkAuthError) { + // Handle errors from the SDK's auth() function itself (e.g., metadata fetch fail) + if (authTimeoutRef.current) clearTimeout(authTimeoutRef.current) + failConnection( + `Failed to initiate authentication: ${sdkAuthError instanceof Error ? sdkAuthError.message : String(sdkAuthError)}`, + sdkAuthError instanceof Error ? sdkAuthError : undefined, + ) + } + } else { + // Handle other connection errors + failConnection(`Failed to connect: ${errorInstance.message}`, errorInstance) + } + // Set connectingRef false only if not waiting for redirect + if (state !== 'authenticating') { + connectingRef.current = false + } + } }, [ - url, storageKeyPrefix, clientName, clientUri, callbackUrl, - addLog, failConnection, disconnect, openAuthPopup, - // Note: state and autoReconnect are accessed but shouldn't trigger recreation of connect - // as they're only used conditionally inside inner functions - ]); - - - // ===== Public API Methods ===== + url, + storageKeyPrefix, + clientName, + clientUri, + callbackUrl, + clientConfig, // Provider/Client configs + addLog, + failConnection, + disconnect, // Core actions + state, + autoReconnect, // State/Config used in handlers + ]) const callTool = useCallback( async (name: string, args?: Record) => { if (state !== 'ready' || !clientRef.current) { - throw new Error(`MCP client is not ready (current state: ${state}). Cannot call tool "${name}".`); + throw new Error(`MCP client is not ready (state: ${state}). Cannot call tool "${name}".`) } - addLog('info', `Calling tool: ${name}`, args); + addLog('info', `Calling tool: ${name}`, args) try { - const result = await clientRef.current.request( - { - method: 'tools/call', - params: { name, arguments: args }, - }, - CallToolResultSchema, // Optional: Validate result against schema - ); - addLog('info', `Tool "${name}" call successful.`); - return result; // Assuming CallToolResultSchema has a `result` field + const result = await clientRef.current.request({ method: 'tools/call', params: { name, arguments: args } }, CallToolResultSchema) + addLog('info', `Tool "${name}" call successful:`, result) + return result } catch (err) { - addLog('error', `Error calling tool "${name}": ${err instanceof Error ? err.message : String(err)}`); - // TODO: Should specific JSON-RPC errors be handled differently? - throw err; // Re-throw the error + addLog('error', `Error calling tool "${name}": ${err instanceof Error ? err.message : String(err)}`, err) + // TODO: Check if error is UnauthorizedError and trigger re-authentication? + if ( + err instanceof UnauthorizedError || + (err instanceof Error && (err.message.includes('Unauthorized') || err.message.includes('401'))) + ) { + addLog('warn', 'Tool call unauthorized, attempting re-authentication...') + setState('authenticating') // Signal auth needed + // Don't disconnect, just trigger auth + assert(authProviderRef.current, 'Auth Provider not available for tool re-auth') + auth(authProviderRef.current, { serverUrl: url }).catch((sdkAuthError) => { + failConnection( + `Re-authentication failed: ${sdkAuthError instanceof Error ? sdkAuthError.message : String(sdkAuthError)}`, + sdkAuthError instanceof Error ? sdkAuthError : undefined, + ) + }) + } + throw err } }, - [state, addLog], - ); + [state, addLog, url], // Add url dependency for re-auth attempt + ) const retry = useCallback(() => { + // Retry should attempt to connect again, regardless of the specific failure reason if (state === 'failed') { - addLog('info', 'Manual retry requested...'); - // Reset attempt counter? Or let connect handle it? Let connect handle it. - disconnect().then(() => { - if (isMountedRef.current) connect(); - }); + addLog('info', 'Retry requested...') + // Reset state and attempt connection + setState('discovering') + setError(undefined) + setAuthUrl(undefined) + connect() } else { - addLog('warn', `Retry called but state is not 'failed' (state: ${state}). Ignoring.`); + addLog('warn', `Retry called but state is not 'failed' (state: ${state}). Ignoring.`) } - }, [state, addLog, disconnect, connect]); + }, [state, addLog, connect]) - const authenticate = useCallback(async (): Promise => { - addLog('info', 'Manual authentication requested...'); - if (state === 'failed' && error?.includes('popup blocked')) { - // If we failed due to popup blocking, try opening the popup again - return openAuthPopup(); - } else if (!metadataRef.current && (state === 'failed' || state === 'discovering')) { - // If we failed before discovering metadata, try full connect again - addLog('info', 'Metadata not yet discovered, attempting full reconnect...'); - await disconnect(); - if (isMountedRef.current) await connect(); // connect will handle auth discovery - // We don't know the auth URL yet in this case - return undefined; - } else if (metadataRef.current) { - // If we have metadata, try opening popup directly - return openAuthPopup(); - } else { - addLog('warn', 'Cannot manually authenticate in current state:', state); - return undefined; + const authenticate = useCallback(() => { + addLog('info', 'Manual authentication requested...') + // Option 1: If failed, just retry connection which includes auth logic + if (state === 'failed') { + addLog('info', 'Attempting to reconnect and authenticate...') + retry() // Retry will call connect() } - }, [state, error, addLog, openAuthPopup, disconnect, connect]); - + // Option 2: If stuck authenticating (popup possibly blocked), maybe re-trigger + else if (state === 'authenticating') { + addLog('warn', 'Already attempting authentication. Check for blocked popups or wait for timeout.') + // Optionally, try retrieving and setting the manual URL again + const manualUrl = authProviderRef.current?.getLastAttemptedAuthUrl() + if (manualUrl && !authUrl) { + // Only set if not already set + setAuthUrl(manualUrl) + addLog('info', 'Manual authentication URL retrieved:', manualUrl) + } + // Maybe try calling SDK auth again? Could be risky if popup is open. + // assert(authProviderRef.current, "Auth Provider not available for manual auth trigger"); + // auth(authProviderRef.current, { serverUrl: url }).catch(failConnection); + } + // Option 3: If already ready, maybe force a refresh/re-auth? Less common. + else { + addLog( + 'info', + 'Client not in a state requiring manual authentication trigger (state: ${state}). If needed, try disconnecting and reconnecting.', + ) + } + }, [state, addLog, retry, authUrl]) // Add authUrl const clearStorage = useCallback(() => { if (authProviderRef.current) { - const count = authProviderRef.current.clearStorage(); - addLog('info', `Cleared ${count} item(s) from localStorage for ${url}.`); - // Also clear related state in the hook - metadataRef.current = undefined; - pendingAuthUrlRef.current = undefined; - setAuthUrl(undefined); - // Consider if a disconnect/reconnect is needed after clearing storage - addLog('info', 'Storage cleared. Recommend disconnecting and reconnecting if currently connected.'); - disconnect(); // Force disconnect as auth state is gone + const count = authProviderRef.current.clearStorage() + addLog('info', `Cleared ${count} item(s) from localStorage for ${url}.`) + setAuthUrl(undefined) // Clear any displayed manual URL + // Force disconnect and state reset as auth info is gone + disconnect().then(() => { + if (isMountedRef.current) { + setState('discovering') // Ready for a fresh connect attempt + } + }) } else { - addLog('warn', 'Auth provider not initialized, cannot clear storage.'); - // Attempt to clear based on prefix/hash anyway? Risky. - // For safety, only clear if provider exists. + addLog('warn', 'Auth provider not initialized, cannot clear storage.') } - }, [addLog, url, disconnect]); - + }, [addLog, url, disconnect]) // ===== Effects ===== // Effect for handling auth callback messages from popup useEffect(() => { const messageHandler = (event: MessageEvent) => { - // Security: Check origin - if (event.origin !== window.location.origin) { - addLog('debug', `Ignoring message from different origin: ${event.origin}`); - return; - } + if (event.origin !== window.location.origin) return if (event.data && event.data.type === 'mcp_auth_callback') { - addLog('info', 'Received auth callback message from popup window.', event.data); - clearTimeout(authTimeoutRef.current!); // Stop the auth timeout - - // Check if the message corresponds to this hook instance if multiple hooks might exist - // (Requires serverUrlHash to be passed back in message data from callback handler) - // const currentHash = authProviderRef.current?.serverUrlHash; // Need to expose serverUrlHash or compare URL - // if (currentHash && event.data.serverUrlHash && currentHash !== event.data.serverUrlHash) { - // addLog('debug', 'Ignoring auth callback for different server instance.'); - // return; - // } + addLog('info', 'Received auth callback message.', event.data) + if (authTimeoutRef.current) clearTimeout(authTimeoutRef.current) // Stop timeout if (event.data.success) { - addLog('info', 'Authentication successful via popup. Reconnecting...'); - // The callback saved the token. Disconnect and reconnect to use it. - disconnect().then(() => { - if (isMountedRef.current) connect(); - }); + addLog('info', 'Authentication successful via popup. Reconnecting client...') + // Tokens are saved by callback handler via provider.saveTokens. + // Re-run connect() to establish the session with the new tokens. + connectingRef.current = false // Allow connect again + connect() } else { // Auth failed in the popup/callback handler - failConnection(`Authentication failed: ${event.data.error || 'Unknown reason.'}`); + failConnection(`Authentication failed in callback: ${event.data.error || 'Unknown reason.'}`) } } - }; - - window.addEventListener('message', messageHandler); - addLog('debug', 'Auth callback message listener added.'); + } + window.addEventListener('message', messageHandler) + addLog('debug', 'Auth callback message listener added.') return () => { - window.removeEventListener('message', messageHandler); - addLog('debug', 'Auth callback message listener removed.'); - clearTimeout(authTimeoutRef.current!); - }; - }, [addLog, failConnection, disconnect, connect]); // Dependencies: functions that handle success/failure + window.removeEventListener('message', messageHandler) + addLog('debug', 'Auth callback message listener removed.') + if (authTimeoutRef.current) clearTimeout(authTimeoutRef.current) + } + }, [addLog, failConnection, connect]) // connect needed for success path - // Effect for initial connection and auto-retry + // Effect for initial connection useEffect(() => { - // Store mount status - isMountedRef.current = true; - addLog('debug', 'Component mounted, initiating connection.'); - connectAttemptRef.current = 0; // Reset attempts on mount/URL change - connect(); // Initial connection attempt + isMountedRef.current = true + addLog('debug', 'useMcp mounted, initiating connection.') + connectAttemptRef.current = 0 + // Ensure provider exists before connecting + if (!authProviderRef.current) { + authProviderRef.current = new BrowserOAuthClientProvider(url, { + storageKeyPrefix, + clientName, + clientUri, + callbackUrl, + }) + addLog('debug', 'BrowserOAuthClientProvider initialized on mount.') + } + connect() // Initial connection attempt - // Cleanup on unmount return () => { - isMountedRef.current = false; - addLog('debug', 'Component unmounting, disconnecting.'); - disconnect(); - }; + isMountedRef.current = false + addLog('debug', 'useMcp unmounting, disconnecting.') + disconnect(true) // Quiet disconnect on unmount + } // eslint-disable-next-line react-hooks/exhaustive-deps }, [ - // Reconnect if these core config options change - url, + url, // Reconnect if URL changes storageKeyPrefix, callbackUrl, - // clientName, clientUri? Only if DCR is used. - // clientConfig.name, clientConfig.version? Usually stable. - // connect, disconnect // connect/disconnect functions are stable due to useCallback - ]); + clientName, + clientUri, // Reconnect if provider config changes + // connect, disconnect are stable due to useCallback + ]) // Effect for auto-retry logic useEffect(() => { - let retryTimeoutId: number | null = null; + let retryTimeoutId: number | null = null if (state === 'failed' && autoRetry && connectAttemptRef.current > 0) { - // Only retry if the *initial* connection or subsequent retries failed - const delay = typeof autoRetry === 'number' ? autoRetry : DEFAULT_RETRY_DELAY; - addLog('info', `Connection failed, auto-retrying in ${delay}ms...`); + const delay = typeof autoRetry === 'number' ? autoRetry : DEFAULT_RETRY_DELAY + addLog('info', `Connection failed, auto-retrying in ${delay}ms...`) retryTimeoutId = setTimeout(() => { - if (isMountedRef.current && state === 'failed') { // Check state again before retrying - retry(); + if (isMountedRef.current && state === 'failed') { + retry() } - }, delay); + }, delay) } - return () => { - if (retryTimeoutId) { - clearTimeout(retryTimeoutId); - } - }; - }, [state, autoRetry, retry, addLog]); // Depends on state and retry config + if (retryTimeoutId) clearTimeout(retryTimeoutId) + } + }, [state, autoRetry, retry, addLog]) - - // Return the public API return { state, tools, error, log, - authUrl, + authUrl, // Expose the manually retrieved auth URL callTool, retry, disconnect, authenticate, clearStorage, - }; + } }