merge upstream
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user