Merge branch 'main' of github.com:taylorwilsdon/google_workspace_mcp into feat/contacts-api

core/server.py

@@ -13,6 +13,7 @@ from fastmcp.server.auth.providers.google import GoogleProvider
 
 from auth.oauth21_session_store import get_oauth21_session_store, set_auth_provider
 from auth.google_auth import handle_auth_callback, start_auth_flow, check_client_secrets
+from auth.oauth_config import is_oauth21_enabled, is_external_oauth21_provider
 from auth.mcp_session_middleware import MCPSessionMiddleware
 from auth.oauth_responses import (
     create_error_response,
@@ -358,15 +359,17 @@ def configure_server_for_http():
                     base_url=config.get_oauth_base_url(),
                     redirect_path=config.redirect_path,
                     required_scopes=required_scopes,
+                    resource_server_url=config.get_oauth_base_url(),
                 )
-                # Disable protocol-level auth, expect bearer tokens in tool calls
-                server.auth = None
-                logger.info(
-                    "OAuth 2.1 enabled with EXTERNAL provider mode - protocol-level auth disabled"
-                )
+                server.auth = provider
+
+                logger.info("OAuth 2.1 enabled with EXTERNAL provider mode")
                 logger.info(
                     "Expecting Authorization bearer tokens in tool call headers"
                 )
+                logger.info(
+                    "Protected resource metadata points to Google's authorization server"
+                )
             else:
                 # Standard OAuth 2.1 mode: use FastMCP's GoogleProvider
                 provider = GoogleProvider(
@@ -384,6 +387,18 @@ def configure_server_for_http():
                     "OAuth 2.1 enabled using FastMCP GoogleProvider with protocol-level auth"
                 )
 
+                # Explicitly mount well-known routes from the OAuth provider
+                # These should be auto-mounted but we ensure they're available
+                try:
+                    well_known_routes = provider.get_well_known_routes()
+                    for route in well_known_routes:
+                        logger.info(f"Mounting OAuth well-known route: {route.path}")
+                        server.custom_route(route.path, methods=list(route.methods))(
+                            route.endpoint
+                        )
+                except Exception as e:
+                    logger.warning(f"Could not mount well-known routes: {e}")
+
             # Always set auth provider for token validation in middleware
             set_auth_provider(provider)
             _auth_provider = provider
@@ -518,9 +533,9 @@ async def start_google_auth(
     """
     Manually initiate Google OAuth authentication flow.
 
-    NOTE: This tool should typically NOT be called directly. The authentication system
-    automatically handles credential checks and prompts for authentication when needed.
-    Only use this tool if:
+    NOTE: This is a legacy OAuth 2.0 tool and is disabled when OAuth 2.1 is enabled.
+    The authentication system automatically handles credential checks and prompts for
+    authentication when needed. Only use this tool if:
     1. You need to re-authenticate with different credentials
     2. You want to proactively authenticate before using other tools
     3. The automatic authentication flow failed and you need to retry
@@ -528,6 +543,19 @@ async def start_google_auth(
     In most cases, simply try calling the Google Workspace tool you need - it will
     automatically handle authentication if required.
     """
+    if is_oauth21_enabled():
+        if is_external_oauth21_provider():
+            return (
+                "start_google_auth is disabled when OAuth 2.1 is enabled. "
+                "Provide a valid OAuth 2.1 bearer token in the Authorization header "
+                "and retry the original tool."
+            )
+        return (
+            "start_google_auth is disabled when OAuth 2.1 is enabled. "
+            "Authenticate through your MCP client's OAuth 2.1 flow and retry the "
+            "original tool."
+        )
+
     if not user_google_email:
         raise ValueError("user_google_email must be provided.")
 

core/tool_registry.py

@@ -8,6 +8,8 @@ based on tier configuration, replacing direct @server.tool() decorators.
 import logging
 from typing import Set, Optional, Callable
 
+from auth.oauth_config import is_oauth21_enabled
+
 logger = logging.getLogger(__name__)
 
 # Global registry of enabled tools
@@ -79,7 +81,8 @@ def wrap_server_tool_method(server):
 def filter_server_tools(server):
     """Remove disabled tools from the server after registration."""
     enabled_tools = get_enabled_tools()
-    if enabled_tools is None:
+    oauth21_enabled = is_oauth21_enabled()
+    if enabled_tools is None and not oauth21_enabled:
         return
 
     tools_removed = 0
@@ -90,16 +93,25 @@ def filter_server_tools(server):
         if hasattr(tool_manager, "_tools"):
             tool_registry = tool_manager._tools
 
-            tools_to_remove = []
-            for tool_name in list(tool_registry.keys()):
-                if not is_tool_enabled(tool_name):
-                    tools_to_remove.append(tool_name)
+            tools_to_remove = set()
+            if enabled_tools is not None:
+                for tool_name in list(tool_registry.keys()):
+                    if not is_tool_enabled(tool_name):
+                        tools_to_remove.add(tool_name)
+
+            if oauth21_enabled and "start_google_auth" in tool_registry:
+                tools_to_remove.add("start_google_auth")
 
             for tool_name in tools_to_remove:
                 del tool_registry[tool_name]
                 tools_removed += 1
+                if tool_name == "start_google_auth":
+                    logger.info("OAuth 2.1 enabled: disabling start_google_auth tool")
 
     if tools_removed > 0:
+        enabled_count = len(enabled_tools) if enabled_tools is not None else "all"
         logger.info(
-            f"Tool tier filtering: removed {tools_removed} tools, {len(enabled_tools)} enabled"
+            "Tool filtering: removed %s tools, %s enabled",
+            tools_removed,
+            enabled_count,
         )

core/tool_tiers.yaml

@@ -31,6 +31,7 @@ drive:
     - get_drive_shareable_link
   extended:
     - list_drive_items
+    - copy_drive_file
     - update_drive_file
     - update_drive_permission
     - remove_drive_permission
@@ -48,6 +49,7 @@ calendar:
     - modify_event
   extended:
     - delete_event
+    - query_freebusy
   complete: []
 
 docs:
@@ -106,6 +108,7 @@ forms:
   complete:
     - set_publish_settings
     - get_form_response
+    - batch_update_form
 
 slides:
   core:

core/utils.py

@@ -12,6 +12,7 @@ from typing import List, Optional
 from googleapiclient.errors import HttpError
 from .api_enablement import get_api_enablement_message
 from auth.google_auth import GoogleAuthenticationError
+from auth.oauth_config import is_oauth21_enabled, is_external_oauth21_provider
 
 logger = logging.getLogger(__name__)
 
@@ -314,10 +315,26 @@ def handle_http_errors(
                             )
                     elif error.resp.status in [401, 403]:
                         # Authentication/authorization errors
+                        if is_oauth21_enabled():
+                            if is_external_oauth21_provider():
+                                auth_hint = (
+                                    "LLM: Ask the user to provide a valid OAuth 2.1 "
+                                    "bearer token in the Authorization header and retry."
+                                )
+                            else:
+                                auth_hint = (
+                                    "LLM: Ask the user to authenticate via their MCP "
+                                    "client's OAuth 2.1 flow and retry."
+                                )
+                        else:
+                            auth_hint = (
+                                "LLM: Try 'start_google_auth' with the user's email "
+                                "and the appropriate service_name."
+                            )
                         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."
+                            f"{auth_hint}"
                         )
                     else:
                         # Other HTTP errors (400 Bad Request, etc.) - don't suggest re-auth