workspace-mcp 1.0.0__py3-none-any.whl → 1.0.2__py3-none-any.whl

auth/service_decorator.py

@@ -193,29 +193,37 @@ def require_google_service(
             # 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':
+            raise TypeError(
+                f"Function '{func.__name__}' decorated with @require_google_service "
+                "must have 'service' as its first parameter."
+            )
+
+        # Create a new signature for the wrapper that excludes the 'service' parameter.
+        # This is the signature that FastMCP will see.
+        wrapper_sig = original_sig.replace(parameters=params[1:])
+
         @wraps(func)
         async def wrapper(*args, **kwargs):
-            # Extract user_google_email from function parameters
-            sig = inspect.signature(func)
-            param_names = list(sig.parameters.keys())
+            # Note: `args` and `kwargs` are now the arguments for the *wrapper*,
+            # which does not include 'service'.
 
-            # Find user_google_email parameter
-            user_google_email = None
-            if 'user_google_email' in kwargs:
-                user_google_email = kwargs['user_google_email']
-            else:
-                # Look for user_google_email in positional args
-                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
+            # 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')
 
             if not user_google_email:
-                raise Exception("user_google_email parameter is required but not found")
+                # 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.")
 
-            # Get service configuration
+            # Get service configuration from the decorator's arguments
             if service_type not in SERVICE_CONFIGS:
                 raise Exception(f"Unknown service type: {service_type}")
 
@@ -226,7 +234,7 @@ def require_google_service(
             # Resolve scopes
             resolved_scopes = _resolve_scopes(scopes)
 
-            # Check cache first if enabled
+            # --- Service Caching and Authentication Logic (largely unchanged) ---
             service = None
             actual_user_email = user_google_email
 
@@ -236,7 +244,6 @@ def require_google_service(
                 if cached_result:
                     service, actual_user_email = cached_result
 
-            # If not cached, authenticate
             if service is None:
                 try:
                     tool_name = func.__name__
@@ -247,30 +254,22 @@ def require_google_service(
                         user_google_email=user_google_email,
                         required_scopes=resolved_scopes,
                     )
-
-                    # Cache the service if caching is enabled
                     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:
                     raise Exception(str(e))
 
-            # Inject service as first parameter
-            if 'service' in param_names:
-                kwargs['service'] = service
-            else:
-                # Insert service as first positional argument
-                args = (service,) + args
-
-            # Call the original function with refresh error handling
+            # --- Call the original function with the service object injected ---
             try:
-                return await func(*args, **kwargs)
+                # Prepend the fetched service object to the original arguments
+                return await func(service, *args, **kwargs)
             except RefreshError as e:
-                # Handle token refresh errors gracefully
                 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
 

core/server.py

@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import Dict, Any, Optional
+from typing import Optional
 from importlib import metadata
 
 from fastapi import Header
@@ -86,18 +86,17 @@ async def health_check(request: Request):
     """Health check endpoint for container orchestration."""
     from fastapi.responses import JSONResponse
     try:
-        version = metadata.version("google-workspace-mcp")
+        version = metadata.version("workspace-mcp")
     except metadata.PackageNotFoundError:
         version = "dev"
     return JSONResponse({
         "status": "healthy",
-        "service": "google-workspace-mcp",
+        "service": "workspace-mcp",
         "version": version,
         "transport": _current_transport_mode
     })
 
 
-# Register OAuth callback as a custom route
 @server.custom_route("/oauth2callback", methods=["GET"])
 async def oauth2_callback(request: Request) -> HTMLResponse:
     """
@@ -105,8 +104,6 @@ async def oauth2_callback(request: Request) -> HTMLResponse:
     This endpoint exchanges the authorization code for credentials and saves them.
     It then displays a success or error page to the user.
     """
-    # State is used by google-auth-library for CSRF protection and should be present.
-    # We don't need to track it ourselves in this simplified flow.
     state = request.query_params.get("state")
     code = request.query_params.get("code")
     error = request.query_params.get("error")
@@ -122,7 +119,6 @@ async def oauth2_callback(request: Request) -> HTMLResponse:
         return create_error_response(error_message)
 
     try:
-        # Use the centralized CONFIG_CLIENT_SECRETS_PATH
         client_secrets_path = CONFIG_CLIENT_SECRETS_PATH
         if not os.path.exists(client_secrets_path):
             logger.error(f"OAuth client secrets file not found at {client_secrets_path}")
@@ -207,13 +203,10 @@ async def start_google_auth(
     if not ensure_oauth_callback_available(_current_transport_mode, WORKSPACE_MCP_PORT, WORKSPACE_MCP_BASE_URI):
         raise Exception("Failed to start OAuth callback server. Please try again.")
 
-    # Use the centralized start_auth_flow from auth.google_auth
     auth_result = await start_auth_flow(
         mcp_session_id=mcp_session_id,
         user_google_email=user_google_email,
         service_name=service_name,
         redirect_uri=redirect_uri
     )
-
-    # auth_result is now a plain string, not a CallToolResult
     return auth_result

core/utils.py

@@ -160,3 +160,39 @@ def extract_office_xml_text(file_bytes: bytes, mime_type: str) -> Optional[str]:
     except Exception as e:
         logger.error(f"Failed to extract office XML text for {mime_type}: {e}", exc_info=True)
         return None
+
+import functools
+from googleapiclient.errors import HttpError
+
+def handle_http_errors(tool_name: str):
+    """
+    A decorator to handle Google API HttpErrors in a standardized way.
+
+    It wraps a tool function, catches HttpError, logs a detailed error message,
+    and raises a generic Exception with a user-friendly message.
+
+    Args:
+        tool_name (str): The name of the tool being decorated (e.g., 'list_calendars').
+                         This is used for logging purposes.
+    """
+    def decorator(func):
+        @functools.wraps(func)
+        async def wrapper(*args, **kwargs):
+            try:
+                return await func(*args, **kwargs)
+            except HttpError as error:
+                user_google_email = kwargs.get('user_google_email', 'N/A')
+                message = (
+                    f"API error in {tool_name}: {error}. "
+                    f"You might need to re-authenticate for user '{user_google_email}'. "
+                    f"LLM: Try 'start_google_auth' with the user's email and the appropriate service_name."
+                )
+                logger.error(message, exc_info=True)
+                raise Exception(message)
+            except Exception as e:
+                # Catch any other unexpected errors
+                message = f"An unexpected error occurred in {tool_name}: {e}"
+                logger.exception(message)
+                raise Exception(message)
+        return wrapper
+    return decorator