Gemini 2.5 refactor

src/auth/callback.ts

@@ -0,0 +1,179 @@
+import { exchangeAuthorization } from '@modelcontextprotocol/sdk/client/auth.js';
+import { OAuthClientInformation, OAuthMetadata } from '@modelcontextprotocol/sdk/shared/auth.js';
+import { StoredState } from './types.js';
+
+/**
+ * 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.
+ */
+export async function onMcpAuthorization(
+  query: Record<string, string>,
+  {
+    storageKeyPrefix = 'mcp:auth',
+  }: {
+    storageKeyPrefix?: string;
+  } = {},
+): Promise<{ success: boolean; error?: string }> {
+  const logPrefix = `[${storageKeyPrefix}-callback]`;
+  console.log(`${logPrefix} Handling OAuth callback with query:`, query);
+
+  try {
+    const code = query.code;
+    const state = query.state;
+    const error = query.error;
+    const errorDescription = query.error_description;
+
+    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.');
+    }
+
+    // Retrieve the stored state using the state parameter from the query
+    const stateKey = `${storageKeyPrefix}:state_${state}`;
+    const storedStateJSON = localStorage.getItem(stateKey);
+    // Remove the state key immediately after retrieval (it's single-use)
+    localStorage.removeItem(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);
+    } catch (e) {
+      throw new Error('Failed to parse stored OAuth state.');
+    }
+
+    // Validate expiry
+    if (storedState.expiry < Date.now()) {
+      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);
+    // Remove the code verifier after retrieving it
+    localStorage.removeItem(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.`);
+    }
+
+    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...`);
+
+    // 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
+      redirectUrl: new URL('/oauth/callback', window.location.origin).toString(), // TODO: Make this configurable or get from provider instance?
+    });
+
+    console.log(`${logPrefix} Token exchange successful. Saving tokens...`);
+
+    // Save the obtained tokens
+    localStorage.setItem(tokensKey, JSON.stringify(tokens));
+
+    // Clean up the persisted manual auth URL if it exists
+    localStorage.removeItem(authUrlKey);
+
+    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
+    }
+
+    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
+    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
+        },
+        window.location.origin,
+      );
+    }
+
+    // Display error in the callback window itself for better debugging
+    try {
+      document.body.innerHTML = `
+        <div style="font-family: sans-serif; padding: 20px;">
+          <h1>Authentication Error</h1>
+          <p style="color: red; background-color: #ffebeb; border: 1px solid red; padding: 10px; border-radius: 4px;">
+            ${errorMessage}
+          </p>
+          <p>You can close this window.</p>
+          <pre style="font-size: 0.8em; color: #555; margin-top: 20px; white-space: pre-wrap;">${
+        err instanceof Error ? err.stack : ''
+      }</pre>
+        </div>
+      `;
+    } 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 };
+  }
+}