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

Google Workspace MCP Response Format Improvements - model-agnostic json envelope format for better compatibility across platforms

Browse Source
This commit is contained in:
Taylor Wilsdon
2025-05-10 17:57:25 -04:00
parent de498150ab
commit 458a1476b0
5 changed files with 711 additions and 337 deletions
Download Patch File Download Diff File Expand all files Collapse all files

343
auth/callback_server.py
View File

@@ -2,21 +2,110 @@
import http.server
import logging
import os
import socketserver
import threading
import urllib.parse
import webbrowser
from typing import Callable, Optional, Dict, Any
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
"""
icon = "✅" if status == "success" else "⚠️"
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">{icon} {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 variable to store the callback function
callback_function: Optional[Callable] = None
# 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."""
@@ -40,95 +129,84 @@ class OAuthCallbackHandler(http.server.BaseHTTPRequestHandler):
self.send_header('Content-Type', 'text/html')
self.end_headers()
html_content = """
<!DOCTYPE html>
<html>
<head>
<title>Google OAuth - Success</title>
<style>
body {
font-family: Arial, sans-serif;
max-width: 600px;
margin: 40px auto;
padding: 20px;
text-align: center;
}
.success {
color: #4CAF50;
font-size: 24px;
margin-bottom: 20px;
}
.info {
color: #555;
margin-bottom: 30px;
}
.close-button {
background-color: #4CAF50;
color: white;
padding: 10px 20px;
border: none;
border-radius: 4px;
cursor: pointer;
font-size: 16px;
}
</style>
</head>
<body>
<div class="success">✅ Authentication Successful!</div>
<div class="info">
You have successfully authenticated with Google.
You can now close this window and return to your application.
</div>
<button class="close-button" onclick="window.close()">Close Window</button>
<script>
// Auto-close after 10 seconds
setTimeout(function() {
window.close();
}, 10000);
</script>
</body>
</html>
"""
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())
# Call the callback function if provided
if OAuthCallbackHandler.callback_function and code:
logger.info(f"[Handler {request_thread_id}] Preparing to call callback function in new thread.")
callback_thread = threading.Thread(
target=OAuthCallbackHandler.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}).")
else:
logger.warning(f"[Handler {request_thread_id}] No callback function set or no code received, skipping callback.")
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'
# Signal the server to shutdown after handling the request
logger.info(f"[Handler {request_thread_id}] Preparing to signal server shutdown in new thread.")
shutdown_thread = threading.Thread(
target=self.server.shutdown,
daemon=True
)
shutdown_thread.start()
logger.info(f"[Handler {request_thread_id}] Server shutdown thread started (ID: {shutdown_thread.ident}). Request handling complete.")
# 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()
self.wfile.write(b"Not Found")
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/plain')
self.send_header('Content-Type', 'text/html')
self.end_headers()
self.wfile.write(f"Internal Server Error: {str(e)}".encode())
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}")
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."""
@@ -137,62 +215,95 @@ class OAuthCallbackHandler(http.server.BaseHTTPRequestHandler):
class OAuthCallbackServer:
"""Server to handle OAuth callbacks."""
def __init__(self,
port: int = 8080,
callback: Optional[Callable] = None,
def __init__(self,
port: int = 8080,
auto_open_browser: bool = True):
"""
Initialize the callback server.
Args:
port: Port to listen on (default: 8080)
callback: Function to call with the code and state
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
# Set the callback function
OAuthCallbackHandler.callback_function = callback
def start(self) -> None:
"""Start the callback server in a separate thread."""
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
return {'success': False, 'port': self.port, 'message': 'Server is already running'}
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}")
# Run the server in a separate thread
def serve():
thread_id = threading.get_ident()
logger.info(f"[Server Thread {thread_id}] Starting serve_forever loop.")
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:
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}")
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}")
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:
logger.error(f"Failed to start callback server: {e}")
raise
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."""
@@ -203,9 +314,9 @@ class OAuthCallbackServer:
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.")
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

61
auth/google_auth.py
View File

@@ -5,13 +5,15 @@ import json
import logging
from typing import List, Optional, Tuple, Dict, Any, Callable
from oauthlib.oauth2.rfc6749.errors import InsecureTransportError
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import Flow, InstalledAppFlow
from google.auth.transport.requests import Request
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from auth.callback_server import OAuthCallbackServer
from auth.callback_server import OAuthCallbackServer, OAuthCallbackHandler
# Configure logging
logging.basicConfig(level=logging.INFO)
@@ -22,6 +24,11 @@ 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
# --- Helper Functions ---
def _get_user_credential_path(user_id: str, base_dir: str = DEFAULT_CREDENTIALS_DIR) -> str:
@@ -111,20 +118,39 @@ def start_auth_flow(
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
starting a local server on the specified port.
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 auto_handle_callback is True.
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:
# Create and start the callback server if auto_handle_callback is enabled
server = None
# 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:
logger.info("Starting OAuth callback server")
server = OAuthCallbackServer(port=port, callback=callback_function, auto_open_browser=False)
server.start()
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(
@@ -143,17 +169,21 @@ 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 server:
server.open_browser(authorization_url)
if auto_handle_callback and _oauth_callback_server:
_oauth_callback_server.open_browser(authorization_url)
return authorization_url, state
except Exception as e:
logger.error(f"Error starting OAuth flow: {e}")
# If we created a server, shut it down
if auto_handle_callback and server:
server.stop()
# We no longer shut down the server after completing the flow
# The persistent server will handle multiple auth flows over time
raise # Re-raise the exception for the caller to handle
def handle_auth_callback(
@@ -183,6 +213,11 @@ def handle_auth_callback(
HttpError: If fetching user info fails.
"""
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'
flow = Flow.from_client_secrets_file(
client_secrets_path,
scopes=scopes,