Initial public upload from Projekte/Connector (part 16)
bf06a0d verified about 2 months ago
56.1 kB
# auth/google_auth.py

import asyncio
import json
import jwt
import logging
import os

from typing import List, Optional, Tuple, Dict, Any
from urllib.parse import parse_qs, urlparse

from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import Flow
from google.auth.transport.requests import Request
from google.auth.exceptions import RefreshError
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from auth.scopes import SCOPES, get_current_scopes, has_required_scopes  # noqa
from auth.oauth21_session_store import get_oauth21_session_store
from auth.credential_store import get_credential_store
from auth.oauth_config import get_oauth_config, is_stateless_mode
from core.config import (
    get_transport_mode,
    get_oauth_redirect_uri,
)
from core.context import get_fastmcp_session_id

# Try to import FastMCP dependencies (may not be available in all environments)
try:
    from fastmcp.server.dependencies import get_context as get_fastmcp_context
except ImportError:
    get_fastmcp_context = None

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# Constants
def get_default_credentials_dir():
    """Get the default credentials directory path, preferring user-specific locations.
    Environment variable priority:
    1. WORKSPACE_MCP_CREDENTIALS_DIR (preferred)
    2. GOOGLE_MCP_CREDENTIALS_DIR (backward compatibility)
    3. ~/.google_workspace_mcp/credentials (default)
    """
    # Check WORKSPACE_MCP_CREDENTIALS_DIR first (preferred)
    workspace_creds_dir = os.getenv("WORKSPACE_MCP_CREDENTIALS_DIR")
    if workspace_creds_dir:
        expanded = os.path.expanduser(workspace_creds_dir)
        logger.info(
            f"Using credentials directory from WORKSPACE_MCP_CREDENTIALS_DIR: {expanded}"
        )
        return expanded

    # Fall back to GOOGLE_MCP_CREDENTIALS_DIR for backward compatibility
    google_creds_dir = os.getenv("GOOGLE_MCP_CREDENTIALS_DIR")
    if google_creds_dir:
        expanded = os.path.expanduser(google_creds_dir)
        logger.info(
            f"Using credentials directory from GOOGLE_MCP_CREDENTIALS_DIR: {expanded}"
        )
        return expanded

    # Use user home directory for credentials storage
    home_dir = os.path.expanduser("~")
    if home_dir and home_dir != "~":  # Valid home directory found
        return os.path.join(home_dir, ".google_workspace_mcp", "credentials")

    # Fallback to current working directory if home directory is not accessible
    return os.path.join(os.getcwd(), ".credentials")


DEFAULT_CREDENTIALS_DIR = get_default_credentials_dir()

# Session credentials now handled by OAuth21SessionStore - no local cache needed
# Centralized Client Secrets Path Logic
_client_secrets_env = os.getenv("GOOGLE_CLIENT_SECRET_PATH") or os.getenv(
    "GOOGLE_CLIENT_SECRETS"
)
if _client_secrets_env:
    CONFIG_CLIENT_SECRETS_PATH = _client_secrets_env
else:
    # Assumes this file is in auth/ and client_secret.json is in the root
    CONFIG_CLIENT_SECRETS_PATH = os.path.join(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
        "client_secret.json",
    )

# --- Helper Functions ---


def _find_any_credentials(
    base_dir: str = DEFAULT_CREDENTIALS_DIR,
) -> tuple[Optional[Credentials], Optional[str]]:
    """
    Find and load any valid credentials from the credentials directory.
    Used in single-user mode to bypass session-to-OAuth mapping.
    Returns:
        Tuple of (Credentials, user_email) or (None, None) if none exist.
        Returns the user email to enable saving refreshed credentials.
    """
    try:
        store = get_credential_store()
        users = store.list_users()
        if not users:
            logger.info(
                "[single-user] No users found with credentials via credential store"
            )
            return None, None

        # Return credentials for the first user found
        first_user = users[0]
        credentials = store.get_credential(first_user)
        if credentials:
            logger.info(
                f"[single-user] Found credentials for {first_user} via credential store"
            )
            return credentials, first_user
        else:
            logger.warning(
                f"[single-user] Could not load credentials for {first_user} via credential store"
            )

    except Exception as e:
        logger.error(
            f"[single-user] Error finding credentials via credential store: {e}"
        )

    logger.info("[single-user] No valid credentials found via credential store")
    return None, None


def save_credentials_to_session(session_id: str, credentials: Credentials):
    """Saves user credentials using OAuth21SessionStore."""
    # Get user email from credentials if possible
    user_email = None
    if credentials and credentials.id_token:
        try:
            decoded_token = jwt.decode(
                credentials.id_token, options={"verify_signature": False}
            )
            user_email = decoded_token.get("email")
        except Exception as e:
            logger.debug(f"Could not decode id_token to get email: {e}")

    if user_email:
        store = get_oauth21_session_store()
        store.store_session(
            user_email=user_email,
            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,
            mcp_session_id=session_id,
        )
        logger.debug(
            f"Credentials saved to OAuth21SessionStore for session_id: {session_id}, user: {user_email}"
        )
    else:
        logger.warning(
            f"Could not save credentials to session store - no user email found for session: {session_id}"
        )


def load_credentials_from_session(session_id: str) -> Optional[Credentials]:
    """Loads user credentials from OAuth21SessionStore."""
    store = get_oauth21_session_store()
    credentials = store.get_credentials_by_mcp_session(session_id)
    if credentials:
        logger.debug(
            f"Credentials loaded from OAuth21SessionStore for session_id: {session_id}"
        )
    else:
        logger.debug(
            f"No credentials found in OAuth21SessionStore for session_id: {session_id}"
        )
    return credentials


