workspace-mcp 0.2.0__py3-none-any.whl

auth/service_decorator.py

@@ -0,0 +1,404 @@
+import inspect
+import logging
+from functools import wraps
+from typing import Dict, List, Optional, Any, Callable, Union
+from datetime import datetime, timedelta
+
+from google.auth.exceptions import RefreshError
+from auth.google_auth import get_authenticated_google_service, GoogleAuthenticationError
+
+logger = logging.getLogger(__name__)
+
+# Import scope constants
+from auth.scopes import (
+    GMAIL_READONLY_SCOPE, GMAIL_SEND_SCOPE, GMAIL_COMPOSE_SCOPE, GMAIL_MODIFY_SCOPE, GMAIL_LABELS_SCOPE,
+    DRIVE_READONLY_SCOPE, DRIVE_FILE_SCOPE,
+    DOCS_READONLY_SCOPE, DOCS_WRITE_SCOPE,
+    CALENDAR_READONLY_SCOPE, CALENDAR_EVENTS_SCOPE,
+    SHEETS_READONLY_SCOPE, SHEETS_WRITE_SCOPE,
+    CHAT_READONLY_SCOPE, CHAT_WRITE_SCOPE, CHAT_SPACES_SCOPE,
+    FORMS_BODY_SCOPE, FORMS_BODY_READONLY_SCOPE, FORMS_RESPONSES_READONLY_SCOPE,
+    SLIDES_SCOPE, SLIDES_READONLY_SCOPE
+)
+
+# Service configuration mapping
+SERVICE_CONFIGS = {
+    "gmail": {"service": "gmail", "version": "v1"},
+    "drive": {"service": "drive", "version": "v3"},
+    "calendar": {"service": "calendar", "version": "v3"},
+    "docs": {"service": "docs", "version": "v1"},
+    "sheets": {"service": "sheets", "version": "v4"},
+    "chat": {"service": "chat", "version": "v1"},
+    "forms": {"service": "forms", "version": "v1"},
+    "slides": {"service": "slides", "version": "v1"}
+}
+
+
+# Scope group definitions for easy reference
+SCOPE_GROUPS = {
+    # Gmail scopes
+    "gmail_read": GMAIL_READONLY_SCOPE,
+    "gmail_send": GMAIL_SEND_SCOPE,
+    "gmail_compose": GMAIL_COMPOSE_SCOPE,
+    "gmail_modify": GMAIL_MODIFY_SCOPE,
+    "gmail_labels": GMAIL_LABELS_SCOPE,
+
+    # Drive scopes
+    "drive_read": DRIVE_READONLY_SCOPE,
+    "drive_file": DRIVE_FILE_SCOPE,
+
+    # Docs scopes
+    "docs_read": DOCS_READONLY_SCOPE,
+    "docs_write": DOCS_WRITE_SCOPE,
+
+    # Calendar scopes
+    "calendar_read": CALENDAR_READONLY_SCOPE,
+    "calendar_events": CALENDAR_EVENTS_SCOPE,
+
+    # Sheets scopes
+    "sheets_read": SHEETS_READONLY_SCOPE,
+    "sheets_write": SHEETS_WRITE_SCOPE,
+
+    # Chat scopes
+    "chat_read": CHAT_READONLY_SCOPE,
+    "chat_write": CHAT_WRITE_SCOPE,
+    "chat_spaces": CHAT_SPACES_SCOPE,
+
+    # Forms scopes
+    "forms": FORMS_BODY_SCOPE,
+    "forms_read": FORMS_BODY_READONLY_SCOPE,
+    "forms_responses_read": FORMS_RESPONSES_READONLY_SCOPE,
+
+    # Slides scopes
+    "slides": SLIDES_SCOPE,
+    "slides_read": SLIDES_READONLY_SCOPE,
+}
+
+# Service cache: {cache_key: (service, cached_time, user_email)}
+_service_cache: Dict[str, tuple[Any, datetime, str]] = {}
+_cache_ttl = timedelta(minutes=30)  # Cache services for 30 minutes
+
+
+def _get_cache_key(user_email: str, service_name: str, version: str, scopes: List[str]) -> str:
+    """Generate a cache key for service instances."""
+    sorted_scopes = sorted(scopes)
+    return f"{user_email}:{service_name}:{version}:{':'.join(sorted_scopes)}"
+
+
+def _is_cache_valid(cached_time: datetime) -> bool:
+    """Check if cached service is still valid."""
+    return datetime.now() - cached_time < _cache_ttl
+
+
+def _get_cached_service(cache_key: str) -> Optional[tuple[Any, str]]:
+    """Retrieve cached service if valid."""
+    if cache_key in _service_cache:
+        service, cached_time, user_email = _service_cache[cache_key]
+        if _is_cache_valid(cached_time):
+            logger.debug(f"Using cached service for key: {cache_key}")
+            return service, user_email
+        else:
+            # Remove expired cache entry
+            del _service_cache[cache_key]
+            logger.debug(f"Removed expired cache entry: {cache_key}")
+    return None
+
+
+def _cache_service(cache_key: str, service: Any, user_email: str) -> None:
+    """Cache a service instance."""
+    _service_cache[cache_key] = (service, datetime.now(), user_email)
+    logger.debug(f"Cached service for key: {cache_key}")
+
+
+def _resolve_scopes(scopes: Union[str, List[str]]) -> List[str]:
+    """Resolve scope names to actual scope URLs."""
+    if isinstance(scopes, str):
+        if scopes in SCOPE_GROUPS:
+            return [SCOPE_GROUPS[scopes]]
+        else:
+            return [scopes]
+
+    resolved = []
+    for scope in scopes:
+        if scope in SCOPE_GROUPS:
+            resolved.append(SCOPE_GROUPS[scope])
+        else:
+            resolved.append(scope)
+    return resolved
+
+
+def _handle_token_refresh_error(error: RefreshError, user_email: str, service_name: str) -> str:
+    """
+    Handle token refresh errors gracefully, particularly expired/revoked tokens.
+
+    Args:
+        error: The RefreshError that occurred
+        user_email: User's email address
+        service_name: Name of the Google service
+
+    Returns:
+        A user-friendly error message with instructions for reauthentication
+    """
+    error_str = str(error)
+
+    if 'invalid_grant' in error_str.lower() or 'expired or revoked' in error_str.lower():
+        logger.warning(f"Token expired or revoked for user {user_email} accessing {service_name}")
+
+        # Clear any cached service for this user to force fresh authentication
+        clear_service_cache(user_email)
+
+        service_display_name = f"Google {service_name.title()}"
+
+        return (
+            f"**Authentication Required: Token Expired/Revoked for {service_display_name}**\n\n"
+            f"Your Google authentication token for {user_email} has expired or been revoked. "
+            f"This commonly happens when:\n"
+            f"- The token has been unused for an extended period\n"
+            f"- You've changed your Google account password\n"
+            f"- You've revoked access to the application\n\n"
+            f"**To resolve this, please:**\n"
+            f"1. Run `start_google_auth` with your email ({user_email}) and service_name='{service_display_name}'\n"
+            f"2. Complete the authentication flow in your browser\n"
+            f"3. Retry your original command\n\n"
+            f"The application will automatically use the new credentials once authentication is complete."
+        )
+    else:
+        # Handle other types of refresh errors
+        logger.error(f"Unexpected refresh error for user {user_email}: {error}")
+        return (
+            f"Authentication error occurred for {user_email}. "
+            f"Please try running `start_google_auth` with your email and the appropriate service name to reauthenticate."
+        )
+
+
+def require_google_service(
+    service_type: str,
+    scopes: Union[str, List[str]],
+    version: Optional[str] = None,
+    cache_enabled: bool = True
+):
+    """
+    Decorator that automatically handles Google service authentication and injection.
+
+    Args:
+        service_type: Type of Google service ("gmail", "drive", "calendar", etc.)
+        scopes: Required scopes (can be scope group names or actual URLs)
+        version: Service version (defaults to standard version for service type)
+        cache_enabled: Whether to use service caching (default: True)
+
+    Usage:
+        @require_google_service("gmail", "gmail_read")
+        async def search_messages(service, user_google_email: str, query: str):
+            # service parameter is automatically injected
+            # Original authentication logic is handled automatically
+    """
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            # Extract user_google_email from function parameters
+            sig = inspect.signature(func)
+            param_names = list(sig.parameters.keys())
+
+            # Find user_google_email parameter
+            user_google_email = None
+            if 'user_google_email' in kwargs:
+                user_google_email = kwargs['user_google_email']
+            else:
+                # Look for user_google_email in positional args
+                try:
+                    user_email_index = param_names.index('user_google_email')
+                    if user_email_index < len(args):
+                        user_google_email = args[user_email_index]
+                except ValueError:
+                    pass
+
+            if not user_google_email:
+                raise Exception("user_google_email parameter is required but not found")
+
+            # Get service configuration
+            if service_type not in SERVICE_CONFIGS:
+                raise Exception(f"Unknown service type: {service_type}")
+
+            config = SERVICE_CONFIGS[service_type]
+            service_name = config["service"]
+            service_version = version or config["version"]
+
+            # Resolve scopes
+            resolved_scopes = _resolve_scopes(scopes)
+
+            # Check cache first if enabled
+            service = None
+            actual_user_email = user_google_email
+
+            if cache_enabled:
+                cache_key = _get_cache_key(user_google_email, service_name, service_version, resolved_scopes)
+                cached_result = _get_cached_service(cache_key)
+                if cached_result:
+                    service, actual_user_email = cached_result
+
+            # If not cached, authenticate
+            if service is None:
+                try:
+                    tool_name = func.__name__
+                    service, actual_user_email = await get_authenticated_google_service(
+                        service_name=service_name,
+                        version=service_version,
+                        tool_name=tool_name,
+                        user_google_email=user_google_email,
+                        required_scopes=resolved_scopes,
+                    )
+
+                    # Cache the service if caching is enabled
+                    if cache_enabled:
+                        cache_key = _get_cache_key(user_google_email, service_name, service_version, resolved_scopes)
+                        _cache_service(cache_key, service, actual_user_email)
+
+                except GoogleAuthenticationError as e:
+                    raise Exception(str(e))
+
+            # Inject service as first parameter
+            if 'service' in param_names:
+                kwargs['service'] = service
+            else:
+                # Insert service as first positional argument
+                args = (service,) + args
+
+            # Call the original function with refresh error handling
+            try:
+                return await func(*args, **kwargs)
+            except RefreshError as e:
+                # Handle token refresh errors gracefully
+                error_message = _handle_token_refresh_error(e, actual_user_email, service_name)
+                raise Exception(error_message)
+
+        return wrapper
+    return decorator
+
+
+def require_multiple_services(service_configs: List[Dict[str, Any]]):
+    """
+    Decorator for functions that need multiple Google services.
+
+    Args:
+        service_configs: List of service configurations, each containing:
+            - service_type: Type of service
+            - scopes: Required scopes
+            - param_name: Name to inject service as (e.g., 'drive_service', 'docs_service')
+            - version: Optional version override
+
+    Usage:
+        @require_multiple_services([
+            {"service_type": "drive", "scopes": "drive_read", "param_name": "drive_service"},
+            {"service_type": "docs", "scopes": "docs_read", "param_name": "docs_service"}
+        ])
+        async def get_doc_with_metadata(drive_service, docs_service, user_google_email: str, doc_id: str):
+            # Both services are automatically injected
+    """
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            # Extract user_google_email
+            sig = inspect.signature(func)
+            param_names = list(sig.parameters.keys())
+
+            user_google_email = None
+            if 'user_google_email' in kwargs:
+                user_google_email = kwargs['user_google_email']
+            else:
+                try:
+                    user_email_index = param_names.index('user_google_email')
+                    if user_email_index < len(args):
+                        user_google_email = args[user_email_index]
+                except ValueError:
+                    pass
+
+            if not user_google_email:
+                raise Exception("user_google_email parameter is required but not found")
+
+            # Authenticate all services
+            for config in service_configs:
+                service_type = config["service_type"]
+                scopes = config["scopes"]
+                param_name = config["param_name"]
+                version = config.get("version")
+
+                if service_type not in SERVICE_CONFIGS:
+                    raise Exception(f"Unknown service type: {service_type}")
+
+                service_config = SERVICE_CONFIGS[service_type]
+                service_name = service_config["service"]
+                service_version = version or service_config["version"]
+                resolved_scopes = _resolve_scopes(scopes)
+
+                try:
+                    tool_name = func.__name__
+                    service, _ = await get_authenticated_google_service(
+                        service_name=service_name,
+                        version=service_version,
+                        tool_name=tool_name,
+                        user_google_email=user_google_email,
+                        required_scopes=resolved_scopes,
+                    )
+
+                    # Inject service with specified parameter name
+                    kwargs[param_name] = service
+
+                except GoogleAuthenticationError as e:
+                    raise Exception(str(e))
+
+            # Call the original function with refresh error handling
+            try:
+                return await func(*args, **kwargs)
+            except RefreshError as e:
+                # Handle token refresh errors gracefully
+                error_message = _handle_token_refresh_error(e, user_google_email, "Multiple Services")
+                raise Exception(error_message)
+
+        return wrapper
+    return decorator
+
+
+def clear_service_cache(user_email: Optional[str] = None) -> int:
+    """
+    Clear service cache entries.
+
+    Args:
+        user_email: If provided, only clear cache for this user. If None, clear all.
+
+    Returns:
+        Number of cache entries cleared.
+    """
+    global _service_cache
+
+    if user_email is None:
+        count = len(_service_cache)
+        _service_cache.clear()
+        logger.info(f"Cleared all {count} service cache entries")
+        return count
+
+    keys_to_remove = [key for key in _service_cache.keys() if key.startswith(f"{user_email}:")]
+    for key in keys_to_remove:
+        del _service_cache[key]
+
+    logger.info(f"Cleared {len(keys_to_remove)} service cache entries for user {user_email}")
+    return len(keys_to_remove)
+
+
+def get_cache_stats() -> Dict[str, Any]:
+    """Get service cache statistics."""
+    now = datetime.now()
+    valid_entries = 0
+    expired_entries = 0
+
+    for _, (_, cached_time, _) in _service_cache.items():
+        if _is_cache_valid(cached_time):
+            valid_entries += 1
+        else:
+            expired_entries += 1
+
+    return {
+        "total_entries": len(_service_cache),
+        "valid_entries": valid_entries,
+        "expired_entries": expired_entries,
+        "cache_ttl_minutes": _cache_ttl.total_seconds() / 60
+    }

