jason/google-mcp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

really dress this puppy up

Browse Source
This commit is contained in:
Taylor Wilsdon
2025-06-07 16:16:48 -04:00
parent aaa39a5b1b
commit 3b8e44b089
5 changed files with 274 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files

121
auth/oauth_callback_server.py
View File

@@ -13,11 +13,11 @@ from typing import Optional, Dict, Any
import socket
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
import uvicorn
from auth.google_auth import handle_auth_callback, CONFIG_CLIENT_SECRETS_PATH
from config.google_config import OAUTH_STATE_TO_SESSION_ID_MAP, SCOPES
from auth.oauth_responses import create_error_response, create_success_response, create_server_error_response
logger = logging.getLogger(__name__)
@@ -26,20 +26,20 @@ class MinimalOAuthServer:
Minimal HTTP server for OAuth callbacks in stdio mode.
Only starts when needed and uses the same port (8000) as streamable-http mode.
"""
def __init__(self, port: int = 8000):
self.port = port
self.app = FastAPI()
self.server = None
self.server_thread = None
self.is_running = False
# Setup the callback route
self._setup_callback_route()
def _setup_callback_route(self):
"""Setup the OAuth callback route."""
@self.app.get("/oauth2callback")
async def oauth_callback(request: Request):
"""Handle OAuth callback - same logic as in core/server.py"""
@@ -50,28 +50,12 @@ class MinimalOAuthServer:
if error:
error_message = f"Authentication failed: Google returned an error: {error}. State: {state}."
logger.error(error_message)
return HTMLResponse(content=f"""
<html><head><title>Authentication Error</title></head>
<body style="font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; max-width: 600px; margin: 40px auto; padding: 20px; text-align: center;">
<h2 style="color: #d32f2f;">Authentication Error</h2>
<p>{error_message}</p>
<p>Please ensure you grant the requested permissions. You can close this window and try again.</p>
<script>setTimeout(function() {{ window.close(); }}, 10000);</script>
</body></html>
""", status_code=400)
return create_error_response(error_message)
if not code:
error_message = "Authentication failed: No authorization code received from Google."
logger.error(error_message)
return HTMLResponse(content=f"""
<html><head><title>Authentication Error</title></head>
<body style="font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; max-width: 600px; margin: 40px auto; padding: 20px; text-align: center;">
<h2 style="color: #d32f2f;">Authentication Error</h2>
<p>{error_message}</p>
<p>You can close this window and try again.</p>
<script>setTimeout(function() {{ window.close(); }}, 10000);</script>
</body></html>
""", status_code=400)
return create_error_response(error_message)
try:
logger.info(f"OAuth callback: Received code (state: {state}). Attempting to exchange for tokens.")
@@ -94,60 +78,25 @@ class MinimalOAuthServer:
log_session_part = f" (linked to session: {mcp_session_id})" if mcp_session_id else ""
logger.info(f"OAuth callback: Successfully authenticated user: {verified_user_id} (state: {state}){log_session_part}.")
# Return success page
return HTMLResponse(content=f"""
<html>
<head>
<title>Authentication Successful</title>
<style>
body {{ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; max-width: 600px; margin: 40px auto; padding: 20px; text-align: center; color: #333; border: 1px solid #ddd; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }}
.status {{ color: #4CAF50; font-size: 24px; margin-bottom: 15px; }}
.message {{ margin-bottom: 20px; line-height: 1.6; }}
.user-id {{ font-weight: bold; color: #2a2a2a; }}
.button {{ background-color: #4CAF50; color: white; padding: 12px 25px; border: none; border-radius: 4px; cursor: pointer; font-size: 16px; text-decoration: none; display: inline-block; margin-top: 10px; }}
</style>
<script> setTimeout(function() {{ window.close(); }}, 10000); </script>
</head>
<body>
<div class="status">✅ Authentication Successful</div>
<div class="message">
You have successfully authenticated as <span class="user-id">{verified_user_id}</span>.
Credentials have been saved.
</div>
<div class="message">
You can now close this window and **retry your original command** in the application.
</div>
<button class="button" onclick="window.close()">Close Window</button>
</body>
</html>
""")
# Return success page using shared template
return create_success_response(verified_user_id)
except Exception as e:
error_message_detail = f"Error processing OAuth callback (state: {state}): {str(e)}"
logger.error(error_message_detail, exc_info=True)
return HTMLResponse(content=f"""
<html>
<head><title>Authentication Processing Error</title></head>
<body style="font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; max-width: 600px; margin: 40px auto; padding: 20px; text-align: center;">
<h2 style="color: #d32f2f;">Authentication Processing Error</h2>
<p>An unexpected error occurred while processing your authentication: {str(e)}</p>
<p>Please try again. You can close this window.</p>
<script>setTimeout(function() {{ window.close(); }}, 10000);</script>
</body>
</html>
""", status_code=500)
return create_server_error_response(str(e))
def start(self) -> bool:
"""
Start the minimal OAuth server.
Returns:
True if started successfully, False otherwise
"""
if self.is_running:
logger.info("Minimal OAuth server is already running")
return True
# Check if port is available
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -155,27 +104,27 @@ class MinimalOAuthServer:
except OSError:
logger.error(f"Port {self.port} is already in use. Cannot start minimal OAuth server.")
return False
def run_server():
"""Run the server in a separate thread."""
try:
config = uvicorn.Config(
self.app,
host="localhost",
self.app,
host="localhost",
port=self.port,
log_level="warning",
access_log=False
)
self.server = uvicorn.Server(config)
asyncio.run(self.server.serve())
except Exception as e:
logger.error(f"Minimal OAuth server error: {e}", exc_info=True)
# Start server in background thread
self.server_thread = threading.Thread(target=run_server, daemon=True)
self.server_thread.start()
# Wait for server to start
max_wait = 3.0
start_time = time.time()
@@ -190,26 +139,26 @@ class MinimalOAuthServer:
except Exception:
pass
time.sleep(0.1)
logger.error(f"Failed to start minimal OAuth server on port {self.port}")
return False
def stop(self):
"""Stop the minimal OAuth server."""
if not self.is_running:
return
try:
if self.server:
if hasattr(self.server, 'should_exit'):
self.server.should_exit = True
if self.server_thread and self.server_thread.is_alive():
self.server_thread.join(timeout=3.0)
self.is_running = False
logger.info(f"Minimal OAuth server stopped (port {self.port})")
except Exception as e:
logger.error(f"Error stopping minimal OAuth server: {e}", exc_info=True)
@@ -220,11 +169,11 @@ _minimal_oauth_server: Optional[MinimalOAuthServer] = None
def get_oauth_redirect_uri(transport_mode: str = "stdio", port: int = 8000) -> str:
"""
Get the appropriate OAuth redirect URI based on transport mode.
Args:
transport_mode: "stdio" or "streamable-http"
port: Port number (default 8000)
Returns:
OAuth redirect URI
"""
@@ -233,30 +182,30 @@ def get_oauth_redirect_uri(transport_mode: str = "stdio", port: int = 8000) -> s
def ensure_oauth_callback_available(transport_mode: str = "stdio", port: int = 8000) -> bool:
"""
Ensure OAuth callback endpoint is available for the given transport mode.
For streamable-http: Assumes the main server is already running
For stdio: Starts a minimal server if needed
Args:
transport_mode: "stdio" or "streamable-http"
transport_mode: "stdio" or "streamable-http"
port: Port number (default 8000)
Returns:
True if callback endpoint is available, False otherwise
"""
global _minimal_oauth_server
if transport_mode == "streamable-http":
# In streamable-http mode, the main FastAPI server should handle callbacks
logger.debug("Using existing FastAPI server for OAuth callbacks (streamable-http mode)")
return True
elif transport_mode == "stdio":
# In stdio mode, start minimal server if not already running
if _minimal_oauth_server is None:
logger.info(f"Creating minimal OAuth server instance for port {port}")
_minimal_oauth_server = MinimalOAuthServer(port)
if not _minimal_oauth_server.is_running:
logger.info("Starting minimal OAuth server for stdio mode")
result = _minimal_oauth_server.start()
@@ -268,7 +217,7 @@ def ensure_oauth_callback_available(transport_mode: str = "stdio", port: int = 8
else:
logger.info("Minimal OAuth server is already running")
return True
else:
logger.error(f"Unknown transport mode: {transport_mode}")
return False