def load_client_secrets_from_env() -> Optional[Dict[str, Any]]:
    """
    Loads the client secrets from environment variables.
    Environment variables used:
        - GOOGLE_OAUTH_CLIENT_ID: OAuth client ID (required)
        - GOOGLE_OAUTH_CLIENT_SECRET: OAuth client secret (optional for public clients)
        - GOOGLE_OAUTH_REDIRECT_URI: (optional) OAuth redirect URI
    Returns:
        Client secrets configuration dict compatible with Google OAuth library,
        or None if required environment variables are not set.
    """
    client_id = os.getenv("GOOGLE_OAUTH_CLIENT_ID")
    client_secret = os.getenv("GOOGLE_OAUTH_CLIENT_SECRET")
    redirect_uri = os.getenv("GOOGLE_OAUTH_REDIRECT_URI")

    if client_id:
        # Create config structure that matches Google client secrets format.
        client_config = {
            "client_id": client_id,
            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
            "token_uri": "https://oauth2.googleapis.com/token",
            "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
        }
        # google-auth-oauthlib token exchange expects this key to exist.
        # Keep it as an empty string for public clients.
        client_config["client_secret"] = client_secret or ""

        # Add redirect_uri if provided via environment variable
        if redirect_uri:
            client_config["redirect_uris"] = [redirect_uri]

        # google-auth-oauthlib supports both "web" and "installed" shapes.
        # Use "installed" for public clients without a secret.
        top_level_key = "web" if client_secret else "installed"
        config = {top_level_key: client_config}

        logger.info("Loaded OAuth client credentials from environment variables")
        return config

    logger.debug("OAuth client credentials not found in environment variables")
    return None


def load_client_secrets(client_secrets_path: str) -> Dict[str, Any]:
    """
    Loads the client secrets from environment variables (preferred) or from the client secrets file.
    Priority order:
    1. Environment variables (GOOGLE_OAUTH_CLIENT_ID, GOOGLE_OAUTH_CLIENT_SECRET)
    2. File-based credentials at the specified path
    Args:
        client_secrets_path: Path to the client secrets JSON file (used as fallback)
    Returns:
        Client secrets configuration dict
    Raises:
        ValueError: If client secrets file has invalid format
        IOError: If file cannot be read and no environment variables are set
    """
    # First, try to load from environment variables
    env_config = load_client_secrets_from_env()
    if env_config:
        # Extract either "web" (confidential) or "installed" (public) config.
        if "web" in env_config:
            return env_config["web"]
        if "installed" in env_config:
            return env_config["installed"]
        raise ValueError(
            "Invalid environment OAuth client config format. Expected 'web' or 'installed'."
        )

    # Fall back to loading from file
    try:
        with open(client_secrets_path, "r") as f:
            client_config = json.load(f)
            # The file usually contains a top-level key like "web" or "installed"
            if "web" in client_config:
                logger.info(
                    f"Loaded OAuth client credentials from file: {client_secrets_path}"
                )
                return client_config["web"]
            elif "installed" in client_config:
                logger.info(
                    f"Loaded OAuth client credentials from file: {client_secrets_path}"
                )
                return client_config["installed"]
            else:
                logger.error(
                    f"Client secrets file {client_secrets_path} has unexpected format."
                )
                raise ValueError("Invalid client secrets file format")
    except (IOError, json.JSONDecodeError) as e:
        logger.error(f"Error loading client secrets file {client_secrets_path}: {e}")
        raise


def check_client_secrets() -> Optional[str]:
    """
    Checks for the presence of OAuth client secrets, either as environment
    variables or as a file.
    Returns:
        An error message string if secrets are not found, otherwise None.
    """
    env_config = load_client_secrets_from_env()
    if not env_config and not os.path.exists(CONFIG_CLIENT_SECRETS_PATH):
        logger.error(
            f"OAuth client credentials not found. No environment variables set and no file at {CONFIG_CLIENT_SECRETS_PATH}"
        )
        return f"OAuth client credentials not found. Please set GOOGLE_OAUTH_CLIENT_ID and GOOGLE_OAUTH_CLIENT_SECRET environment variables or provide a client secrets file at {CONFIG_CLIENT_SECRETS_PATH}."
    return None


def create_oauth_flow(
    scopes: List[str],
    redirect_uri: str,
    state: Optional[str] = None,
    code_verifier: Optional[str] = None,
    autogenerate_code_verifier: bool = True,
) -> Flow:
    """Creates an OAuth flow using environment variables or client secrets file."""
    flow_kwargs = {
        "scopes": scopes,
        "redirect_uri": redirect_uri,
        "state": state,
    }
    if code_verifier:
        flow_kwargs["code_verifier"] = code_verifier
        # Preserve the original verifier when re-creating the flow in callback.
        flow_kwargs["autogenerate_code_verifier"] = False
    else:
        # Generate PKCE code verifier for the initial auth flow.
        # google-auth-oauthlib's from_client_* helpers pass
        # autogenerate_code_verifier=None unless explicitly provided, which
        # prevents Flow from generating and storing a code_verifier.
        flow_kwargs["autogenerate_code_verifier"] = autogenerate_code_verifier

    # Try environment variables first
    env_config = load_client_secrets_from_env()
    if env_config:
        # Use client config directly
        flow = Flow.from_client_config(env_config, **flow_kwargs)
        logger.debug("Created OAuth flow from environment variables")
        return flow

    # Fall back to file-based config
    if not os.path.exists(CONFIG_CLIENT_SECRETS_PATH):
        raise FileNotFoundError(
            f"OAuth client secrets file not found at {CONFIG_CLIENT_SECRETS_PATH} and no environment variables set"
        )

    flow = Flow.from_client_secrets_file(
        CONFIG_CLIENT_SECRETS_PATH,
        **flow_kwargs,
    )
    logger.debug(
        f"Created OAuth flow from client secrets file: {CONFIG_CLIENT_SECRETS_PATH}"
    )
    return flow


