From 23f89e49b810842b82d7fdc19b650c89a4bf43ac Mon Sep 17 00:00:00 2001
From: Justin Spahr-Summers <justin@anthropic.com>
Date: Fri, 24 Jan 2025 13:08:39 +0000
Subject: [PATCH] Implement OAuth callback

---
 client/src/App.tsx                      | 11 ++++++-
 client/src/components/OAuthCallback.tsx | 39 +++++++++++++++++++++++++
 client/src/lib/hooks/useConnection.ts   |  2 ++
 client/vite.config.ts                   |  3 ++
 4 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 client/src/components/OAuthCallback.tsx

diff --git a/client/src/App.tsx b/client/src/App.tsx
index f3791b2..9fad257 100644
--- a/client/src/App.tsx
+++ b/client/src/App.tsx
@@ -18,7 +18,7 @@ import {
   ServerNotification,
   Tool,
 } from "@modelcontextprotocol/sdk/types.js";
-import { useEffect, useRef, useState } from "react";
+import { useEffect, useRef, useState, Suspense } from "react";
 
 import { StdErrNotification } from "./lib/notificationTypes";
 
@@ -49,6 +49,15 @@ const PROXY_PORT = params.get("proxyPort") ?? "3000";
 const PROXY_SERVER_URL = `http://localhost:${PROXY_PORT}`;
 
 const App = () => {
+  // Handle OAuth callback route
+  if (window.location.pathname === '/oauth/callback') {
+    const OAuthCallback = React.lazy(() => import('./components/OAuthCallback'));
+    return (
+      <Suspense fallback={<div>Loading...</div>}>
+        <OAuthCallback />
+      </Suspense>
+    );
+  }
   const [resources, setResources] = useState<Resource[]>([]);
   const [resourceTemplates, setResourceTemplates] = useState<
     ResourceTemplate[]
diff --git a/client/src/components/OAuthCallback.tsx b/client/src/components/OAuthCallback.tsx
new file mode 100644
index 0000000..67fc91f
--- /dev/null
+++ b/client/src/components/OAuthCallback.tsx
@@ -0,0 +1,39 @@
+import { useEffect } from 'react';
+import { handleOAuthCallback } from '../lib/auth';
+
+const OAuthCallback = () => {
+  useEffect(() => {
+    const handleCallback = async () => {
+      const params = new URLSearchParams(window.location.search);
+      const code = params.get('code');
+      const serverUrl = sessionStorage.getItem('mcp_server_url');
+
+      if (!code || !serverUrl) {
+        console.error('Missing code or server URL');
+        window.location.href = '/';
+        return;
+      }
+
+      try {
+        const accessToken = await handleOAuthCallback(serverUrl, code);
+        // Store the access token for future use
+        sessionStorage.setItem('mcp_access_token', accessToken);
+        // Redirect back to the main app
+        window.location.href = '/';
+      } catch (error) {
+        console.error('OAuth callback error:', error);
+        window.location.href = '/';
+      }
+    };
+
+    void handleCallback();
+  }, []);
+
+  return (
+    <div className="flex items-center justify-center h-screen">
+      <p className="text-lg text-gray-500">Processing OAuth callback...</p>
+    </div>
+  );
+};
+
+export default OAuthCallback;
\ No newline at end of file
diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts
index 78d3a7d..f12d3cd 100644
--- a/client/src/lib/hooks/useConnection.ts
+++ b/client/src/lib/hooks/useConnection.ts
@@ -166,6 +166,8 @@ export function useConnection({
       } catch (error) {
         console.error("Failed to connect to MCP server:", error);
         if (error instanceof SseError && error.code === 401) {
+          // Store the server URL for the callback handler
+          sessionStorage.setItem('mcp_server_url', sseUrl);
           const redirectUrl = await startOAuthFlow(sseUrl);
           window.location.href = redirectUrl;
         }
diff --git a/client/vite.config.ts b/client/vite.config.ts
index dd3bd01..8ef6cbc 100644
--- a/client/vite.config.ts
+++ b/client/vite.config.ts
@@ -5,6 +5,9 @@ import react from "@vitejs/plugin-react";
 // https://vitejs.dev/config/
 export default defineConfig({
   plugins: [react()],
+  server: {
+    historyApiFallback: true,
+  },
   resolve: {
     alias: {
       "@": path.resolve(__dirname, "./src"),