small fixes

Author: ihrpr

examples/clients/simple-auth-client/mcp_simple_auth_client/main.py

@@ -45,18 +45,19 @@ async def set_client_info(self, client_info: OAuthClientInformationFull) -> None
 class CallbackHandler(BaseHTTPRequestHandler):
     """Simple HTTP handler to capture OAuth callback."""
 
-    authorization_code = None
-    state = None
-    error = None
+    def __init__(self, request, client_address, server, callback_data):
+        """Initialize with callback data storage."""
+        self.callback_data = callback_data
+        super().__init__(request, client_address, server)
 
     def do_GET(self):
         """Handle GET request from OAuth redirect."""
         parsed = urlparse(self.path)
         query_params = parse_qs(parsed.query)
 
         if "code" in query_params:
-            CallbackHandler.authorization_code = query_params["code"][0]
-            CallbackHandler.state = query_params.get("state", [None])[0]
+            self.callback_data["authorization_code"] = query_params["code"][0]
+            self.callback_data["state"] = query_params.get("state", [None])[0]
             self.send_response(200)
             self.send_header("Content-type", "text/html")
             self.end_headers()
@@ -70,7 +71,7 @@ def do_GET(self):
             </html>
             """)
         elif "error" in query_params:
-            CallbackHandler.error = query_params["error"][0]
+            self.callback_data["error"] = query_params["error"][0]
             self.send_response(400)
             self.send_header("Content-type", "text/html")
             self.end_headers()
@@ -101,10 +102,26 @@ def __init__(self, port=3000):
         self.port = port
         self.server = None
         self.thread = None
+        self.callback_data = {
+            "authorization_code": None,
+            "state": None,
+            "error": None
+        }
+
+    def _create_handler_with_data(self):
+        """Create a handler class with access to callback data."""
+        callback_data = self.callback_data
+        
+        class DataCallbackHandler(CallbackHandler):
+            def __init__(self, request, client_address, server):
+                super().__init__(request, client_address, server, callback_data)
+        
+        return DataCallbackHandler
 
     def start(self):
         """Start the callback server in a background thread."""
-        self.server = HTTPServer(("localhost", self.port), CallbackHandler)
+        handler_class = self._create_handler_with_data()
+        self.server = HTTPServer(("localhost", self.port), handler_class)
         self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
         self.thread.start()
         print(f"🖥️  Started callback server on http://localhost:{self.port}")
@@ -121,12 +138,16 @@ def wait_for_callback(self, timeout=300):
         """Wait for OAuth callback with timeout."""
         start_time = time.time()
         while time.time() - start_time < timeout:
-            if CallbackHandler.authorization_code:
-                return CallbackHandler.authorization_code
-            elif CallbackHandler.error:
-                raise Exception(f"OAuth error: {CallbackHandler.error}")
+            if self.callback_data["authorization_code"]:
+                return self.callback_data["authorization_code"]
+            elif self.callback_data["error"]:
+                raise Exception(f"OAuth error: {self.callback_data['error']}")
             time.sleep(0.1)
         raise Exception("Timeout waiting for OAuth callback")
+        
+    def get_state(self):
+        """Get the received state parameter."""
+        return self.callback_data["state"]
 
 
 class SimpleAuthClient:
@@ -153,7 +174,7 @@ async def callback_handler() -> tuple[str, str | None]:
                 print("⏳ Waiting for authorization callback...")
                 try:
                     auth_code = callback_server.wait_for_callback(timeout=300)
-                    return auth_code, CallbackHandler.state
+                    return auth_code, callback_server.get_state()
                 finally:
                     callback_server.stop()