def _is_pkce_verifier_not_needed_error(error: Exception) -> bool:
    """Detect Google's legacy desktop-client response when PKCE is unnecessary."""
    message = str(error).lower()
    return (
        "invalid_grant" in message
        and "code_verifier" in message
        and "not needed" in message
    )


async def _determine_oauth_prompt(
    user_google_email: Optional[str],
    required_scopes: List[str],
    session_id: Optional[str] = None,
) -> str:
    """
    Determine which OAuth prompt to use for a new authorization URL.
    Uses `select_account` for re-auth when existing credentials already cover
    required scopes. Uses `consent` for first-time auth and scope expansion.
    """
    normalized_email = (
        user_google_email.strip()
        if user_google_email
        and user_google_email.strip()
        and user_google_email.lower() != "default"
        else None
    )

    # If no explicit email was provided, attempt to resolve it from session mapping.
    if not normalized_email and session_id:
        try:
            session_user = get_oauth21_session_store().get_user_by_mcp_session(
                session_id
            )
            if session_user:
                normalized_email = session_user
        except Exception as e:
            logger.debug(f"Could not resolve user from session for prompt choice: {e}")

    if not normalized_email:
        logger.info(
            "[start_auth_flow] Using prompt='consent' (no known user email for re-auth detection)."
        )
        return "consent"

    existing_credentials: Optional[Credentials] = None

    # Prefer credentials bound to the current session when available.
    if session_id:
        try:
            session_store = get_oauth21_session_store()
            mapped_user = session_store.get_user_by_mcp_session(session_id)
            if mapped_user == normalized_email:
                existing_credentials = session_store.get_credentials_by_mcp_session(
                    session_id
                )
        except Exception as e:
            logger.debug(
                f"Could not read OAuth 2.1 session store for prompt choice: {e}"
            )

    # Fall back to credential file store in stateful mode.
    if not existing_credentials and not is_stateless_mode():
        try:
            existing_credentials = await asyncio.to_thread(
                get_credential_store().get_credential, normalized_email
            )
        except Exception as e:
            logger.debug(f"Could not read credential store for prompt choice: {e}")

    if not existing_credentials:
        logger.info(
            f"[start_auth_flow] Using prompt='consent' (no existing credentials for {normalized_email})."
        )
        return "consent"

    if has_required_scopes(existing_credentials.scopes, required_scopes):
        # Verify the credentials can still be refreshed before using select_account.
        # When credentials are revoked, Google's select_account prompt may produce
        # incomplete callbacks (missing state parameter, partial scopes).
        if existing_credentials.valid:
            logger.info(
                f"[start_auth_flow] Using prompt='select_account' for re-auth of {normalized_email}."
            )
            return "select_account"
        if existing_credentials.refresh_token:
            try:
                await asyncio.to_thread(existing_credentials.refresh, Request())
                logger.info(
                    f"[start_auth_flow] Using prompt='select_account' for re-auth of {normalized_email}."
                )
                return "select_account"
            except Exception:
                logger.info(
                    f"[start_auth_flow] Credentials for {normalized_email} cannot be refreshed; "
                    "using prompt='consent' to ensure full re-authorization."
                )
                return "consent"

    logger.info(
        f"[start_auth_flow] Using prompt='consent' (existing credentials for {normalized_email} "
        "are not refreshable or are missing required scopes)."
    )
    return "consent"


# --- Core OAuth Logic ---