core/__init__.py

@@ -0,0 +1 @@
+# Make the core directory a Python package

core/server.py

@@ -0,0 +1,214 @@
+import logging
+import os
+from typing import Dict, Any, Optional
+
+from fastapi import Header
+from fastapi.responses import HTMLResponse
+
+from mcp import types
+
+from mcp.server.fastmcp import FastMCP
+from starlette.requests import Request
+
+from auth.google_auth import handle_auth_callback, start_auth_flow, CONFIG_CLIENT_SECRETS_PATH
+from auth.oauth_callback_server import get_oauth_redirect_uri, ensure_oauth_callback_available
+from auth.oauth_responses import create_error_response, create_success_response, create_server_error_response
+
+# Import shared configuration
+from auth.scopes import (
+    OAUTH_STATE_TO_SESSION_ID_MAP,
+    USERINFO_EMAIL_SCOPE,
+    OPENID_SCOPE,
+    CALENDAR_READONLY_SCOPE,
+    CALENDAR_EVENTS_SCOPE,
+    DRIVE_READONLY_SCOPE,
+    DRIVE_FILE_SCOPE,
+    GMAIL_READONLY_SCOPE,
+    GMAIL_SEND_SCOPE,
+    GMAIL_COMPOSE_SCOPE,
+    GMAIL_MODIFY_SCOPE,
+    GMAIL_LABELS_SCOPE,
+    BASE_SCOPES,
+    CALENDAR_SCOPES,
+    DRIVE_SCOPES,
+    GMAIL_SCOPES,
+    DOCS_READONLY_SCOPE,
+    DOCS_WRITE_SCOPE,
+    CHAT_READONLY_SCOPE,
+    CHAT_WRITE_SCOPE,
+    CHAT_SPACES_SCOPE,
+    CHAT_SCOPES,
+    SHEETS_READONLY_SCOPE,
+    SHEETS_WRITE_SCOPE,
+    SHEETS_SCOPES,
+    FORMS_BODY_SCOPE,
+    FORMS_BODY_READONLY_SCOPE,
+    FORMS_RESPONSES_READONLY_SCOPE,
+    FORMS_SCOPES,
+    SLIDES_SCOPE,
+    SLIDES_READONLY_SCOPE,
+    SLIDES_SCOPES,
+    SCOPES
+)
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+WORKSPACE_MCP_PORT = int(os.getenv("PORT", os.getenv("WORKSPACE_MCP_PORT", 8000)))
+WORKSPACE_MCP_BASE_URI = os.getenv("WORKSPACE_MCP_BASE_URI", "http://localhost")
+
+# Transport mode detection (will be set by main.py)
+_current_transport_mode = "stdio"  # Default to stdio
+
+# Basic MCP server instance
+server = FastMCP(
+    name="google_workspace",
+    server_url=f"{WORKSPACE_MCP_BASE_URI}:{WORKSPACE_MCP_PORT}/mcp",
+    port=WORKSPACE_MCP_PORT,
+    host="0.0.0.0"
+)
+
+def set_transport_mode(mode: str):
+    """Set the current transport mode for OAuth callback handling."""
+    global _current_transport_mode
+    _current_transport_mode = mode
+    logger.info(f"Transport mode set to: {mode}")
+
+def get_oauth_redirect_uri_for_current_mode() -> str:
+    """Get OAuth redirect URI based on current transport mode."""
+    return get_oauth_redirect_uri(_current_transport_mode, WORKSPACE_MCP_PORT, WORKSPACE_MCP_BASE_URI)
+
+# Health check endpoint
+@server.custom_route("/health", methods=["GET"])
+async def health_check(request: Request):
+    """Health check endpoint for container orchestration."""
+    from fastapi.responses import JSONResponse
+    return JSONResponse({
+        "status": "healthy",
+        "service": "google-workspace-mcp", 
+        "version": "0.1.1",
+        "transport": _current_transport_mode
+    })
+
+
+# Register OAuth callback as a custom route
+@server.custom_route("/oauth2callback", methods=["GET"])
+async def oauth2_callback(request: Request) -> HTMLResponse:
+    """
+    Handle OAuth2 callback from Google via a custom route.
+    This endpoint exchanges the authorization code for credentials and saves them.
+    It then displays a success or error page to the user.
+    """
+    # State is used by google-auth-library for CSRF protection and should be present.
+    # We don't need to track it ourselves in this simplified flow.
+    state = request.query_params.get("state")
+    code = request.query_params.get("code")
+    error = request.query_params.get("error")
+
+    if error:
+        error_message = f"Authentication failed: Google returned an error: {error}. State: {state}."
+        logger.error(error_message)
+        return create_error_response(error_message)
+
+    if not code:
+        error_message = "Authentication failed: No authorization code received from Google."
+        logger.error(error_message)
+        return create_error_response(error_message)
+
+    try:
+        # Use the centralized CONFIG_CLIENT_SECRETS_PATH
+        client_secrets_path = CONFIG_CLIENT_SECRETS_PATH
+        if not os.path.exists(client_secrets_path):
+            logger.error(f"OAuth client secrets file not found at {client_secrets_path}")
+            # This is a server configuration error, should not happen in a deployed environment.
+            return HTMLResponse(content="Server Configuration Error: Client secrets not found.", status_code=500)
+
+        logger.info(f"OAuth callback: Received code (state: {state}). Attempting to exchange for tokens.")
+
+        mcp_session_id: Optional[str] = OAUTH_STATE_TO_SESSION_ID_MAP.pop(state, None)
+        if mcp_session_id:
+            logger.info(f"OAuth callback: Retrieved MCP session ID '{mcp_session_id}' for state '{state}'.")
+        else:
+            logger.warning(f"OAuth callback: No MCP session ID found for state '{state}'. Auth will not be tied to a specific session directly via this callback.")
+
+        # Exchange code for credentials. handle_auth_callback will save them.
+        # The user_id returned here is the Google-verified email.
+        verified_user_id, credentials = handle_auth_callback(
+            client_secrets_path=client_secrets_path,
+            scopes=SCOPES, # Ensure all necessary scopes are requested
+            authorization_response=str(request.url),
+            redirect_uri=get_oauth_redirect_uri_for_current_mode(),
+            session_id=mcp_session_id # Pass session_id if available
+        )
+
+        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 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)
+        # Generic error page for any other issues during token exchange or credential saving
+        return create_server_error_response(str(e))
+
+@server.tool()
+async def start_google_auth(
+    user_google_email: str,
+    service_name: str,
+    mcp_session_id: Optional[str] = Header(None, alias="Mcp-Session-Id")
+) -> str:
+    """
+    Initiates the Google OAuth 2.0 authentication flow for the specified user email and service.
+    This is the primary method to establish credentials when no valid session exists or when targeting a specific account for a particular service.
+    It generates an authorization URL that the LLM must present to the user.
+    The authentication attempt is linked to the current MCP session via `mcp_session_id`.
+
+    LLM Guidance:
+    - Use this tool when you need to authenticate a user for a specific Google service (e.g., "Google Calendar", "Google Docs", "Gmail", "Google Drive")
+      and don't have existing valid credentials for the session or specified email.
+    - You MUST provide the `user_google_email` and the `service_name`. If you don't know the email, ask the user first.
+    - Valid `service_name` values typically include "Google Calendar", "Google Docs", "Gmail", "Google Drive".
+    - After calling this tool, present the returned authorization URL clearly to the user and instruct them to:
+        1. Click the link and complete the sign-in/consent process in their browser.
+        2. Note the authenticated email displayed on the success page.
+        3. Provide that email back to you (the LLM).
+        4. Retry their original request, including the confirmed `user_google_email`.
+
+    Args:
+        user_google_email (str): The user's full Google email address (e.g., 'example@gmail.com'). This is REQUIRED.
+        service_name (str): The name of the Google service for which authentication is being requested (e.g., "Google Calendar", "Google Docs"). This is REQUIRED.
+        mcp_session_id (Optional[str]): The active MCP session ID (automatically injected by FastMCP from the Mcp-Session-Id header). Links the OAuth flow state to the session.
+
+    Returns:
+        str: A detailed message for the LLM with the authorization URL and instructions to guide the user through the authentication process.
+    """
+    if not user_google_email or not isinstance(user_google_email, str) or '@' not in user_google_email:
+        error_msg = "Invalid or missing 'user_google_email'. This parameter is required and must be a valid email address. LLM, please ask the user for their Google email address."
+        logger.error(f"[start_google_auth] {error_msg}")
+        raise Exception(error_msg)
+
+    if not service_name or not isinstance(service_name, str):
+        error_msg = "Invalid or missing 'service_name'. This parameter is required (e.g., 'Google Calendar', 'Google Docs'). LLM, please specify the service name."
+        logger.error(f"[start_google_auth] {error_msg}")
+        raise Exception(error_msg)
+
+    logger.info(f"Tool 'start_google_auth' invoked for user_google_email: '{user_google_email}', service: '{service_name}', session: '{mcp_session_id}'.")
+
+    # Ensure OAuth callback is available for current transport mode
+    redirect_uri = get_oauth_redirect_uri_for_current_mode()
+    if not ensure_oauth_callback_available(_current_transport_mode, WORKSPACE_MCP_PORT, WORKSPACE_MCP_BASE_URI):
+        raise Exception("Failed to start OAuth callback server. Please try again.")
+
+    # Use the centralized start_auth_flow from auth.google_auth
+    auth_result = await start_auth_flow(
+        mcp_session_id=mcp_session_id,
+        user_google_email=user_google_email,
+        service_name=service_name,
+        redirect_uri=redirect_uri
+    )
+
+    # auth_result is now a plain string, not a CallToolResult
+    return auth_result