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

Merge branch 'main' of https://github.com/taylorwilsdon/google_workspace_mcp into credential-provider

Browse Source
This commit is contained in:
Taylor Wilsdon
2025-08-18 13:04:36 -04:00
parent 8d2167222f 97b59a2b4d
commit 03cb54a2e1
23 changed files with 3240 additions and 2022 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
auth/oauth_common_handlers.py
View File

@@ -30,7 +30,7 @@ logger = logging.getLogger(__name__)
async def handle_oauth_authorize(request: Request):
"""Common handler for OAuth authorization proxy."""
origin = request.headers.get("origin")
if request.method == "OPTIONS":
cors_headers = get_development_cors_headers(origin)
return JSONResponse(content={}, headers=cors_headers)
@@ -69,7 +69,7 @@ async def handle_oauth_authorize(request: Request):
async def handle_proxy_token_exchange(request: Request):
"""Common handler for OAuth token exchange proxy with comprehensive error handling."""
origin = request.headers.get("origin")
if request.method == "OPTIONS":
cors_headers = get_development_cors_headers(origin)
return JSONResponse(content={}, headers=cors_headers)
@@ -201,7 +201,7 @@ async def handle_proxy_token_exchange(request: Request):
"Cache-Control": "no-store"
}
response_headers.update(cors_headers)
return JSONResponse(
status_code=response.status,
content=response_data,
@@ -228,7 +228,7 @@ async def handle_oauth_protected_resource(request: Request):
Handle OAuth protected resource metadata requests.
"""
origin = request.headers.get("origin")
# Handle preflight
if request.method == "OPTIONS":
cors_headers = get_development_cors_headers(origin)
@@ -262,7 +262,7 @@ async def handle_oauth_protected_resource(request: Request):
"Cache-Control": "public, max-age=3600"
}
response_headers.update(cors_headers)
return JSONResponse(
content=metadata,
headers=response_headers
@@ -274,13 +274,13 @@ async def handle_oauth_authorization_server(request: Request):
Handle OAuth authorization server metadata.
"""
origin = request.headers.get("origin")
if request.method == "OPTIONS":
cors_headers = get_development_cors_headers(origin)
return JSONResponse(content={}, headers=cors_headers)
config = get_oauth_config()
# Get authorization server metadata from centralized config
# Pass scopes directly to keep all metadata generation in one place
metadata = config.get_authorization_server_metadata(scopes=get_current_scopes())
@@ -294,7 +294,7 @@ async def handle_oauth_authorization_server(request: Request):
"Cache-Control": "public, max-age=3600"
}
response_headers.update(cors_headers)
return JSONResponse(
content=metadata,
headers=response_headers
@@ -304,7 +304,7 @@ async def handle_oauth_authorization_server(request: Request):
async def handle_oauth_client_config(request: Request):
"""Common handler for OAuth client configuration."""
origin = request.headers.get("origin")
if request.method == "OPTIONS":
cors_headers = get_development_cors_headers(origin)
return JSONResponse(content={}, headers=cors_headers)
@@ -328,7 +328,6 @@ async def handle_oauth_client_config(request: Request):
"client_uri": config.base_url,
"redirect_uris": [
f"{config.base_url}/oauth2callback",
"http://localhost:5173/auth/callback"
],
"grant_types": ["authorization_code", "refresh_token"],
"response_types": ["code"],
@@ -347,7 +346,7 @@ async def handle_oauth_client_config(request: Request):
async def handle_oauth_register(request: Request):
"""Common handler for OAuth dynamic client registration with comprehensive error handling."""
origin = request.headers.get("origin")
if request.method == "OPTIONS":
cors_headers = get_development_cors_headers(origin)
return JSONResponse(content={}, headers=cors_headers)

12
auth/oauth_config.py
View File

@@ -203,6 +203,18 @@ class OAuthConfig:
if params.has_pkce:
return "oauth21"
# Additional detection: Check if we have an active OAuth 2.1 session
# This is important for tool calls where PKCE params aren't available
authenticated_user = request_params.get("authenticated_user")
if authenticated_user:
try:
from auth.oauth21_session_store import get_oauth21_session_store
store = get_oauth21_session_store()
if store.has_session(authenticated_user):
return "oauth21"
except (ImportError, AttributeError, RuntimeError):
pass # Fall back to OAuth 2.0 if session check fails
# For public clients in OAuth 2.1 mode, we require PKCE
# But since they didn't send PKCE, fall back to OAuth 2.0
# This ensures backward compatibility

590
auth/service_decorator.py
View File

@@ -1,29 +1,199 @@
import inspect
import logging
import re
from functools import wraps
from typing import Dict, List, Optional, Any, Callable, Union
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Callable, Union, Tuple
from google.auth.exceptions import RefreshError
from googleapiclient.discovery import build
from fastmcp.server.dependencies import get_context
from auth.google_auth import get_authenticated_google_service, GoogleAuthenticationError
from auth.oauth21_session_store import get_oauth21_session_store
from auth.oauth_config import is_oauth21_enabled, get_oauth_config
from core.context import set_fastmcp_session_id
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,
TASKS_SCOPE, TASKS_READONLY_SCOPE,
CUSTOM_SEARCH_SCOPE
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,
TASKS_SCOPE,
TASKS_READONLY_SCOPE,
CUSTOM_SEARCH_SCOPE,
)
# OAuth 2.1 integration is now handled by FastMCP auth
OAUTH21_INTEGRATION_AVAILABLE = True
# REMOVED: _extract_and_verify_bearer_token function. This functionality is now handled by AuthInfoMiddleware.
# Authentication helper functions
def _get_auth_context(
tool_name: str,
) -> Tuple[Optional[str], Optional[str], Optional[str]]:
"""
Get authentication context from FastMCP.
Returns:
Tuple of (authenticated_user, auth_method, mcp_session_id)
"""
try:
ctx = get_context()
if not ctx:
return None, None, None
authenticated_user = ctx.get_state("authenticated_user_email")
auth_method = ctx.get_state("authenticated_via")
mcp_session_id = ctx.session_id if hasattr(ctx, "session_id") else None
if mcp_session_id:
set_fastmcp_session_id(mcp_session_id)
logger.debug(
f"[{tool_name}] Auth from middleware: {authenticated_user} via {auth_method}"
)
return authenticated_user, auth_method, mcp_session_id
except Exception as e:
logger.debug(f"[{tool_name}] Could not get FastMCP context: {e}")
return None, None, None
def _detect_oauth_version(
authenticated_user: Optional[str], mcp_session_id: Optional[str], tool_name: str
) -> bool:
"""
Detect whether to use OAuth 2.1 based on configuration and context.
Returns:
True if OAuth 2.1 should be used, False otherwise
"""
if not is_oauth21_enabled():
return False
# When OAuth 2.1 is enabled globally, ALWAYS use OAuth 2.1 for authenticated users
if authenticated_user:
logger.info(
f"[{tool_name}] OAuth 2.1 mode: Using OAuth 2.1 for authenticated user '{authenticated_user}'"
)
return True
# Only use version detection for unauthenticated requests
config = get_oauth_config()
request_params = {}
if mcp_session_id:
request_params["session_id"] = mcp_session_id
oauth_version = config.detect_oauth_version(request_params)
use_oauth21 = oauth_version == "oauth21"
logger.info(
f"[{tool_name}] OAuth version detected: {oauth_version}, will use OAuth 2.1: {use_oauth21}"
)
return use_oauth21
def _update_email_in_args(args: tuple, index: int, new_email: str) -> tuple:
"""Update email at specific index in args tuple."""
if index < len(args):
args_list = list(args)
args_list[index] = new_email
return tuple(args_list)
return args
def _override_oauth21_user_email(
use_oauth21: bool,
authenticated_user: Optional[str],
current_user_email: str,
args: tuple,
kwargs: dict,
param_names: List[str],
tool_name: str,
service_type: str = "",
) -> Tuple[str, tuple]:
"""
Override user_google_email with authenticated user when using OAuth 2.1.
Returns:
Tuple of (updated_user_email, updated_args)
"""
if not (use_oauth21 and authenticated_user and current_user_email != authenticated_user):
return current_user_email, args
service_suffix = f" for service '{service_type}'" if service_type else ""
logger.info(
f"[{tool_name}] OAuth 2.1: Overriding user_google_email from '{current_user_email}' to authenticated user '{authenticated_user}'{service_suffix}"
)
# Update in kwargs if present
if "user_google_email" in kwargs:
kwargs["user_google_email"] = authenticated_user
# Update in args if user_google_email is passed positionally
try:
user_email_index = param_names.index("user_google_email")
args = _update_email_in_args(args, user_email_index, authenticated_user)
except ValueError:
pass # user_google_email not in positional parameters
return authenticated_user, args
async def _authenticate_service(
use_oauth21: bool,
service_name: str,
service_version: str,
tool_name: str,
user_google_email: str,
resolved_scopes: List[str],
mcp_session_id: Optional[str],
authenticated_user: Optional[str],
) -> Tuple[Any, str]:
"""
Authenticate and get Google service using appropriate OAuth version.
Returns:
Tuple of (service, actual_user_email)
"""
if use_oauth21:
logger.debug(f"[{tool_name}] Using OAuth 2.1 flow")
return await get_authenticated_google_service_oauth21(
service_name=service_name,
version=service_version,
tool_name=tool_name,
user_google_email=user_google_email,
required_scopes=resolved_scopes,
session_id=mcp_session_id,
auth_token_email=authenticated_user,
allow_recent_auth=False,
)
else:
logger.debug(f"[{tool_name}] Using legacy OAuth 2.0 flow")
return 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,
session_id=mcp_session_id,
)
async def get_authenticated_google_service_oauth21(
service_name: str,
version: str,
@@ -37,9 +207,6 @@ async def get_authenticated_google_service_oauth21(
"""
OAuth 2.1 authentication using the session store with security validation.
"""
from auth.oauth21_session_store import get_oauth21_session_store
from googleapiclient.discovery import build
store = get_oauth21_session_store()
# Use the new validation method to ensure session can only access its own credentials
@@ -47,11 +214,10 @@ async def get_authenticated_google_service_oauth21(
requested_user_email=user_google_email,
session_id=session_id,
auth_token_email=auth_token_email,
allow_recent_auth=allow_recent_auth
allow_recent_auth=allow_recent_auth,
)
if not credentials:
from auth.google_auth import GoogleAuthenticationError
raise GoogleAuthenticationError(
f"Access denied: Cannot retrieve credentials for {user_google_email}. "
f"You can only access credentials for your authenticated account."
@@ -59,8 +225,9 @@ async def get_authenticated_google_service_oauth21(
# Check scopes
if not all(scope in credentials.scopes for scope in required_scopes):
from auth.google_auth import GoogleAuthenticationError
raise GoogleAuthenticationError(f"OAuth 2.1 credentials lack required scopes. Need: {required_scopes}, Have: {credentials.scopes}")
raise GoogleAuthenticationError(
f"OAuth 2.1 credentials lack required scopes. Need: {required_scopes}, Have: {credentials.scopes}"
)
# Build service
service = build(service_name, version, credentials=credentials)
@@ -68,8 +235,42 @@ async def get_authenticated_google_service_oauth21(
return service, user_google_email
logger = logging.getLogger(__name__)
def _remove_user_email_arg_from_docstring(docstring: str) -> str:
"""
Remove user_google_email parameter documentation from docstring.
Args:
docstring: The original function docstring
Returns:
Modified docstring with user_google_email parameter removed
"""
if not docstring:
return docstring
# Pattern to match user_google_email parameter documentation
# Handles various formats like:
# - user_google_email (str): The user's Google email address. Required.
# - user_google_email: Description
# - user_google_email (str) - Description
patterns = [
r'^\s*user_google_email\s*\([^)]*\)\s*:\s*[^\n]*\.?\s*(?:Required\.?)?\s*\n',
r'^\s*user_google_email\s*:\s*[^\n]*\n',
r'^\s*user_google_email\s*\([^)]*\)\s*-\s*[^\n]*\n',
]
modified_docstring = docstring
for pattern in patterns:
modified_docstring = re.sub(pattern, '', modified_docstring, flags=re.MULTILINE)
# Clean up any sequence of 3 or more newlines that might have been created
modified_docstring = re.sub(r'\n{3,}', '\n\n', modified_docstring)
return modified_docstring
# Service configuration mapping
SERVICE_CONFIGS = {
"gmail": {"service": "gmail", "version": "v1"},
@@ -81,7 +282,7 @@ SERVICE_CONFIGS = {
"forms": {"service": "forms", "version": "v1"},
"slides": {"service": "slides", "version": "v1"},
"tasks": {"service": "tasks", "version": "v1"},
"customsearch": {"service": "customsearch", "version": "v1"}
"customsearch": {"service": "customsearch", "version": "v1"},
}
@@ -93,79 +294,36 @@ SCOPE_GROUPS = {
"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,
# Tasks scopes
"tasks": TASKS_SCOPE,
"tasks_read": TASKS_READONLY_SCOPE,
# Custom Search scope
"customsearch": CUSTOM_SEARCH_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]:
@@ -185,7 +343,9 @@ def _resolve_scopes(scopes: Union[str, List[str]]) -> List[str]:
return resolved
def _handle_token_refresh_error(error: RefreshError, user_email: str, service_name: str) -> str:
def _handle_token_refresh_error(
error: RefreshError, user_email: str, service_name: str
) -> str:
"""
Handle token refresh errors gracefully, particularly expired/revoked tokens.
@@ -199,11 +359,14 @@ def _handle_token_refresh_error(error: RefreshError, user_email: str, service_na
"""
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}")
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()}"
@@ -233,7 +396,6 @@ 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.
@@ -242,7 +404,6 @@ def require_google_service(
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")
@@ -250,13 +411,13 @@ def require_google_service(
# service parameter is automatically injected
# Original authentication logic is handled automatically
"""
def decorator(func: Callable) -> Callable:
# Inspect the original function signature
original_sig = inspect.signature(func)
params = list(original_sig.parameters.values())
# The decorated function must have 'service' as its first parameter.
if not params or params[0].name != 'service':
if not params or params[0].name != "service":
raise TypeError(
f"Function '{func.__name__}' decorated with @require_google_service "
"must have 'service' as its first parameter."
@@ -274,12 +435,13 @@ def require_google_service(
# Extract user_google_email from the arguments passed to the wrapper
bound_args = wrapper_sig.bind(*args, **kwargs)
bound_args.apply_defaults()
user_google_email = bound_args.arguments.get('user_google_email')
user_google_email = bound_args.arguments.get("user_google_email")
if not user_google_email:
# This should ideally not be reached if 'user_google_email' is a required parameter
# in the function signature, but it's a good safeguard.
raise Exception("'user_google_email' parameter is required but was not found.")
raise Exception(
"'user_google_email' parameter is required but was not found."
)
# Get service configuration from the decorator's arguments
if service_type not in SERVICE_CONFIGS:
@@ -292,122 +454,80 @@ def require_google_service(
# Resolve scopes
resolved_scopes = _resolve_scopes(scopes)
# --- Service Caching and Authentication Logic (largely unchanged) ---
service = None
actual_user_email = user_google_email
try:
tool_name = func.__name__
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
# Get authentication context
authenticated_user, auth_method, mcp_session_id = _get_auth_context(
tool_name
)
if service is None:
try:
tool_name = func.__name__
# Log authentication status
logger.debug(
f"[{tool_name}] Auth: {authenticated_user or 'none'} via {auth_method or 'none'} (session: {mcp_session_id[:8] if mcp_session_id else 'none'})"
)
# SIMPLIFIED: Just get the authenticated user from the context
# The AuthInfoMiddleware has already done all the authentication checks
authenticated_user = None
auth_method = None
mcp_session_id = None
# Detect OAuth version
use_oauth21 = _detect_oauth_version(
authenticated_user, mcp_session_id, tool_name
)
try:
from fastmcp.server.dependencies import get_context
ctx = get_context()
if ctx:
# Get the authenticated user email set by AuthInfoMiddleware
authenticated_user = ctx.get_state("authenticated_user_email")
auth_method = ctx.get_state("authenticated_via")
# Override user_google_email with authenticated user when using OAuth 2.1
wrapper_params = list(wrapper_sig.parameters.keys())
user_google_email, args = _override_oauth21_user_email(
use_oauth21,
authenticated_user,
user_google_email,
args,
kwargs,
wrapper_params,
tool_name,
)
# Get session ID for logging
if hasattr(ctx, 'session_id'):
mcp_session_id = ctx.session_id
# Set FastMCP session ID in context variable for propagation
from core.context import set_fastmcp_session_id
set_fastmcp_session_id(mcp_session_id)
# Update bound_args for consistency
if use_oauth21 and authenticated_user and user_google_email == authenticated_user:
bound_args.arguments["user_google_email"] = authenticated_user
logger.debug(f"[{tool_name}] Auth from middleware: {authenticated_user} via {auth_method}")
except Exception as e:
logger.debug(f"[{tool_name}] Could not get FastMCP context: {e}")
# Authenticate service
service, actual_user_email = await _authenticate_service(
use_oauth21,
service_name,
service_version,
tool_name,
user_google_email,
resolved_scopes,
mcp_session_id,
authenticated_user,
)
except GoogleAuthenticationError as e:
logger.error(
f"[{tool_name}] GoogleAuthenticationError during authentication. "
f"Method={auth_method or 'none'}, User={authenticated_user or 'none'}, "
f"Service={service_name} v{service_version}, MCPSessionID={mcp_session_id or 'none'}: {e}"
)
# Re-raise the original error without wrapping it
raise
# Log authentication status
logger.debug(f"[{tool_name}] Auth: {authenticated_user or 'none'} via {auth_method or 'none'} (session: {mcp_session_id[:8] if mcp_session_id else 'none'})")
from auth.oauth_config import is_oauth21_enabled, get_oauth_config
# Smart OAuth version detection and fallback
use_oauth21 = False
oauth_version = "oauth20" # Default
if is_oauth21_enabled():
# OAuth 2.1 is enabled globally, check client capabilities
# Try to detect from context if this is an OAuth 2.1 capable client
config = get_oauth_config()
# Build request params from context for version detection
request_params = {}
if authenticated_user:
request_params["authenticated_user"] = authenticated_user
if mcp_session_id:
request_params["session_id"] = mcp_session_id
# Detect OAuth version based on client capabilities
oauth_version = config.detect_oauth_version(request_params)
use_oauth21 = (oauth_version == "oauth21")
logger.debug(f"[{tool_name}] OAuth version detected: {oauth_version}, will use OAuth 2.1: {use_oauth21}")
if use_oauth21:
logger.debug(f"[{tool_name}] Using OAuth 2.1 flow")
# The downstream get_authenticated_google_service_oauth21 will handle
# whether the user's token is valid for the requested resource.
# This decorator should not block the call here.
service, actual_user_email = await get_authenticated_google_service_oauth21(
service_name=service_name,
version=service_version,
tool_name=tool_name,
user_google_email=user_google_email,
required_scopes=resolved_scopes,
session_id=mcp_session_id,
auth_token_email=authenticated_user,
allow_recent_auth=False,
)
else:
# Use legacy OAuth 2.0 authentication
logger.debug(f"[{tool_name}] Using legacy OAuth 2.0 flow")
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,
session_id=mcp_session_id,
)
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:
logger.error(
f"[{tool_name}] GoogleAuthenticationError during authentication. "
f"Method={auth_method or 'none'}, User={authenticated_user or 'none'}, "
f"Service={service_name} v{service_version}, MCPSessionID={mcp_session_id or 'none'}: {e}"
)
# Re-raise the original error without wrapping it
raise
# --- Call the original function with the service object injected ---
try:
# Prepend the fetched service object to the original arguments
return await func(service, *args, **kwargs)
except RefreshError as e:
error_message = _handle_token_refresh_error(e, actual_user_email, service_name)
error_message = _handle_token_refresh_error(
e, actual_user_email, service_name
)
raise Exception(error_message)
# Set the wrapper's signature to the one without 'service'
wrapper.__signature__ = wrapper_sig
# Conditionally modify docstring to remove user_google_email parameter documentation
if is_oauth21_enabled():
logger.debug('OAuth 2.1 mode enabled, removing user_google_email from docstring')
if func.__doc__:
wrapper.__doc__ = _remove_user_email_arg_from_docstring(func.__doc__)
return wrapper
return decorator
@@ -430,6 +550,7 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
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):
@@ -438,11 +559,11 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
param_names = list(sig.parameters.keys())
user_google_email = None
if 'user_google_email' in kwargs:
user_google_email = kwargs['user_google_email']
if "user_google_email" in kwargs:
user_google_email = kwargs["user_google_email"]
else:
try:
user_email_index = param_names.index('user_google_email')
user_email_index = param_names.index("user_google_email")
if user_email_index < len(args):
user_google_email = args[user_email_index]
except ValueError:
@@ -469,46 +590,37 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
try:
tool_name = func.__name__
# SIMPLIFIED: Get authentication state from context (set by AuthInfoMiddleware)
authenticated_user = None
mcp_session_id = None
# Get authentication context
authenticated_user, _, mcp_session_id = _get_auth_context(tool_name)
try:
from fastmcp.server.dependencies import get_context
ctx = get_context()
if ctx:
authenticated_user = ctx.get_state("authenticated_user_email")
if hasattr(ctx, 'session_id'):
mcp_session_id = ctx.session_id
except Exception as e:
logger.debug(f"[{tool_name}] Could not get FastMCP context: {e}")
# Detect OAuth version (simplified for multiple services)
use_oauth21 = (
is_oauth21_enabled() and authenticated_user is not None
)
# Use the same logic as single service decorator
from auth.oauth_config import is_oauth21_enabled
# Override user_google_email with authenticated user when using OAuth 2.1
user_google_email, args = _override_oauth21_user_email(
use_oauth21,
authenticated_user,
user_google_email,
args,
kwargs,
param_names,
tool_name,
service_type,
)
if is_oauth21_enabled():
logger.debug(f"[{tool_name}] Attempting OAuth 2.1 authentication flow for {service_type}.")
service, _ = await get_authenticated_google_service_oauth21(
service_name=service_name,
version=service_version,
tool_name=tool_name,
user_google_email=user_google_email,
required_scopes=resolved_scopes,
session_id=mcp_session_id,
auth_token_email=authenticated_user,
allow_recent_auth=False,
)
else:
# If OAuth 2.1 is not enabled, always use the legacy authentication method.
logger.debug(f"[{tool_name}] Using legacy authentication flow for {service_type} (OAuth 2.1 disabled).")
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,
session_id=mcp_session_id,
)
# Authenticate service
service, _ = await _authenticate_service(
use_oauth21,
service_name,
service_version,
tool_name,
user_google_email,
resolved_scopes,
mcp_session_id,
authenticated_user,
)
# Inject service with specified parameter name
kwargs[param_name] = service
@@ -525,53 +637,15 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
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")
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."""
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
}