async def start_auth_flow(
    user_google_email: Optional[str],
    service_name: str,  # e.g., "Google Calendar", "Gmail" for user messages
    redirect_uri: str,  # Added redirect_uri as a required parameter
) -> str:
    """
    Initiates the Google OAuth flow and returns an actionable message for the user.
    Args:
        user_google_email: The user's specified Google email, if provided.
        service_name: The name of the Google service requiring auth (for user messages).
        redirect_uri: The URI Google will redirect to after authorization.
    Returns:
        A formatted string containing guidance for the LLM/user.
    Raises:
        Exception: If the OAuth flow cannot be initiated.
    """
    initial_email_provided = bool(
        user_google_email
        and user_google_email.strip()
        and user_google_email.lower() != "default"
    )
    user_display_name = (
        f"{service_name} for '{user_google_email}'"
        if initial_email_provided
        else service_name
    )

    logger.info(
        f"[start_auth_flow] Initiating auth for {user_display_name} with scopes for enabled tools."
    )

    # Note: Caller should ensure OAuth callback is available before calling this function

    try:
        if "OAUTHLIB_INSECURE_TRANSPORT" not in os.environ and (
            "localhost" in redirect_uri or "127.0.0.1" in redirect_uri
        ):  # Use passed redirect_uri
            logger.warning(
                "OAUTHLIB_INSECURE_TRANSPORT not set. Setting it for localhost/local development."
            )
            os.environ["OAUTHLIB_INSECURE_TRANSPORT"] = "1"

        oauth_state = os.urandom(16).hex()
        current_scopes = get_current_scopes()

        flow = create_oauth_flow(
            scopes=current_scopes,  # Use scopes for enabled tools only
            redirect_uri=redirect_uri,  # Use passed redirect_uri
            state=oauth_state,
        )

        session_id = None
        try:
            session_id = get_fastmcp_session_id()
        except Exception as e:
            logger.debug(
                f"Could not retrieve FastMCP session ID for state binding: {e}"
            )

        prompt_type = await _determine_oauth_prompt(
            user_google_email=user_google_email,
            required_scopes=current_scopes,
            session_id=session_id,
        )
        auth_url, _ = flow.authorization_url(access_type="offline", prompt=prompt_type)

        store = get_oauth21_session_store()
        store.store_oauth_state(
            oauth_state,
            session_id=session_id,
            code_verifier=flow.code_verifier,
        )

        logger.info(
            f"Auth flow started for {user_display_name}. Advise user to visit: {auth_url}"
        )

        message_lines = [
            f"**ACTION REQUIRED: Google Authentication Needed for {user_display_name}**\n",
            f"To proceed, the user must authorize this application for {service_name} access using all required permissions.",
            "**LLM, please present this exact authorization URL to the user as a clickable hyperlink:**",
            f"Authorization URL: {auth_url}",
            f"Markdown for hyperlink: [Click here to authorize {service_name} access]({auth_url})\n",
            "**LLM, after presenting the link, instruct the user as follows:**",
            "1. Click the link and complete the authorization in their browser.",
        ]
        session_info_for_llm = ""

        if not initial_email_provided:
            message_lines.extend(
                [
                    f"2. After successful authorization{session_info_for_llm}, the browser page will display the authenticated email address.",
                    "   **LLM: Instruct the user to provide you with this email address.**",
                    "3. Once you have the email, **retry their original command, ensuring you include this `user_google_email`.**",
                ]
            )
        else:
            message_lines.append(
                f"2. After successful authorization{session_info_for_llm}, **retry their original command**."
            )

        message_lines.append(
            f"\nThe application will use the new credentials. If '{user_google_email}' was provided, it must match the authenticated account."
        )
        return "\n".join(message_lines)

    except FileNotFoundError as e:
        error_text = f"OAuth client credentials not found: {e}. Please either:\n1. Set environment variables: GOOGLE_OAUTH_CLIENT_ID and GOOGLE_OAUTH_CLIENT_SECRET\n2. Ensure '{CONFIG_CLIENT_SECRETS_PATH}' file exists"
        logger.error(error_text, exc_info=True)
        raise Exception(error_text)
    except Exception as e:
        error_text = f"Could not initiate authentication for {user_display_name} due to an unexpected error: {str(e)}"
        logger.error(
            f"Failed to start the OAuth flow for {user_display_name}: {e}",
            exc_info=True,
        )
        raise Exception(error_text)


