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

working build for all functionality on streamablehttp with auth_session_id

Browse Source
This commit is contained in:
Taylor Wilsdon
2025-05-11 15:37:44 -04:00
parent 984d4065ff
commit e681d265bf
7 changed files with 441 additions and 726 deletions
Download Patch File Download Diff File Expand all files Collapse all files

337
auth/callback_server.py
View File

@@ -1,337 +0,0 @@
# auth/callback_server.py
import http.server
import logging
import os
import socketserver
import threading
import urllib.parse
import webbrowser
from typing import Callable, Optional, Dict, Any, Literal
from oauthlib.oauth2.rfc6749.errors import InsecureTransportError
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def get_response_html(
title: str,
status: Literal["success", "error"],
message: str,
show_close_button: bool = True,
auto_close_seconds: int = 0
) -> str:
"""Generate HTML response for OAuth callback.
Args:
title: Page title
status: 'success' or 'error'
message: Message to display to the user
show_close_button: Whether to show a close button
auto_close_seconds: Auto-close after this many seconds (0 to disable)
Returns:
HTML content as a string
"""
color = "#4CAF50" if status == "success" else "#d32f2f"
close_button = """
<button class="button" onclick="window.close()">Close Window</button>
""" if show_close_button else ""
auto_close_script = f"""
<script>
setTimeout(function() {{ window.close(); }}, {auto_close_seconds * 1000});
</script>
""" if auto_close_seconds > 0 else ""
return f"""<!DOCTYPE html>
<html>
<head>
<title>{title}</title>
<style>
body {{
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
max-width: 500px;
margin: 40px auto;
padding: 20px;
text-align: center;
color: #333;
}}
.status {{
color: {color};
font-size: 24px;
margin-bottom: 20px;
}}
.message {{
margin-bottom: 30px;
line-height: 1.5;
}}
.button {{
background-color: {color};
color: white;
padding: 10px 20px;
border: none;
border-radius: 4px;
cursor: pointer;
font-size: 16px;
}}
</style>
</head>
<body>
<div class="status">{title}</div>
<div class="message">{message}</div>
{close_button}
{auto_close_script}
</body>
</html>"""
class OAuthCallbackHandler(http.server.BaseHTTPRequestHandler):
"""Handler for OAuth callback requests."""
# Class variables to store callback functions by state
callback_registry: Dict[str, Callable] = {}
@classmethod
def register_callback(cls, state: str, callback: Callable) -> None:
"""Register a callback function for a specific state parameter."""
logger.info(f"Registering callback for state: {state}")
cls.callback_registry[state] = callback
@classmethod
def unregister_callback(cls, state: str) -> None:
"""Unregister a callback function for a specific state parameter."""
if state in cls.callback_registry:
logger.info(f"Unregistering callback for state: {state}")
del cls.callback_registry[state]
def do_GET(self):
"""Handle GET requests to the callback endpoint."""
request_thread_id = threading.get_ident()
logger.info(f"[Handler {request_thread_id}] GET request received for path: {self.path}")
try:
# Parse the URL and extract query parameters
parsed_url = urllib.parse.urlparse(self.path)
query_params = urllib.parse.parse_qs(parsed_url.query)
# Check if we're handling the OAuth callback
if parsed_url.path == '/callback':
# Extract authorization code and state
code = query_params.get('code', [''])[0]
state = query_params.get('state', [''])[0]
logger.info(f"Received OAuth callback with code: {code[:10]}... and state: {state}")
# Show success page to the user
self.send_response(200)
self.send_header('Content-Type', 'text/html')
self.end_headers()
html_content = get_response_html(
title="Authentication Successful",
status="success",
message="You have successfully authenticated with Google. You can now close this window and return to your application.",
show_close_button=True,
auto_close_seconds=10
)
self.wfile.write(html_content.encode())
try:
# Ensure OAUTHLIB_INSECURE_TRANSPORT is set
if 'OAUTHLIB_INSECURE_TRANSPORT' not in os.environ:
logger.warning("OAUTHLIB_INSECURE_TRANSPORT not set. Setting it for localhost development.")
os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = '1'
# Call the appropriate callback function based on state
if state in OAuthCallbackHandler.callback_registry and code:
logger.info(f"[Handler {request_thread_id}] Found callback for state: {state}")
callback_function = OAuthCallbackHandler.callback_registry[state]
logger.info(f"[Handler {request_thread_id}] Preparing to call callback function in new thread.")
callback_thread = threading.Thread(
target=callback_function,
args=(code, state),
daemon=True
)
callback_thread.start()
logger.info(f"[Handler {request_thread_id}] Callback function thread started (ID: {callback_thread.ident}).")
# Unregister the callback after it's been called
OAuthCallbackHandler.unregister_callback(state)
else:
logger.warning(f"[Handler {request_thread_id}] No callback registered for state: {state} or no code received.")
except InsecureTransportError as e:
logger.error(f"[Handler {request_thread_id}] InsecureTransportError: {e}. Ensure OAUTHLIB_INSECURE_TRANSPORT is set for localhost development.")
self.send_response(400)
self.send_header('Content-Type', 'text/html')
self.end_headers()
error_html = get_response_html(
title="OAuth Error: Insecure Transport",
status="error",
message="The OAuth flow requires HTTPS or explicit allowance of HTTP for localhost development. Please ensure OAUTHLIB_INSECURE_TRANSPORT is set in your environment.",
show_close_button=False
)
self.wfile.write(error_html.encode())
return
# Note: We no longer shut down the server after handling a callback
# This allows it to handle multiple auth flows over time
else:
# Handle other paths with a 404 response
self.send_response(404)
self.send_header('Content-Type', 'text/html')
self.end_headers()
error_html = get_response_html(
title="Not Found",
status="error",
message="The requested resource was not found.",
show_close_button=False
)
self.wfile.write(error_html.encode())
except Exception as e:
logger.error(f"[Handler {request_thread_id}] Error handling callback request: {e}", exc_info=True)
try:
self.send_response(500)
self.send_header('Content-Type', 'text/html')
self.end_headers()
error_html = get_response_html(
title="Internal Server Error",
status="error",
message=f"An error occurred while processing your request: {str(e)}",
show_close_button=False
)
self.wfile.write(error_html.encode())
except Exception as send_error:
logger.error(f"[Handler {request_thread_id}] Error sending 500 response: {send_error}")
def log_message(self, format, *args):
"""Override to use our logger instead of printing to stderr."""
logger.info(f"{self.address_string()} - {format%args}")
class OAuthCallbackServer:
"""Server to handle OAuth callbacks."""
def __init__(self,
port: int = 8080,
auto_open_browser: bool = True):
"""
Initialize the callback server.
Args:
port: Port to listen on (default: 8080)
auto_open_browser: Whether to automatically open the browser
"""
self.port = port
self.server = None
self.server_thread = None
self.auto_open_browser = auto_open_browser
def start(self) -> Dict[str, Any]:
"""
Start the callback server in a separate thread.
Returns:
Dict containing server status and port information:
{
'success': bool,
'port': int,
'message': str
}
"""
if self.server:
logger.warning("Server is already running")
return {'success': False, 'port': self.port, 'message': 'Server is already running'}
original_port = self.port
max_port = 8090 # Try ports 8080-8090
def serve():
thread_id = threading.get_ident()
logger.info(f"[Server Thread {thread_id}] Starting serve_forever loop.")
try:
self.server.serve_forever()
except Exception as serve_e:
logger.error(f"[Server Thread {thread_id}] Exception in serve_forever: {serve_e}", exc_info=True)
finally:
logger.info(f"[Server Thread {thread_id}] serve_forever loop finished.")
# Ensure server_close is called even if shutdown wasn't clean
try:
if self.server:
self.server.server_close()
except Exception as close_e:
logger.error(f"[Server Thread {thread_id}] Error during server_close: {close_e}")
try:
while self.port <= max_port:
try:
# Create and start the server
self.server = socketserver.TCPServer(('localhost', self.port), OAuthCallbackHandler)
logger.info(f"Starting OAuth callback server on port {self.port}")
if self.port != original_port:
logger.info(f"Successfully reassigned from port {original_port} to {self.port}")
# Start the server thread
self.server_thread = threading.Thread(target=serve, daemon=True)
self.server_thread.start()
logger.info(f"OAuth callback server thread started (ID: {self.server_thread.ident}) on http://localhost:{self.port}")
return {
'success': True,
'port': self.port,
'message': f"Server started successfully on port {self.port}"
}
except OSError as e:
if e.errno == 48: # Address already in use
logger.warning(f"Port {self.port} is already in use, trying next port")
self.port += 1
if self.port > max_port:
error_msg = f"Failed to find available port in range {original_port}-{max_port}"
logger.error(error_msg)
return {'success': False, 'port': None, 'message': error_msg}
continue
else:
logger.error(f"Failed to start server: {e}")
return {'success': False, 'port': None, 'message': str(e)}
except Exception as e:
error_msg = f"Failed to start callback server: {e}"
logger.error(error_msg)
return {'success': False, 'port': None, 'message': error_msg}
def stop(self) -> None:
"""Stop the callback server."""
if self.server:
logger.info("Stopping OAuth callback server")
# shutdown() signals serve_forever to stop
self.server.shutdown()
logger.info("Server shutdown() called.")
# Wait briefly for the server thread to finish
if self.server_thread:
self.server_thread.join(timeout=2.0) # Wait up to 2 seconds
if self.server_thread.is_alive():
logger.warning("Server thread did not exit cleanly after shutdown.")
# server_close() is now called in the 'finally' block of the serve() function
self.server = None
self.server_thread = None
logger.info("Server resources released.")
else:
logger.warning("Server is not running")
def open_browser(self, url: str) -> bool:
"""Open the default web browser to the given URL."""
if not self.auto_open_browser:
return False
try:
logger.info(f"Opening browser to: {url}")
webbrowser.open(url)
return True
except Exception as e:
logger.error(f"Failed to open browser: {e}")
return False

