minimal fastmcp based oauth working
This commit is contained in:
Taylor Wilsdon
216
auth/service_decorator_oauth21.py.bak
Normal file
216
auth/service_decorator_oauth21.py.bak
Normal file
@@ -0,0 +1,216 @@
|
||||
"""
|
||||
Enhanced Service Decorator with OAuth 2.1 Support
|
||||
|
||||
This module provides an enhanced version of the service decorator that can
|
||||
extract and use OAuth 2.1 session context from FastMCP.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
import logging
|
||||
from functools import wraps
|
||||
from typing import Dict, List, Optional, Any, Callable, Union
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from google.auth.exceptions import RefreshError
|
||||
|
||||
from auth.service_decorator import (
|
||||
SERVICE_CONFIGS,
|
||||
SCOPE_GROUPS,
|
||||
_resolve_scopes,
|
||||
_get_cache_key,
|
||||
_is_cache_valid,
|
||||
_handle_token_refresh_error,
|
||||
_get_cached_service,
|
||||
_cache_service,
|
||||
GoogleAuthenticationError,
|
||||
)
|
||||
from auth.oauth21_integration import get_authenticated_google_service_oauth21
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _extract_context_from_args(args: tuple, kwargs: dict, sig: inspect.Signature) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Extract FastMCP Context from function arguments.
|
||||
|
||||
Args:
|
||||
args: Positional arguments
|
||||
kwargs: Keyword arguments
|
||||
sig: Function signature
|
||||
|
||||
Returns:
|
||||
Context information if found
|
||||
"""
|
||||
param_names = list(sig.parameters.keys())
|
||||
|
||||
# Check for Context type annotation
|
||||
for param_name, param in sig.parameters.items():
|
||||
if param.annotation and "Context" in str(param.annotation):
|
||||
# Found Context parameter
|
||||
if param_name in kwargs:
|
||||
ctx = kwargs[param_name]
|
||||
else:
|
||||
try:
|
||||
param_index = param_names.index(param_name)
|
||||
if param_index < len(args):
|
||||
ctx = args[param_index]
|
||||
else:
|
||||
continue
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
# Extract relevant information from Context
|
||||
if ctx:
|
||||
context_info = {}
|
||||
|
||||
# Try to get session_id
|
||||
if hasattr(ctx, "session_id"):
|
||||
context_info["session_id"] = ctx.session_id
|
||||
|
||||
# Try to get request object
|
||||
if hasattr(ctx, "request"):
|
||||
context_info["request"] = ctx.request
|
||||
|
||||
# Try to get auth context from request state
|
||||
if hasattr(ctx, "request") and hasattr(ctx.request, "state"):
|
||||
if hasattr(ctx.request.state, "auth"):
|
||||
context_info["auth_context"] = ctx.request.state.auth
|
||||
|
||||
return context_info if context_info else None
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def require_google_service_oauth21(
|
||||
service_type: str,
|
||||
scopes: Union[str, List[str]],
|
||||
version: Optional[str] = None,
|
||||
cache_enabled: bool = True,
|
||||
fallback_to_legacy: bool = True
|
||||
):
|
||||
"""
|
||||
Enhanced decorator that injects authenticated Google service with OAuth 2.1 support.
|
||||
|
||||
This decorator checks for FastMCP Context in the function parameters and uses
|
||||
OAuth 2.1 session information if available, otherwise falls back to legacy auth.
|
||||
|
||||
Args:
|
||||
service_type: Type of Google service (e.g., 'gmail', 'drive')
|
||||
scopes: Required scopes or scope aliases
|
||||
version: API version (optional, uses default if not specified)
|
||||
cache_enabled: Whether to cache service instances
|
||||
fallback_to_legacy: Whether to fall back to legacy auth if OAuth 2.1 fails
|
||||
|
||||
Usage:
|
||||
@require_google_service_oauth21("gmail", "gmail_read")
|
||||
async def search_emails(service, user_google_email: str, ctx: Context):
|
||||
# service is automatically injected
|
||||
# ctx provides OAuth 2.1 session context
|
||||
"""
|
||||
def decorator(func: Callable) -> Callable:
|
||||
# Get service configuration
|
||||
if service_type not in SERVICE_CONFIGS:
|
||||
raise ValueError(f"Unknown service type: {service_type}")
|
||||
|
||||
service_config = SERVICE_CONFIGS[service_type]
|
||||
service_name = service_config["service"]
|
||||
service_version = version or service_config["version"]
|
||||
|
||||
# Resolve scopes
|
||||
resolved_scopes = _resolve_scopes(scopes)
|
||||
|
||||
# Create wrapper with modified signature
|
||||
sig = inspect.signature(func)
|
||||
params = list(sig.parameters.values())
|
||||
|
||||
# Remove 'service' parameter from signature
|
||||
wrapper_params = [p for p in params if p.name != 'service']
|
||||
wrapper_sig = sig.replace(parameters=wrapper_params)
|
||||
|
||||
@wraps(func)
|
||||
async def wrapper(*args, **kwargs):
|
||||
# Extract user_google_email
|
||||
user_google_email = None
|
||||
if 'user_google_email' in kwargs:
|
||||
user_google_email = kwargs['user_google_email']
|
||||
else:
|
||||
param_names = list(sig.parameters.keys())
|
||||
try:
|
||||
user_email_index = param_names.index('user_google_email')
|
||||
if user_email_index < len(args):
|
||||
user_google_email = args[user_email_index]
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
if not user_google_email:
|
||||
raise ValueError("user_google_email parameter is required")
|
||||
|
||||
# Extract context information
|
||||
context = _extract_context_from_args(args, kwargs, sig)
|
||||
|
||||
service = None
|
||||
actual_user_email = user_google_email
|
||||
|
||||
# Check cache if enabled
|
||||
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
|
||||
logger.debug(f"Using cached service for {user_google_email}")
|
||||
|
||||
if service is None:
|
||||
try:
|
||||
tool_name = func.__name__
|
||||
|
||||
# Try OAuth 2.1 authentication with context
|
||||
if context:
|
||||
logger.debug(f"Attempting OAuth 2.1 authentication for {tool_name}")
|
||||
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,
|
||||
context=context,
|
||||
)
|
||||
elif fallback_to_legacy:
|
||||
# Fall back to legacy authentication
|
||||
logger.debug(f"Using legacy authentication for {tool_name}")
|
||||
from auth.google_auth import get_authenticated_google_service
|
||||
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,
|
||||
)
|
||||
else:
|
||||
raise GoogleAuthenticationError(
|
||||
"OAuth 2.1 context required but not found"
|
||||
)
|
||||
|
||||
# Cache the service if enabled
|
||||
if cache_enabled and service:
|
||||
cache_key = _get_cache_key(user_google_email, service_name, service_version, resolved_scopes)
|
||||
_cache_service(cache_key, service, actual_user_email)
|
||||
|
||||
except GoogleAuthenticationError as e:
|
||||
raise Exception(str(e))
|
||||
|
||||
# Call the original function with the service object injected
|
||||
try:
|
||||
return await func(service, *args, **kwargs)
|
||||
except RefreshError as e:
|
||||
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
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
# Alias for backward compatibility
|
||||
require_google_service = require_google_service_oauth21
|
||||
Reference in New Issue
Block a user