async def handle_auth_callback(
    scopes: List[str],
    authorization_response: str,
    redirect_uri: str,
    credentials_base_dir: str = DEFAULT_CREDENTIALS_DIR,
    session_id: Optional[str] = None,
    client_secrets_path: Optional[
        str
    ] = None,  # Deprecated: kept for backward compatibility
) -> Tuple[str, Credentials]:
    """
    Handles the callback from Google, exchanges the code for credentials,
    fetches user info, determines user_google_email, saves credentials (file & session),
    and returns them.
    Args:
        scopes: List of OAuth scopes requested.
        authorization_response: The full callback URL from Google.
        redirect_uri: The redirect URI.
        credentials_base_dir: Base directory for credential files.
        session_id: Optional MCP session ID to associate with the credentials.
        client_secrets_path: (Deprecated) Path to client secrets file. Ignored if environment variables are set.
    Returns:
        A tuple containing the user_google_email and the obtained Credentials object.
    Raises:
        ValueError: If the state is missing or doesn't match.
        FlowExchangeError: If the code exchange fails.
        HttpError: If fetching user info fails.
    """
    try:
        # Log deprecation warning if old parameter is used
        if client_secrets_path:
            logger.warning(
                "The 'client_secrets_path' parameter is deprecated. Use GOOGLE_OAUTH_CLIENT_ID and GOOGLE_OAUTH_CLIENT_SECRET environment variables instead."
            )

        # 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"

        # Allow partial scope grants without raising an exception.
        # When users decline some scopes on Google's consent screen,
        # oauthlib raises because the granted scopes differ from requested.
        if "OAUTHLIB_RELAX_TOKEN_SCOPE" not in os.environ:
            os.environ["OAUTHLIB_RELAX_TOKEN_SCOPE"] = "1"

        store = get_oauth21_session_store()
        parsed_response = urlparse(authorization_response)
        state_values = parse_qs(parsed_response.query).get("state")
        state = state_values[0] if state_values else None

        if state:
            state_info = store.validate_and_consume_oauth_state(
                state, session_id=session_id
            )
        elif session_id is None:
            # stdio mode fallback: state may be absent from Google's redirect
            # (e.g. when prompt=select_account is used with revoked credentials).
            # Use the most recently stored state to recover the PKCE code_verifier.
            logger.warning(
                "OAuth callback missing state parameter; using most recent stored state (stdio fallback)"
            )
            state_info = store.consume_latest_oauth_state(
                initiating_session_id=session_id
            )
            if not state_info:
                raise ValueError(
                    "Missing OAuth state parameter and no stored state available"
                )
        else:
            raise ValueError("Missing OAuth state parameter")

        logger.debug(
            "OAuth callback state %s for session %s",
            (state[:8] if state else "<fallback>"),
            state_info.get("session_id") or "<unknown>",
        )

        flow = create_oauth_flow(
            scopes=scopes,
            redirect_uri=redirect_uri,
            state=state,
            code_verifier=state_info.get("code_verifier"),
            autogenerate_code_verifier=False,
        )

        # Exchange the authorization code for credentials
        # Note: fetch_token will use the redirect_uri configured in the flow
        try:
            await asyncio.to_thread(
                flow.fetch_token, authorization_response=authorization_response
            )
            credentials = flow.credentials
        except Exception as exc:
            if _is_pkce_verifier_not_needed_error(exc):
                logger.error(
                    "OAuth token exchange rejected PKCE verifier. "
                    "The authorization code has been consumed and cannot be reused. "
                    "Please restart the authentication flow from the beginning."
                )
            raise
        logger.info("Successfully exchanged authorization code for tokens.")

        # Handle partial OAuth grants: if the user declined some scopes on
        # Google's consent screen, credentials.granted_scopes contains only
        # what was actually authorized. Store those instead of the inflated
        # requested scopes so that refresh() sends the correct scope set.
        granted = getattr(credentials, "granted_scopes", None)
        if granted and set(granted) != set(credentials.scopes or []):
            logger.warning(
                "Partial OAuth grant detected. Requested: %s, Granted: %s",
                credentials.scopes,
                granted,
            )
            credentials = Credentials(
                token=credentials.token,
                refresh_token=credentials.refresh_token,
                id_token=getattr(credentials, "id_token", None),
                token_uri=credentials.token_uri,
                client_id=credentials.client_id,
                client_secret=credentials.client_secret,
                scopes=list(granted),
                expiry=credentials.expiry,
                quota_project_id=getattr(credentials, "quota_project_id", None),
            )

        # Get user info to determine user_id (using email here)
        user_info = await asyncio.to_thread(get_user_info, credentials)
        if not user_info or "email" not in user_info:
            logger.error("Could not retrieve user email from Google.")
            raise ValueError("Failed to get user email for identification.")

        user_google_email = user_info["email"]
        logger.info(f"Identified user_google_email: {user_google_email}")

        stateless_mode = is_stateless_mode()
        credential_store = None
        if not stateless_mode:
            credential_store = get_credential_store()
        if not credentials.refresh_token:
            fallback_refresh_token = None

            if session_id:
                try:
                    session_credentials = store.get_credentials_by_mcp_session(
                        session_id
                    )
                    if session_credentials and session_credentials.refresh_token:
                        fallback_refresh_token = session_credentials.refresh_token
                        logger.info(
                            "OAuth callback response omitted refresh token; preserving existing refresh token from session store."
                        )
                except Exception as e:
                    logger.debug(
                        f"Could not check session store for existing refresh token: {e}"
                    )

            if not fallback_refresh_token and not stateless_mode:
                try:
                    existing_credentials = await asyncio.to_thread(
                        credential_store.get_credential, user_google_email
                    )
                    if existing_credentials and existing_credentials.refresh_token:
                        fallback_refresh_token = existing_credentials.refresh_token
                        logger.info(
                            "OAuth callback response omitted refresh token; preserving existing refresh token from credential store."
                        )
                except Exception as e:
                    logger.debug(
                        f"Could not check credential store for existing refresh token: {e}"
                    )

            if fallback_refresh_token:
                credentials = Credentials(
                    token=credentials.token,
                    refresh_token=fallback_refresh_token,
                    id_token=getattr(credentials, "id_token", None),
                    token_uri=credentials.token_uri,
                    client_id=credentials.client_id,
                    client_secret=credentials.client_secret,
                    scopes=credentials.scopes,
                    expiry=credentials.expiry,
                    quota_project_id=getattr(credentials, "quota_project_id", None),
                )
            else:
                logger.warning(
                    "OAuth callback did not include a refresh token and no previous refresh token was available to preserve."
                )

        if not stateless_mode:
            # Save the credentials before updating session state so both stores stay in sync.
            stored = await asyncio.to_thread(
                credential_store.store_credential, user_google_email, credentials
            )
            if not stored:
                logger.warning(
                    "Credential store rejected updated credentials for %s; aborting session persistence.",
                    user_google_email,
                )
                raise RuntimeError(
                    f"Failed to persist credentials for {user_google_email}; "
                    "session state was not updated."
                )

        # Always save to OAuth21SessionStore for centralized management
        store.store_session(
            user_email=user_google_email,
            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,
            mcp_session_id=session_id,
            issuer="https://accounts.google.com",  # Add issuer for Google tokens
        )

        # If session_id is provided, also save to session cache for compatibility
        if session_id:
            save_credentials_to_session(session_id, credentials)

        return user_google_email, credentials

    except Exception as e:  # Catch specific exceptions like FlowExchangeError if needed
        logger.error(f"Error handling auth callback: {e}")
        raise  # Re-raise for the caller