74
auth/google_auth.py
View File

@@ -13,21 +13,13 @@ from google.auth.transport.requests import Request
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from auth.callback_server import OAuthCallbackServer, OAuthCallbackHandler
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Constants
DEFAULT_CREDENTIALS_DIR = ".credentials"
DEFAULT_REDIRECT_URI = "http://localhost:8080/callback"
DEFAULT_SERVER_PORT = 8080
# --- Global Variables ---
# Singleton OAuth callback server instance
_oauth_callback_server = None
DEFAULT_REDIRECT_URI = "http://localhost:8000/oauth2callback"
# --- Helper Functions ---
@@ -105,10 +97,7 @@ def load_client_secrets(client_secrets_path: str) -> Dict[str, Any]:
def start_auth_flow(
client_secrets_path: str,
scopes: List[str],
redirect_uri: str = DEFAULT_REDIRECT_URI,
auto_handle_callback: bool = False,
callback_function: Optional[Callable] = None,
port: int = DEFAULT_SERVER_PORT
redirect_uri: str = DEFAULT_REDIRECT_URI
) -> Tuple[str, str]:
"""
Initiates the OAuth 2.0 flow and returns the authorization URL and state.
@@ -117,41 +106,16 @@ def start_auth_flow(
client_secrets_path: Path to the Google client secrets JSON file.
scopes: List of OAuth scopes required.
redirect_uri: The URI Google will redirect to after authorization.
auto_handle_callback: Whether to automatically handle the callback by
using the persistent callback server.
callback_function: Function to call with the code and state when received.
port: Port to run the callback server on, if one is not already running.
Returns:
A tuple containing the authorization URL and the state parameter.
"""
global _oauth_callback_server
try:
# Allow HTTP for localhost in development
if 'OAUTHLIB_INSECURE_TRANSPORT' not in os.environ:
logger.warning("OAUTHLIB_INSECURE_TRANSPORT not set. Setting it for localhost development.")
os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = '1'
# Use or initialize the persistent callback server
if auto_handle_callback:
if _oauth_callback_server is None:
logger.info("Starting OAuth callback server (persistent instance)")
_oauth_callback_server = OAuthCallbackServer(port=port, auto_open_browser=False)
result = _oauth_callback_server.start()
if not result['success']:
logger.error(f"Failed to start callback server: {result['message']}")
return None, None
# Store the actual port being used for future redirect_uri construction
port = result['port']
else:
logger.info(f"Using existing OAuth callback server on port {port}")
# Always use the port from the running server
redirect_uri = f"http://localhost:{port}/callback"
# Set up the OAuth flow
flow = Flow.from_client_secrets_file(
client_secrets_path,
@@ -169,15 +133,6 @@ def start_auth_flow(
)
logger.info(f"Generated authorization URL. State: {state}")
# Register the callback function with the state
if auto_handle_callback and callback_function:
OAuthCallbackHandler.register_callback(state, callback_function)
logger.info(f"Registered callback function for state: {state}")
# Auto-open the browser if requested
if auto_handle_callback and _oauth_callback_server:
_oauth_callback_server.open_browser(authorization_url)
return authorization_url, state
except Exception as e:
@@ -266,44 +221,53 @@ def get_credentials(
Returns:
Valid Credentials object if found and valid/refreshed, otherwise None.
"""
logger.info(f"[get_credentials] Called for user_id: '{user_id}', required_scopes: {required_scopes}")
credential_file_path = _get_user_credential_path(user_id, credentials_base_dir)
logger.info(f"[get_credentials] Attempting to load credentials from: {credential_file_path}")
credentials = _load_credentials(user_id, credentials_base_dir)
if not credentials:
logger.info(f"No stored credentials found for user {user_id}.")
logger.info(f"[get_credentials] No stored credentials found for user '{user_id}' at {credential_file_path}.")
return None
logger.info(f"[get_credentials] Successfully loaded credentials for user '{user_id}'. Scopes: {credentials.scopes}, Valid: {credentials.valid}, Expired: {credentials.expired}")
# Check if scopes are sufficient
if not all(scope in credentials.scopes for scope in required_scopes):
logger.warning(f"Stored credentials for user {user_id} lack required scopes. Need: {required_scopes}, Have: {credentials.scopes}")
logger.warning(f"[get_credentials] Stored credentials for user '{user_id}' lack required scopes. Need: {required_scopes}, Have: {credentials.scopes}")
# Re-authentication is needed to grant missing scopes
return None
logger.info(f"[get_credentials] Stored credentials for user '{user_id}' have sufficient scopes.")
# Check if credentials are still valid or need refresh
if credentials.valid:
logger.info(f"Stored credentials for user {user_id} are valid.")
logger.info(f"[get_credentials] Stored credentials for user '{user_id}' are valid.")
return credentials
elif credentials.expired and credentials.refresh_token:
logger.info(f"Credentials for user {user_id} expired. Attempting refresh.")
logger.info(f"[get_credentials] Credentials for user '{user_id}' expired. Attempting refresh.")
if not client_secrets_path:
logger.error("Client secrets path is required to refresh credentials but was not provided.")
logger.error("[get_credentials] Client secrets path is required to refresh credentials but was not provided.")
# Cannot refresh without client secrets info
return None
try:
# Load client secrets to provide info for refresh
# Note: Credentials object holds client_id/secret if available from initial flow,
# but loading from file is safer if they weren't stored or if using InstalledAppFlow secrets.
logger.info(f"[get_credentials] Attempting to refresh token for user '{user_id}' using client_secrets_path: {client_secrets_path}")
client_config = load_client_secrets(client_secrets_path)
credentials.refresh(Request()) # Pass client_id/secret if needed and not in creds
logger.info(f"Credentials for user {user_id} refreshed successfully.")
logger.info(f"[get_credentials] Credentials for user '{user_id}' refreshed successfully.")
# Save the updated credentials (with potentially new access token)
_save_credentials(user_id, credentials, credentials_base_dir)
return credentials
except Exception as e: # Catch specific refresh errors like google.auth.exceptions.RefreshError
logger.error(f"Error refreshing credentials for user {user_id}: {e}")
logger.error(f"[get_credentials] Error refreshing credentials for user '{user_id}': {e}", exc_info=True)
# Failed to refresh, re-authentication is needed
return None
else:
logger.warning(f"Credentials for user {user_id} are invalid or missing refresh token.")
logger.warning(f"[get_credentials] Credentials for user '{user_id}' are invalid and/or missing refresh token. Valid: {credentials.valid}, Refresh Token Present: {credentials.refresh_token is not None}")
# Invalid and cannot be refreshed, re-authentication needed
return None