Download google-workspace/upstream/source/core/server.py from smlflg/Connector: direct link, hf CLI and curl.
- Browser
- Download file 30.2 kB
-
https://huggingface.co/smlflg/Connector/resolve/main/google-workspace/upstream/source/core/server.py
- Command line
-
hf download hf://smlflg/Connector/google-workspace/upstream/source/core/server.py
-
curl -L -o server.py https://huggingface.co/smlflg/Connector/resolve/main/google-workspace/upstream/source/core/server.py
30.2 kB
| import asyncio | |
| import hashlib | |
| import logging | |
| import os | |
| from typing import List, Optional | |
| from importlib import metadata | |
| from fastapi.responses import HTMLResponse, JSONResponse, FileResponse | |
| from starlette.applications import Starlette | |
| from starlette.datastructures import MutableHeaders | |
| from starlette.types import Scope, Receive, Send | |
| from starlette.requests import Request | |
| from starlette.middleware import Middleware | |
| from fastmcp import FastMCP | |
| from fastmcp.server.auth.providers.google import GoogleProvider | |
| from auth.oauth21_session_store import get_oauth21_session_store, set_auth_provider | |
| from auth.google_auth import handle_auth_callback, start_auth_flow, check_client_secrets | |
| from auth.oauth_config import is_oauth21_enabled, is_external_oauth21_provider | |
| from auth.mcp_session_middleware import MCPSessionMiddleware | |
| from auth.oauth_responses import ( | |
| create_error_response, | |
| create_success_response, | |
| create_server_error_response, | |
| ) | |
| from auth.auth_info_middleware import AuthInfoMiddleware | |
| from auth.scopes import BASE_SCOPES, SCOPES, get_current_scopes # noqa | |
| from core.config import ( | |
| USER_GOOGLE_EMAIL, | |
| get_transport_mode, | |
| set_transport_mode as _set_transport_mode, | |
| get_oauth_redirect_uri as get_oauth_redirect_uri_for_current_mode, | |
| ) | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| _auth_provider: Optional[GoogleProvider] = None | |
| _legacy_callback_registered = False | |
| session_middleware = Middleware(MCPSessionMiddleware) | |
| class WellKnownCacheControlMiddleware: | |
| """Force no-cache headers for OAuth well-known discovery endpoints.""" | |
| def __init__(self, app): | |
| self.app = app | |
| async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: | |
| if scope["type"] != "http": | |
| await self.app(scope, receive, send) | |
| return | |
| path = scope.get("path", "") | |
| is_oauth_well_known = ( | |
| path == "/.well-known/oauth-authorization-server" | |
| or path.startswith("/.well-known/oauth-authorization-server/") | |
| or path == "/.well-known/oauth-protected-resource" | |
| or path.startswith("/.well-known/oauth-protected-resource/") | |
| ) | |
| if not is_oauth_well_known: | |
| await self.app(scope, receive, send) | |
| return | |
| async def send_with_no_cache_headers(message): | |
| if message["type"] == "http.response.start": | |
| headers = MutableHeaders(raw=message.setdefault("headers", [])) | |
| headers["Cache-Control"] = "no-store, must-revalidate" | |
| headers["ETag"] = f'"{_compute_scope_fingerprint()}"' | |
| await send(message) | |
| await self.app(scope, receive, send_with_no_cache_headers) | |
| well_known_cache_control_middleware = Middleware(WellKnownCacheControlMiddleware) | |
| def _compute_scope_fingerprint() -> str: | |
| """Compute a short hash of the current scope configuration for cache-busting.""" | |
| scopes_str = ",".join(sorted(get_current_scopes())) | |
| return hashlib.sha256(scopes_str.encode()).hexdigest()[:12] | |
| # Custom FastMCP that adds secure middleware stack for OAuth 2.1 | |
| class SecureFastMCP(FastMCP): | |
| def http_app(self, **kwargs) -> "Starlette": | |
| """Override to add secure middleware stack for OAuth 2.1.""" | |
| app = super().http_app(**kwargs) | |
| # Add middleware in order (first added = outermost layer) | |
| app.user_middleware.insert(0, well_known_cache_control_middleware) | |
| # Session Management - extracts session info for MCP context | |
| app.user_middleware.insert(1, session_middleware) | |
| # Rebuild middleware stack | |
| app.middleware_stack = app.build_middleware_stack() | |
| logger.info("Added middleware stack: WellKnownCacheControl, Session Management") | |
| return app | |
| async def list_tools(self, *, run_middleware: bool = True): | |
| """Override to mark user_google_email as optional when USER_GOOGLE_EMAIL is set. | |
| In single-user / self-hosted mode the env var provides the default email, so | |
| callers (agents, MCP adapters) should not be required to supply it. We patch | |
| the JSON schema returned by list_tools to remove 'user_google_email' from the | |
| ``required`` array and inject the env-var value as the ``default``. The | |
| runtime still resolves the email correctly via the service decorator. | |
| """ | |
| tools = list(await super().list_tools(run_middleware=run_middleware)) | |
| if not USER_GOOGLE_EMAIL or is_oauth21_enabled(): | |
| return tools | |
| patched = [] | |
| for tool in tools: | |
| schema = dict(tool.parameters) | |
| required = list(schema.get("required", [])) | |
| if "user_google_email" in required: | |
| required = [r for r in required if r != "user_google_email"] | |
| props = {k: dict(v) for k, v in schema.get("properties", {}).items()} | |
| if "user_google_email" in props: | |
| props["user_google_email"]["default"] = USER_GOOGLE_EMAIL | |
| schema = dict(schema, required=required, properties=props) | |
| patched.append(tool.model_copy(update={"parameters": schema})) | |
| else: | |
| patched.append(tool) | |
| return patched | |
| async def call_tool(self, name: str, arguments: Optional[dict], *args, **kwargs): | |
| """Inject user_google_email before pydantic validates the call arguments. | |
| When USER_GOOGLE_EMAIL is configured and OAuth 2.1 is not active, callers | |
| (agents, adapters) are allowed to omit user_google_email. FastMCP validates | |
| arguments against the function signature BEFORE calling the tool, so we must | |
| inject the default BEFORE that validation step. | |
| """ | |
| arguments = arguments or {} | |
| if ( | |
| not is_oauth21_enabled() | |
| and USER_GOOGLE_EMAIL | |
| and "user_google_email" not in arguments | |
| ): | |
| arguments = {**arguments, "user_google_email": USER_GOOGLE_EMAIL} | |
| return await super().call_tool(name, arguments, *args, **kwargs) | |
| # Build server instructions with user email context for single-user mode | |
| _server_instructions = None | |
| if USER_GOOGLE_EMAIL: | |
| _server_instructions = f"""Connected Google account: {USER_GOOGLE_EMAIL} | |
| When using Google Workspace tools, always use `{USER_GOOGLE_EMAIL}` as the `user_google_email` parameter. Do not ask the user for their email address.""" | |
| logger.info(f"Server instructions configured for user: {USER_GOOGLE_EMAIL}") | |
| server = SecureFastMCP( | |
| name="google_workspace", | |
| auth=None, | |
| instructions=_server_instructions, | |
| ) | |
| # Add the AuthInfo middleware to inject authentication into FastMCP context | |
| auth_info_middleware = AuthInfoMiddleware() | |
| server.add_middleware(auth_info_middleware) | |
| def _parse_bool_env(value: str) -> bool: | |
| """Parse environment variable string to boolean.""" | |
| return value.lower() in ("1", "true", "yes", "on") | |
| def _parse_allowed_redirect_uris(value: Optional[str]) -> Optional[List[str]]: | |
| """Parse a comma-separated list of OAuth client redirect URIs. | |
| Returns a list of non-empty, trimmed URIs, or None when the input is | |
| empty/None. Returning None preserves FastMCP's default behaviour of | |
| accepting any client-supplied redirect URI during DCR — callers that | |
| want to lock down registration must supply a non-empty list. | |
| Patterns supported by FastMCP's matcher (see | |
| ``fastmcp.server.auth.redirect_validation``) include ``*`` for port | |
| and path globs (e.g. ``http://localhost:*/callback``) and ``*.example.com`` | |
| for subdomain wildcards. | |
| """ | |
| if not value: | |
| return None | |
| uris = [u.strip() for u in value.split(",") if u.strip()] | |
| return uris or None | |
| def set_transport_mode(mode: str): | |
| """Sets the transport mode for the server.""" | |
| _set_transport_mode(mode) | |
| logger.info(f"Transport: {mode}") | |
| def _ensure_legacy_callback_route() -> None: | |
| global _legacy_callback_registered | |
| if _legacy_callback_registered: | |
| return | |
| server.custom_route("/oauth2callback", methods=["GET"])(legacy_oauth2_callback) | |
| _legacy_callback_registered = True | |
| def configure_server_for_http(): | |
| """ | |
| Configures the authentication provider for HTTP transport. | |
| This must be called BEFORE server.run(). | |
| """ | |
| global _auth_provider | |
| transport_mode = get_transport_mode() | |
| if transport_mode != "streamable-http": | |
| return | |
| # Use centralized OAuth configuration | |
| from auth.oauth_config import get_oauth_config | |
| config = get_oauth_config() | |
| # Check if OAuth 2.1 is enabled via centralized config | |
| oauth21_enabled = config.is_oauth21_enabled() | |
| if oauth21_enabled: | |
| if not config.is_configured(): | |
| logger.warning("OAuth 2.1 enabled but OAuth credentials not configured") | |
| return | |
| def validate_and_derive_jwt_key( | |
| jwt_signing_key_override: str | None, client_secret: str | None | |
| ) -> bytes: | |
| """Validate JWT signing key override and derive the final JWT key.""" | |
| if jwt_signing_key_override: | |
| if len(jwt_signing_key_override) < 12: | |
| logger.warning( | |
| "OAuth 2.1: FASTMCP_SERVER_AUTH_GOOGLE_JWT_SIGNING_KEY is less than 12 characters; " | |
| "use a longer secret to improve key derivation strength." | |
| ) | |
| return derive_jwt_key( | |
| low_entropy_material=jwt_signing_key_override, | |
| salt="fastmcp-jwt-signing-key", | |
| ) | |
| if client_secret: | |
| return derive_jwt_key( | |
| high_entropy_material=client_secret, | |
| salt="fastmcp-jwt-signing-key", | |
| ) | |
| raise ValueError( | |
| "Public client OAuth 2.1 requires FASTMCP_SERVER_AUTH_GOOGLE_JWT_SIGNING_KEY " | |
| "when GOOGLE_OAUTH_CLIENT_SECRET is not set." | |
| ) | |
| try: | |
| # Import common dependencies for storage backends | |
| from key_value.aio.wrappers.encryption import FernetEncryptionWrapper | |
| from cryptography.fernet import Fernet | |
| from fastmcp.server.auth.jwt_issuer import derive_jwt_key | |
| provider_valid_scopes: List[str] = sorted(get_current_scopes()) | |
| provider_required_scopes: List[str] = sorted(BASE_SCOPES) | |
| client_storage = None | |
| jwt_signing_key_override = ( | |
| os.getenv("FASTMCP_SERVER_AUTH_GOOGLE_JWT_SIGNING_KEY", "").strip() | |
| or None | |
| ) | |
| storage_backend = ( | |
| os.getenv("WORKSPACE_MCP_OAUTH_PROXY_STORAGE_BACKEND", "") | |
| .strip() | |
| .lower() | |
| ) | |
| valkey_host = os.getenv("WORKSPACE_MCP_OAUTH_PROXY_VALKEY_HOST", "").strip() | |
| # Determine storage backend: valkey, disk, memory (default) | |
| use_valkey = storage_backend == "valkey" or bool(valkey_host) | |
| use_disk = storage_backend == "disk" | |
| if use_valkey: | |
| try: | |
| from key_value.aio.stores.valkey import ValkeyStore | |
| valkey_port_raw = os.getenv( | |
| "WORKSPACE_MCP_OAUTH_PROXY_VALKEY_PORT", "6379" | |
| ).strip() | |
| valkey_db_raw = os.getenv( | |
| "WORKSPACE_MCP_OAUTH_PROXY_VALKEY_DB", "0" | |
| ).strip() | |
| valkey_port = int(valkey_port_raw) | |
| valkey_db = int(valkey_db_raw) | |
| valkey_use_tls_raw = os.getenv( | |
| "WORKSPACE_MCP_OAUTH_PROXY_VALKEY_USE_TLS", "" | |
| ).strip() | |
| valkey_use_tls = ( | |
| _parse_bool_env(valkey_use_tls_raw) | |
| if valkey_use_tls_raw | |
| else valkey_port == 6380 | |
| ) | |
| valkey_request_timeout_ms_raw = os.getenv( | |
| "WORKSPACE_MCP_OAUTH_PROXY_VALKEY_REQUEST_TIMEOUT_MS", "" | |
| ).strip() | |
| valkey_connection_timeout_ms_raw = os.getenv( | |
| "WORKSPACE_MCP_OAUTH_PROXY_VALKEY_CONNECTION_TIMEOUT_MS", "" | |
| ).strip() | |
| valkey_request_timeout_ms = ( | |
| int(valkey_request_timeout_ms_raw) | |
| if valkey_request_timeout_ms_raw | |
| else None | |
| ) | |
| valkey_connection_timeout_ms = ( | |
| int(valkey_connection_timeout_ms_raw) | |
| if valkey_connection_timeout_ms_raw | |
| else None | |
| ) | |
| valkey_username = ( | |
| os.getenv( | |
| "WORKSPACE_MCP_OAUTH_PROXY_VALKEY_USERNAME", "" | |
| ).strip() | |
| or None | |
| ) | |
| valkey_password = ( | |
| os.getenv( | |
| "WORKSPACE_MCP_OAUTH_PROXY_VALKEY_PASSWORD", "" | |
| ).strip() | |
| or None | |
| ) | |
| if not valkey_host: | |
| valkey_host = "localhost" | |
| client_storage = ValkeyStore( | |
| host=valkey_host, | |
| port=valkey_port, | |
| db=valkey_db, | |
| username=valkey_username, | |
| password=valkey_password, | |
| ) | |
| # Configure TLS and timeouts on the underlying Glide client config. | |
| # ValkeyStore currently doesn't expose these settings directly. | |
| glide_config = getattr(client_storage, "_client_config", None) | |
| if glide_config is not None: | |
| glide_config.use_tls = valkey_use_tls | |
| is_remote_host = valkey_host not in {"localhost", "127.0.0.1"} | |
| if valkey_request_timeout_ms is None and ( | |
| valkey_use_tls or is_remote_host | |
| ): | |
| # Glide defaults to 250ms if unset; increase for remote/TLS endpoints. | |
| valkey_request_timeout_ms = 5000 | |
| if valkey_request_timeout_ms is not None: | |
| glide_config.request_timeout = valkey_request_timeout_ms | |
| if valkey_connection_timeout_ms is None and ( | |
| valkey_use_tls or is_remote_host | |
| ): | |
| valkey_connection_timeout_ms = 10000 | |
| if valkey_connection_timeout_ms is not None: | |
| from glide_shared.config import ( | |
| AdvancedGlideClientConfiguration, | |
| ) | |
| glide_config.advanced_config = ( | |
| AdvancedGlideClientConfiguration( | |
| connection_timeout=valkey_connection_timeout_ms | |
| ) | |
| ) | |
| jwt_signing_key = validate_and_derive_jwt_key( | |
| jwt_signing_key_override, config.client_secret | |
| ) | |
| storage_encryption_key = derive_jwt_key( | |
| high_entropy_material=jwt_signing_key.decode(), | |
| salt="fastmcp-storage-encryption-key", | |
| ) | |
| client_storage = FernetEncryptionWrapper( | |
| key_value=client_storage, | |
| fernet=Fernet(key=storage_encryption_key), | |
| ) | |
| logger.info( | |
| "OAuth 2.1: Using ValkeyStore for FastMCP OAuth proxy client_storage (host=%s, port=%s, db=%s, tls=%s)", | |
| valkey_host, | |
| valkey_port, | |
| valkey_db, | |
| valkey_use_tls, | |
| ) | |
| if valkey_request_timeout_ms is not None: | |
| logger.info( | |
| "OAuth 2.1: Valkey request timeout set to %sms", | |
| valkey_request_timeout_ms, | |
| ) | |
| if valkey_connection_timeout_ms is not None: | |
| logger.info( | |
| "OAuth 2.1: Valkey connection timeout set to %sms", | |
| valkey_connection_timeout_ms, | |
| ) | |
| logger.info( | |
| "OAuth 2.1: Applied Fernet encryption wrapper to Valkey client_storage (key derived from FASTMCP_SERVER_AUTH_GOOGLE_JWT_SIGNING_KEY or GOOGLE_OAUTH_CLIENT_SECRET)." | |
| ) | |
| except ImportError as exc: | |
| logger.warning( | |
| "OAuth 2.1: Valkey client_storage requested but Valkey dependencies are not installed (%s). " | |
| "Install 'workspace-mcp[valkey]' (or 'py-key-value-aio[valkey]', which includes 'valkey-glide') " | |
| "or unset WORKSPACE_MCP_OAUTH_PROXY_STORAGE_BACKEND/WORKSPACE_MCP_OAUTH_PROXY_VALKEY_HOST.", | |
| exc, | |
| ) | |
| except ValueError as exc: | |
| logger.warning( | |
| "OAuth 2.1: Invalid Valkey configuration; falling back to default storage (%s).", | |
| exc, | |
| ) | |
| elif use_disk: | |
| try: | |
| from core.storage import make_sanitized_file_store | |
| disk_directory = os.getenv( | |
| "WORKSPACE_MCP_OAUTH_PROXY_DISK_DIRECTORY", "" | |
| ).strip() | |
| if not disk_directory: | |
| # Default to FASTMCP_HOME/oauth-proxy or ~/.fastmcp/oauth-proxy | |
| fastmcp_home = os.getenv("FASTMCP_HOME", "").strip() | |
| if fastmcp_home: | |
| disk_directory = os.path.join(fastmcp_home, "oauth-proxy") | |
| else: | |
| disk_directory = os.path.expanduser( | |
| "~/.fastmcp/oauth-proxy" | |
| ) | |
| client_storage = make_sanitized_file_store(disk_directory) | |
| jwt_signing_key = validate_and_derive_jwt_key( | |
| jwt_signing_key_override, config.client_secret | |
| ) | |
| storage_encryption_key = derive_jwt_key( | |
| high_entropy_material=jwt_signing_key.decode(), | |
| salt="fastmcp-storage-encryption-key", | |
| ) | |
| client_storage = FernetEncryptionWrapper( | |
| key_value=client_storage, | |
| fernet=Fernet(key=storage_encryption_key), | |
| ) | |
| logger.info( | |
| "OAuth 2.1: Using FileTreeStore for FastMCP OAuth proxy client_storage (directory=%s)", | |
| disk_directory, | |
| ) | |
| except ImportError as exc: | |
| logger.warning( | |
| "OAuth 2.1: Disk storage requested but dependencies not available (%s). " | |
| "Falling back to default storage.", | |
| exc, | |
| ) | |
| elif storage_backend == "memory": | |
| from key_value.aio.stores.memory import MemoryStore | |
| client_storage = MemoryStore() | |
| logger.info( | |
| "OAuth 2.1: Using MemoryStore for FastMCP OAuth proxy client_storage" | |
| ) | |
| # else: client_storage remains None, FastMCP uses its default | |
| # Ensure JWT signing key is always derived for all storage backends | |
| if "jwt_signing_key" not in locals(): | |
| jwt_signing_key = validate_and_derive_jwt_key( | |
| jwt_signing_key_override, config.client_secret | |
| ) | |
| # Check if external OAuth provider is configured | |
| if config.is_external_oauth21_provider(): | |
| # External OAuth mode: use custom provider that handles ya29.* access tokens | |
| from auth.external_oauth_provider import ExternalOAuthProvider | |
| provider = ExternalOAuthProvider( | |
| client_id=config.client_id, | |
| client_secret=config.client_secret, | |
| base_url=config.get_oauth_base_url(), | |
| redirect_path=config.redirect_path, | |
| required_scopes=provider_valid_scopes, | |
| resource_server_url=config.get_oauth_base_url(), | |
| ) | |
| server.auth = provider | |
| logger.info("OAuth 2.1 enabled with EXTERNAL provider mode") | |
| logger.info( | |
| "Expecting Authorization bearer tokens in tool call headers" | |
| ) | |
| logger.info( | |
| "Protected resource metadata points to Google's authorization server" | |
| ) | |
| else: | |
| # Standard OAuth 2.1 mode: use FastMCP's GoogleProvider | |
| allowed_client_redirect_uris = _parse_allowed_redirect_uris( | |
| os.getenv("WORKSPACE_MCP_ALLOWED_CLIENT_REDIRECT_URIS") | |
| ) | |
| if allowed_client_redirect_uris: | |
| logger.info( | |
| "OAuth 2.1: restricting DCR client redirect URIs to allowlist: %s", | |
| allowed_client_redirect_uris, | |
| ) | |
| provider = GoogleProvider( | |
| client_id=config.client_id, | |
| client_secret=config.client_secret, | |
| base_url=config.get_oauth_base_url(), | |
| redirect_path=config.redirect_path, | |
| required_scopes=provider_required_scopes, | |
| valid_scopes=provider_valid_scopes, | |
| client_storage=client_storage, | |
| jwt_signing_key=jwt_signing_key, | |
| allowed_client_redirect_uris=allowed_client_redirect_uris, | |
| ) | |
| if provider.client_registration_options is not None: | |
| # Keep protocol-level auth limited to base identity scopes, but | |
| # allow dynamically registered MCP clients to request any scope | |
| # needed by enabled tools during subsequent authorization flows. | |
| provider.client_registration_options.default_scopes = ( | |
| provider_valid_scopes | |
| ) | |
| # Enable protocol-level auth | |
| server.auth = provider | |
| logger.info( | |
| "OAuth 2.1 enabled using FastMCP GoogleProvider with protocol-level auth" | |
| ) | |
| # Always set auth provider for token validation in middleware | |
| set_auth_provider(provider) | |
| _auth_provider = provider | |
| except Exception as exc: | |
| logger.error( | |
| "Failed to initialize FastMCP GoogleProvider: %s", exc, exc_info=True | |
| ) | |
| raise | |
| else: | |
| logger.info("OAuth 2.0 mode - Server will use legacy authentication.") | |
| server.auth = None | |
| _auth_provider = None | |
| set_auth_provider(None) | |
| _ensure_legacy_callback_route() | |
| def get_auth_provider() -> Optional[GoogleProvider]: | |
| """Gets the global authentication provider instance.""" | |
| return _auth_provider | |
| async def health_check(request: Request): | |
| try: | |
| version = metadata.version("workspace-mcp") | |
| except metadata.PackageNotFoundError: | |
| version = "dev" | |
| return JSONResponse( | |
| { | |
| "status": "healthy", | |
| "service": "workspace-mcp", | |
| "version": version, | |
| "transport": get_transport_mode(), | |
| } | |
| ) | |
| async def serve_attachment(request: Request): | |
| """Serve a stored attachment file.""" | |
| from core.attachment_storage import get_attachment_storage | |
| file_id = request.path_params["file_id"] | |
| storage = get_attachment_storage() | |
| metadata = storage.get_attachment_metadata(file_id) | |
| if not metadata: | |
| return JSONResponse( | |
| {"error": "Attachment not found or expired"}, status_code=404 | |
| ) | |
| file_path = storage.get_attachment_path(file_id) | |
| if not file_path: | |
| return JSONResponse({"error": "Attachment file not found"}, status_code=404) | |
| return FileResponse( | |
| path=str(file_path), | |
| filename=metadata["filename"], | |
| media_type=metadata["mime_type"], | |
| ) | |
| async def legacy_oauth2_callback(request: Request) -> HTMLResponse: | |
| state = request.query_params.get("state") | |
| code = request.query_params.get("code") | |
| error = request.query_params.get("error") | |
| if error: | |
| msg = ( | |
| f"Authentication failed: Google returned an error: {error}. State: {state}." | |
| ) | |
| logger.error(msg) | |
| return create_error_response(msg) | |
| if not code: | |
| msg = "Authentication failed: No authorization code received from Google." | |
| logger.error(msg) | |
| return create_error_response(msg) | |
| try: | |
| error_message = check_client_secrets() | |
| if error_message: | |
| return create_server_error_response(error_message) | |
| logger.info("OAuth callback: Received authorization code.") | |
| mcp_session_id = None | |
| if hasattr(request, "state") and hasattr(request.state, "session_id"): | |
| mcp_session_id = request.state.session_id | |
| verified_user_id, credentials = await handle_auth_callback( | |
| scopes=get_current_scopes(), | |
| authorization_response=str(request.url), | |
| redirect_uri=get_oauth_redirect_uri_for_current_mode(), | |
| session_id=mcp_session_id, | |
| ) | |
| logger.info( | |
| f"OAuth callback: Successfully authenticated user: {verified_user_id}." | |
| ) | |
| try: | |
| store = get_oauth21_session_store() | |
| store.store_session( | |
| user_email=verified_user_id, | |
| access_token=credentials.token, | |
| refresh_token=credentials.refresh_token, | |
| token_uri=credentials.token_uri, | |
| client_id=credentials.client_id, | |
| client_secret=credentials.client_secret, | |
| scopes=credentials.scopes, | |
| expiry=credentials.expiry, | |
| session_id=f"google-{state}", | |
| mcp_session_id=mcp_session_id, | |
| ) | |
| logger.info( | |
| f"Stored Google credentials in OAuth 2.1 session store for {verified_user_id}" | |
| ) | |
| except Exception as e: | |
| logger.error(f"Failed to store credentials in OAuth 2.1 store: {e}") | |
| return create_success_response(verified_user_id) | |
| except Exception as e: | |
| logger.error(f"Error processing OAuth callback: {str(e)}", exc_info=True) | |
| return create_server_error_response(str(e)) | |
| async def start_google_auth( | |
| service_name: str, user_google_email: str = USER_GOOGLE_EMAIL | |
| ) -> str: | |
| """ | |
| Manually initiate Google OAuth authentication flow. | |
| NOTE: This is a legacy OAuth 2.0 tool and is disabled when OAuth 2.1 is enabled. | |
| The authentication system automatically handles credential checks and prompts for | |
| authentication when needed. Only use this tool if: | |
| 1. You need to re-authenticate with different credentials | |
| 2. You want to proactively authenticate before using other tools | |
| 3. The automatic authentication flow failed and you need to retry | |
| In most cases, simply try calling the Google Workspace tool you need - it will | |
| automatically handle authentication if required. | |
| """ | |
| if is_oauth21_enabled(): | |
| if is_external_oauth21_provider(): | |
| return ( | |
| "start_google_auth is disabled when OAuth 2.1 is enabled. " | |
| "Provide a valid OAuth 2.1 bearer token in the Authorization header " | |
| "and retry the original tool." | |
| ) | |
| return ( | |
| "start_google_auth is disabled when OAuth 2.1 is enabled. " | |
| "Authenticate through your MCP client's OAuth 2.1 flow and retry the " | |
| "original tool." | |
| ) | |
| if not user_google_email: | |
| raise ValueError("user_google_email must be provided.") | |
| error_message = check_client_secrets() | |
| if error_message: | |
| return f"**Authentication Error:** {error_message}" | |
| try: | |
| transport_mode = get_transport_mode() | |
| if transport_mode == "stdio": | |
| # Only stdio legacy OAuth depends on the standalone callback server. | |
| from auth.oauth_callback_server import ensure_oauth_callback_available | |
| from auth.oauth_config import get_oauth_config | |
| config = get_oauth_config() | |
| success, error_msg = await asyncio.to_thread( | |
| ensure_oauth_callback_available, | |
| transport_mode, | |
| config.port, | |
| config.base_uri, | |
| ) | |
| if not success: | |
| error_detail = f" ({error_msg})" if error_msg else "" | |
| return f"**Error:** Cannot initiate OAuth flow - callback server unavailable{error_detail}" | |
| auth_message = await start_auth_flow( | |
| user_google_email=user_google_email, | |
| service_name=service_name, | |
| redirect_uri=get_oauth_redirect_uri_for_current_mode(), | |
| ) | |
| return auth_message | |
| except Exception as e: | |
| logger.error(f"Failed to start Google authentication flow: {e}", exc_info=True) | |
| return f"**Error:** An unexpected error occurred: {e}" | |