def get_credentials(
    user_google_email: Optional[str],  # Can be None if relying on session_id
    required_scopes: List[str],
    client_secrets_path: Optional[str] = None,
    credentials_base_dir: str = DEFAULT_CREDENTIALS_DIR,
    session_id: Optional[str] = None,
) -> Optional[Credentials]:
    """
    Retrieves stored credentials, prioritizing OAuth 2.1 store, then session, then file. Refreshes if necessary.
    If credentials are loaded from file and a session_id is present, they are cached in the session.
    In single-user mode, bypasses session mapping. If user_google_email is provided, only credentials
    for that email are used and the function returns None instead of falling back to any available
    credentials. If user_google_email is not provided, any available credentials may be used.
    Args:
        user_google_email: Optional user's Google email.
        required_scopes: List of scopes the credentials must have.
        client_secrets_path: Optional path to client secrets (legacy; refresh uses embedded client info).
        credentials_base_dir: Base directory for credential files.
        session_id: Optional MCP session ID.
    Returns:
        Valid Credentials object or None.
    """
    skip_session_cache = False
    # First, try OAuth 2.1 session store if we have a session_id (FastMCP session)
    if session_id:
        try:
            store = get_oauth21_session_store()

            session_user = store.get_user_by_mcp_session(session_id)
            if user_google_email and session_user and session_user != user_google_email:
                logger.info(
                    f"[get_credentials] Session user {session_user} doesn't match requested {user_google_email}; "
                    "skipping session store"
                )
                skip_session_cache = True
            else:
                # Try to get credentials by MCP session
                credentials = store.get_credentials_by_mcp_session(session_id)
                if credentials:
                    logger.info(
                        f"[get_credentials] Found OAuth 2.1 credentials for MCP session {session_id}"
                    )

                    # Refresh invalid credentials before checking scopes
                    if (not credentials.valid) and credentials.refresh_token:
                        try:
                            credentials.refresh(Request())
                            logger.info(
                                f"[get_credentials] Refreshed OAuth 2.1 credentials for session {session_id}"
                            )
                            # Update stored credentials
                            user_email = store.get_user_by_mcp_session(session_id)
                            if user_email:
                                # Persist to file so rotated refresh tokens survive restarts
                                persist_succeeded = True
                                if not is_stateless_mode():
                                    try:
                                        credential_store = get_credential_store()
                                        persist_succeeded = (
                                            credential_store.store_credential(
                                                user_email, credentials
                                            )
                                        )
                                        if not persist_succeeded:
                                            logger.warning(
                                                "[get_credentials] Credential store rejected refreshed OAuth 2.1 credentials for user %s; skipping session update.",
                                                user_email,
                                            )
                                    except Exception as persist_error:
                                        persist_succeeded = False
                                        logger.warning(
                                            f"[get_credentials] Failed to persist refreshed OAuth 2.1 credentials for user {user_email}: {persist_error}"
                                        )

                                if not persist_succeeded and not is_stateless_mode():
                                    logger.warning(
                                        "[get_credentials] Refreshed OAuth 2.1 credentials for user %s were not persisted; discarding in-memory refresh result.",
                                        user_email,
                                    )
                                    return None

                                if persist_succeeded or is_stateless_mode():
                                    store.store_session(
                                        user_email=user_email,
                                        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,
                                        mcp_session_id=session_id,
                                        issuer="https://accounts.google.com",
                                    )
                        except Exception as e:
                            logger.error(
                                f"[get_credentials] Failed to refresh OAuth 2.1 credentials: {e}"
                            )
                            return None

                    # Check scopes after refresh so stale metadata doesn't block valid tokens
                    if not has_required_scopes(credentials.scopes, required_scopes):
                        logger.warning(
                            f"[get_credentials] OAuth 2.1 credentials lack required scopes. Need: {required_scopes}, Have: {credentials.scopes}"
                        )
                        return None

                    if credentials.valid:
                        return credentials

                    return None
        except ImportError:
            pass  # OAuth 2.1 store not available
        except Exception as e:
            logger.debug(f"[get_credentials] Error checking OAuth 2.1 store: {e}")

    # Check for single-user mode
    if os.getenv("MCP_SINGLE_USER_MODE") == "1":
        logger.info(
            "[get_credentials] Single-user mode: bypassing session mapping, finding any credentials"
        )
        # If a specific email was requested, try to load that user's credentials first
        # to avoid session binding conflicts when multiple credential files exist
        if user_google_email:
            credential_store = get_credential_store()
            credentials = credential_store.get_credential(user_google_email)
            if credentials:
                logger.info(
                    f"[get_credentials] Single-user mode: found credentials for requested user {user_google_email}"
                )
                found_user_email = user_google_email
            else:
                logger.info(
                    "[get_credentials] Single-user mode: no credentials for requested "
                    f"user {user_google_email}; not falling back to another user"
                )
                return None
        else:
            credentials, found_user_email = _find_any_credentials(credentials_base_dir)
        if not credentials:
            logger.info(
                f"[get_credentials] Single-user mode: No credentials found in {credentials_base_dir}"
            )
            return None

        # Use the email from the credential file if not provided
        # This ensures we can save refreshed credentials even when the token is expired
        if not user_google_email and found_user_email:
            user_google_email = found_user_email
            logger.debug(
                f"[get_credentials] Single-user mode: using email {user_google_email} from credential file"
            )
    else:
        credentials: Optional[Credentials] = None

        # Session ID should be provided by the caller
        if not session_id:
            logger.debug("[get_credentials] No session_id provided")

        logger.debug(
            f"[get_credentials] Called for user_google_email: '{user_google_email}', session_id: '{session_id}', required_scopes: {required_scopes}"
        )

        if session_id and not skip_session_cache:
            credentials = load_credentials_from_session(session_id)
            if credentials:
                logger.debug(
                    f"[get_credentials] Loaded credentials from session for session_id '{session_id}'."
                )

        if not credentials and user_google_email:
            if not is_stateless_mode():
                logger.debug(
                    f"[get_credentials] No session credentials, trying credential store for user_google_email '{user_google_email}'."
                )
                store = get_credential_store()
                credentials = store.get_credential(user_google_email)
            else:
                logger.debug(
                    f"[get_credentials] No session credentials, skipping file store in stateless mode for user_google_email '{user_google_email}'."
                )

            if credentials and session_id:
                logger.debug(
                    f"[get_credentials] Loaded from file for user '{user_google_email}', caching to session '{session_id}'."
                )
                if not skip_session_cache:
                    save_credentials_to_session(
                        session_id, credentials
                    )  # Cache for current session

        if not credentials:
            logger.info(
                f"[get_credentials] No credentials found for user '{user_google_email}' or session '{session_id}'."
            )
            return None

    logger.debug(
        f"[get_credentials] Credentials found. Scopes: {credentials.scopes}, Valid: {credentials.valid}, Expired: {credentials.expired}"
    )

    # Attempt refresh before checking scopes — the scope check validates against
    # credentials.scopes which is set at authorization time and not updated by the
    # google-auth library on refresh. Checking scopes first would block a valid
    # refresh attempt when stored scope metadata is stale.
    if credentials.valid:
        logger.debug(
            f"[get_credentials] Credentials are valid. User: '{user_google_email}', Session: '{session_id}'"
        )
    elif credentials.refresh_token:
        logger.info(
            f"[get_credentials] Credentials not valid. Attempting refresh. User: '{user_google_email}', Session: '{session_id}'"
        )
        try:
            logger.debug(
                "[get_credentials] Refreshing token using embedded client credentials"
            )
            credentials.refresh(Request())
            logger.info(
                f"[get_credentials] Credentials refreshed successfully. User: '{user_google_email}', Session: '{session_id}'"
            )

            # Save refreshed credentials (skip file save in stateless mode)
            persist_succeeded = True
            if user_google_email:  # Always save to credential store if email is known
                if not is_stateless_mode():
                    try:
                        credential_store = get_credential_store()
                        persist_succeeded = credential_store.store_credential(
                            user_google_email, credentials
                        )
                    except Exception as persist_error:
                        persist_succeeded = False
                        logger.warning(
                            "[get_credentials] Failed to persist refreshed credentials for user %s: %s",
                            user_google_email,
                            persist_error,
                        )

                    if not persist_succeeded:
                        logger.warning(
                            "[get_credentials] Credential store rejected refreshed credentials for user %s; skipping session update.",
                            user_google_email,
                        )
                else:
                    logger.info(
                        f"Skipping credential file save in stateless mode for {user_google_email}"
                    )

                if not persist_succeeded and not is_stateless_mode():
                    logger.warning(
                        "[get_credentials] Refreshed credentials for user %s were not persisted; discarding in-memory refresh result.",
                        user_google_email,
                    )
                    return None

                if persist_succeeded or is_stateless_mode():
                    # Also update OAuth21SessionStore
                    store = get_oauth21_session_store()
                    store.store_session(
                        user_email=user_google_email,
                        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,
                        mcp_session_id=session_id,
                        issuer="https://accounts.google.com",  # Add issuer for Google tokens
                    )

            if session_id and (persist_succeeded or is_stateless_mode()):
                # Update session cache if it was the source or is active
                save_credentials_to_session(session_id, credentials)
        except RefreshError as e:
            logger.warning(
                f"[get_credentials] RefreshError - token expired/revoked: {e}. User: '{user_google_email}', Session: '{session_id}'"
            )
            # For RefreshError, we should return None to trigger reauthentication
            return None
        except Exception as e:
            logger.error(
                f"[get_credentials] Error refreshing credentials: {e}. User: '{user_google_email}', Session: '{session_id}'",
                exc_info=True,
            )
            return None  # Failed to refresh
    else:
        logger.warning(
            f"[get_credentials] Credentials invalid/cannot refresh. Valid: {credentials.valid}, Refresh Token: {credentials.refresh_token is not None}. User: '{user_google_email}', Session: '{session_id}'"
        )
        return None

    # Check scopes after refresh so stale scope metadata doesn't block valid tokens.
    # Uses hierarchy-aware check (e.g. gmail.modify satisfies gmail.readonly).
    if not has_required_scopes(credentials.scopes, required_scopes):
        logger.warning(
            f"[get_credentials] Credentials lack required scopes. Need: {required_scopes}, Have: {credentials.scopes}. User: '{user_google_email}', Session: '{session_id}'"
        )
        return None  # Re-authentication needed for scopes

    logger.debug(
        f"[get_credentials] Credentials have sufficient scopes. User: '{user_google_email}', Session: '{session_id}'"
    )
    return credentials


def get_user_info(
    credentials: Credentials, *, skip_valid_check: bool = False
) -> Optional[Dict[str, Any]]:
    """Fetches basic user profile information (requires userinfo.email scope)."""
    if not credentials:
        logger.error("Cannot get user info: Missing credentials.")
        return None
    if not skip_valid_check and not credentials.valid:
        logger.error("Cannot get user info: Invalid credentials.")
        return None
    service = None
    try:
        # Using googleapiclient discovery to get user info
        # Requires 'google-api-python-client' library
        service = build("oauth2", "v2", credentials=credentials)
        user_info = service.userinfo().get().execute()
        logger.info(f"Successfully fetched user info: {user_info.get('email')}")
        return user_info
    except HttpError as e:
        logger.error(f"HttpError fetching user info: {e.status_code} {e.reason}")
        # Handle specific errors, e.g., 401 Unauthorized might mean token issue
        return None
    except Exception as e:
        logger.error(f"Unexpected error fetching user info: {e}")
        return None
    finally:
        if service:
            service.close()


# --- Centralized Google Service Authentication ---


class GoogleAuthenticationError(Exception):
    """Exception raised when Google authentication is required or fails."""

    def __init__(self, message: str, auth_url: Optional[str] = None):
        super().__init__(message)
        self.auth_url = auth_url


async def get_authenticated_google_service(
    service_name: str,  # "gmail", "calendar", "drive", "docs"
    version: str,  # "v1", "v3"
    tool_name: str,  # For logging/debugging
    user_google_email: str,  # Required - no more Optional
    required_scopes: List[str],
    session_id: Optional[str] = None,  # Session context for logging
) -> tuple[Any, str]:
    """
    Centralized Google service authentication for all MCP tools.
    Returns (service, user_email) on success or raises GoogleAuthenticationError.
    Args:
        service_name: The Google service name ("gmail", "calendar", "drive", "docs")
        version: The API version ("v1", "v3", etc.)
        tool_name: The name of the calling tool (for logging/debugging)
        user_google_email: The user's Google email address (required)
        required_scopes: List of required OAuth scopes
    Returns:
        tuple[service, user_email] on success
    Raises:
        GoogleAuthenticationError: When authentication is required or fails
    """

    # Try to get FastMCP session ID if not provided
    if not session_id:
        try:
            # First try context variable (works in async context)
            session_id = get_fastmcp_session_id()
            if session_id:
                logger.debug(
                    f"[{tool_name}] Got FastMCP session ID from context: {session_id}"
                )
            else:
                logger.debug(
                    f"[{tool_name}] Context variable returned None/empty session ID"
                )
        except Exception as e:
            logger.debug(
                f"[{tool_name}] Could not get FastMCP session from context: {e}"
            )

        # Fallback to direct FastMCP context if context variable not set
        if not session_id and get_fastmcp_context:
            try:
                fastmcp_ctx = get_fastmcp_context()
                if fastmcp_ctx and hasattr(fastmcp_ctx, "session_id"):
                    session_id = fastmcp_ctx.session_id
                    logger.debug(
                        f"[{tool_name}] Got FastMCP session ID directly: {session_id}"
                    )
                else:
                    logger.debug(
                        f"[{tool_name}] FastMCP context exists but no session_id attribute"
                    )
            except Exception as e:
                logger.debug(
                    f"[{tool_name}] Could not get FastMCP context directly: {e}"
                )

        # Final fallback: log if we still don't have session_id
        if not session_id:
            logger.warning(
                f"[{tool_name}] Unable to obtain FastMCP session ID from any source"
            )

    logger.info(
        f"[{tool_name}] Attempting to get authenticated {service_name} service. Email: '{user_google_email}', Session: '{session_id}'"
    )

    # Validate email format
    if not user_google_email or "@" not in user_google_email:
        error_msg = f"Authentication required for {tool_name}. No valid 'user_google_email' provided. Please provide a valid Google email address."
        logger.info(f"[{tool_name}] {error_msg}")
        raise GoogleAuthenticationError(error_msg)

    credentials = await asyncio.to_thread(
        get_credentials,
        user_google_email=user_google_email,
        required_scopes=required_scopes,
        client_secrets_path=CONFIG_CLIENT_SECRETS_PATH,
        session_id=session_id,  # Pass through session context
    )

    if not credentials or not credentials.valid:
        logger.warning(
            f"[{tool_name}] No valid credentials. Email: '{user_google_email}'."
        )
        logger.info(
            f"[{tool_name}] Valid email '{user_google_email}' provided, initiating auth flow."
        )

        redirect_uri = get_oauth_redirect_uri()
        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

            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 ""
                raise GoogleAuthenticationError(
                    f"Cannot initiate OAuth flow - callback server unavailable{error_detail}"
                )

        # Generate auth URL and raise exception with it
        auth_response = await start_auth_flow(
            user_google_email=user_google_email,
            service_name=f"Google {service_name.title()}",
            redirect_uri=redirect_uri,
        )

        # Extract the auth URL from the response and raise with it
        raise GoogleAuthenticationError(auth_response)

    try:
        service = build(service_name, version, credentials=credentials)
        log_user_email = user_google_email

        # Try to get email from credentials if needed for validation
        if credentials and credentials.id_token:
            try:
                # Decode without verification (just to get email for logging)
                decoded_token = jwt.decode(
                    credentials.id_token, options={"verify_signature": False}
                )
                token_email = decoded_token.get("email")
                if token_email:
                    log_user_email = token_email
                    logger.info(f"[{tool_name}] Token email: {token_email}")
            except Exception as e:
                logger.debug(f"[{tool_name}] Could not decode id_token: {e}")

        logger.info(
            f"[{tool_name}] Successfully authenticated {service_name} service for user: {log_user_email}"
        )
        return service, log_user_email

    except Exception as e:
        error_msg = f"[{tool_name}] Failed to build {service_name} service: {str(e)}"
        logger.error(error_msg, exc_info=True)
        raise GoogleAuthenticationError(error_msg)