Merge pull request #394 from taylorwilsdon/pr_fix_start_google_oauth

enh: Remove start_google_auth when OAuth2.1 enabled

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,
@@ -518,9 +519,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 +529,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/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