GEMINI TO THE RESCUE
This commit is contained in:
Glen Maddern
parent
78d6583160
commit
d0372fd8d8
2 changed files with 573 additions and 757 deletions
|
|
@ -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<string | undefined>;
|
||||
authenticate: () => void;
|
||||
/** Clears all stored authentication data (tokens, client info, etc.) for this server URL from localStorage. */
|
||||
clearStorage: () => void;
|
||||
};
|
||||
|
|
@ -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<UseMcpResult['state']>('discovering');
|
||||
const [tools, setTools] = useState<Tool[]>([]);
|
||||
const [error, setError] = useState<string | undefined>(undefined);
|
||||
const [log, setLog] = useState<UseMcpResult['log']>([]);
|
||||
const [authUrl, setAuthUrl] = useState<string | undefined>(undefined); // For manual auth if popup blocked
|
||||
const [state, setState] = useState<UseMcpResult['state']>('discovering')
|
||||
const [tools, setTools] = useState<Tool[]>([])
|
||||
const [error, setError] = useState<string | undefined>(undefined)
|
||||
const [log, setLog] = useState<UseMcpResult['log']>([])
|
||||
// Store the URL for manual fallback if needed
|
||||
const [authUrl, setAuthUrl] = useState<string | undefined>(undefined)
|
||||
|
||||
const clientRef = useRef<Client | null>(null);
|
||||
const transportRef = useRef<SSEClientTransport | null>(null);
|
||||
const authProviderRef = useRef<BrowserOAuthClientProvider | null>(null);
|
||||
const metadataRef = useRef<OAuthMetadata | undefined>(undefined);
|
||||
const pendingAuthUrlRef = useRef<URL | undefined>(undefined); // Stores the URL while waiting for popup
|
||||
const connectingRef = useRef<boolean>(false); // Prevent concurrent connection attempts
|
||||
const isMountedRef = useRef<boolean>(true); // Track component mount status
|
||||
const connectAttemptRef = useRef<number>(0); // Track connection attempts for retry logic
|
||||
const authTimeoutRef = useRef<number | null>(null); // Timer for auth popup
|
||||
|
||||
// ===== Logging Utility =====
|
||||
const clientRef = useRef<Client | null>(null)
|
||||
const transportRef = useRef<SSEClientTransport | null>(null)
|
||||
const authProviderRef = useRef<BrowserOAuthClientProvider | null>(null)
|
||||
const connectingRef = useRef<boolean>(false)
|
||||
const isMountedRef = useRef<boolean>(true)
|
||||
const connectAttemptRef = useRef<number>(0)
|
||||
const authTimeoutRef = useRef<number | null>(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<URL | undefined> => {
|
||||
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<string, unknown>) => {
|
||||
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<string | undefined> => {
|
||||
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,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue