merge upstream

This commit is contained in:
Taylor Wilsdon
2026-01-31 13:14:27 -05:00
43 changed files with 6531 additions and 481 deletions

View File

@@ -4,6 +4,7 @@ import logging
import re
from functools import wraps
from typing import Dict, List, Optional, Any, Callable, Union, Tuple
from contextlib import ExitStack
from google.auth.exceptions import RefreshError
from googleapiclient.discovery import build
@@ -14,7 +15,11 @@ from auth.oauth21_session_store import (
get_oauth21_session_store,
ensure_session_from_access_token,
)
from auth.oauth_config import is_oauth21_enabled, get_oauth_config
from auth.oauth_config import (
is_oauth21_enabled,
get_oauth_config,
is_external_oauth21_provider,
)
from core.context import set_fastmcp_session_id
from auth.scopes import (
GMAIL_READONLY_SCOPE,
@@ -41,7 +46,13 @@ from auth.scopes import (
SLIDES_READONLY_SCOPE,
TASKS_SCOPE,
TASKS_READONLY_SCOPE,
CONTACTS_SCOPE,
CONTACTS_READONLY_SCOPE,
CUSTOM_SEARCH_SCOPE,
SCRIPT_PROJECTS_SCOPE,
SCRIPT_PROJECTS_READONLY_SCOPE,
SCRIPT_DEPLOYMENTS_SCOPE,
SCRIPT_DEPLOYMENTS_READONLY_SCOPE,
)
logger = logging.getLogger(__name__)
@@ -69,8 +80,8 @@ def _get_auth_context(
if mcp_session_id:
set_fastmcp_session_id(mcp_session_id)
logger.debug(
f"[{tool_name}] Auth from middleware: {authenticated_user} via {auth_method}"
logger.info(
f"[{tool_name}] Auth from middleware: authenticated_user={authenticated_user}, auth_method={auth_method}, session_id={mcp_session_id}"
)
return authenticated_user, auth_method, mcp_session_id
@@ -98,6 +109,19 @@ def _detect_oauth_version(
)
return True
# If FastMCP protocol-level auth is enabled, a validated access token should
# be available even if middleware state wasn't populated.
try:
if get_access_token() is not None:
logger.info(
f"[{tool_name}] OAuth 2.1 mode: Using OAuth 2.1 based on validated access token"
)
return True
except Exception as e:
logger.debug(
f"[{tool_name}] Could not inspect access token for OAuth mode: {e}"
)
# Only use version detection for unauthenticated requests
config = get_oauth_config()
request_params = {}
@@ -384,7 +408,9 @@ SERVICE_CONFIGS = {
"forms": {"service": "forms", "version": "v1"},
"slides": {"service": "slides", "version": "v1"},
"tasks": {"service": "tasks", "version": "v1"},
"people": {"service": "people", "version": "v1"},
"customsearch": {"service": "customsearch", "version": "v1"},
"script": {"service": "script", "version": "v1"},
}
@@ -423,8 +449,16 @@ SCOPE_GROUPS = {
# Tasks scopes
"tasks": TASKS_SCOPE,
"tasks_read": TASKS_READONLY_SCOPE,
# Contacts scopes
"contacts": CONTACTS_SCOPE,
"contacts_read": CONTACTS_READONLY_SCOPE,
# Custom Search scope
"customsearch": CUSTOM_SEARCH_SCOPE,
# Apps Script scopes
"script_readonly": SCRIPT_PROJECTS_READONLY_SCOPE,
"script_projects": SCRIPT_PROJECTS_SCOPE,
"script_deployments": SCRIPT_DEPLOYMENTS_SCOPE,
"script_deployments_readonly": SCRIPT_DEPLOYMENTS_READONLY_SCOPE,
}
@@ -470,6 +504,26 @@ def _handle_token_refresh_error(
)
service_display_name = f"Google {service_name.title()}"
if is_oauth21_enabled():
if is_external_oauth21_provider():
oauth21_step = (
"Provide a valid OAuth 2.1 bearer token in the Authorization header"
)
else:
oauth21_step = "Sign in through your MCP client's OAuth 2.1 flow"
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. {oauth21_step}\n"
f"2. Retry your original command\n\n"
f"The application will automatically use the new credentials once authentication is complete."
)
return (
f"**Authentication Required: Token Expired/Revoked for {service_display_name}**\n\n"
@@ -487,6 +541,16 @@ def _handle_token_refresh_error(
else:
# Handle other types of refresh errors
logger.error(f"Unexpected refresh error for user {user_email}: {error}")
if is_oauth21_enabled():
if is_external_oauth21_provider():
return (
f"Authentication error occurred for {user_email}. "
"Please provide a valid OAuth 2.1 bearer token and retry."
)
return (
f"Authentication error occurred for {user_email}. "
"Please sign in via your MCP client's OAuth 2.1 flow and retry."
)
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."
@@ -623,7 +687,10 @@ def require_google_service(
error_message = _handle_token_refresh_error(
e, actual_user_email, service_name
)
raise Exception(error_message)
raise GoogleAuthenticationError(error_message)
finally:
if service:
service.close()
# Set the wrapper's signature to the one without 'service'
wrapper.__signature__ = wrapper_sig
@@ -697,74 +764,76 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
)
# 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")
with ExitStack() as stack:
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}")
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)
service_config = SERVICE_CONFIGS[service_type]
service_name = service_config["service"]
service_version = version or service_config["version"]
resolved_scopes = _resolve_scopes(scopes)
try:
# Detect OAuth version (simplified for multiple services)
use_oauth21 = (
is_oauth21_enabled() and authenticated_user is not None
)
# In OAuth 2.0 mode, we may need to override user_google_email
if not is_oauth21_enabled():
user_google_email, args = _override_oauth21_user_email(
use_oauth21,
authenticated_user,
user_google_email,
args,
kwargs,
wrapper_param_names,
tool_name,
service_type,
try:
# Detect OAuth version (simplified for multiple services)
use_oauth21 = (
is_oauth21_enabled() and authenticated_user is not None
)
# Authenticate service
service, _ = await _authenticate_service(
use_oauth21,
service_name,
service_version,
tool_name,
user_google_email,
resolved_scopes,
mcp_session_id,
authenticated_user,
# In OAuth 2.0 mode, we may need to override user_google_email
if not is_oauth21_enabled():
user_google_email, args = _override_oauth21_user_email(
use_oauth21,
authenticated_user,
user_google_email,
args,
kwargs,
wrapper_param_names,
tool_name,
service_type,
)
# 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
stack.callback(service.close)
except GoogleAuthenticationError as e:
logger.error(
f"[{tool_name}] GoogleAuthenticationError for service '{service_type}' (user: {user_google_email}): {e}"
)
# Re-raise the original error without wrapping it
raise
# Call the original function with refresh error handling
try:
# In OAuth 2.1 mode, we need to add user_google_email to kwargs since it was removed from signature
if is_oauth21_enabled():
kwargs["user_google_email"] = user_google_email
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"
)
# Inject service with specified parameter name
kwargs[param_name] = service
except GoogleAuthenticationError as e:
logger.error(
f"[{tool_name}] GoogleAuthenticationError for service '{service_type}' (user: {user_google_email}): {e}"
)
# Re-raise the original error without wrapping it
raise
# Call the original function with refresh error handling
try:
# In OAuth 2.1 mode, we need to add user_google_email to kwargs since it was removed from signature
if is_oauth21_enabled():
kwargs["user_google_email"] = user_google_email
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)
raise GoogleAuthenticationError(error_message)
# Set the wrapper's signature
wrapper.__signature__ = wrapper_sig