This commit is contained in:
Taylor Wilsdon
2025-08-01 13:19:09 -04:00
25 changed files with 1648 additions and 591 deletions

104
core/api_enablement.py Normal file
View File

@@ -0,0 +1,104 @@
import re
from typing import Dict, Optional, Tuple
API_ENABLEMENT_LINKS: Dict[str, str] = {
"calendar-json.googleapis.com": "https://console.cloud.google.com/flows/enableapi?apiid=calendar-json.googleapis.com",
"drive.googleapis.com": "https://console.cloud.google.com/flows/enableapi?apiid=drive.googleapis.com",
"gmail.googleapis.com": "https://console.cloud.google.com/flows/enableapi?apiid=gmail.googleapis.com",
"docs.googleapis.com": "https://console.cloud.google.com/flows/enableapi?apiid=docs.googleapis.com",
"sheets.googleapis.com": "https://console.cloud.google.com/flows/enableapi?apiid=sheets.googleapis.com",
"slides.googleapis.com": "https://console.cloud.google.com/flows/enableapi?apiid=slides.googleapis.com",
"forms.googleapis.com": "https://console.cloud.google.com/flows/enableapi?apiid=forms.googleapis.com",
"tasks.googleapis.com": "https://console.cloud.google.com/flows/enableapi?apiid=tasks.googleapis.com",
"chat.googleapis.com": "https://console.cloud.google.com/flows/enableapi?apiid=chat.googleapis.com",
"customsearch.googleapis.com": "https://console.cloud.google.com/flows/enableapi?apiid=customsearch.googleapis.com",
}
SERVICE_NAME_TO_API: Dict[str, str] = {
"Google Calendar": "calendar-json.googleapis.com",
"Google Drive": "drive.googleapis.com",
"Gmail": "gmail.googleapis.com",
"Google Docs": "docs.googleapis.com",
"Google Sheets": "sheets.googleapis.com",
"Google Slides": "slides.googleapis.com",
"Google Forms": "forms.googleapis.com",
"Google Tasks": "tasks.googleapis.com",
"Google Chat": "chat.googleapis.com",
"Google Custom Search": "customsearch.googleapis.com",
}
INTERNAL_SERVICE_TO_API: Dict[str, str] = {
"calendar": "calendar-json.googleapis.com",
"drive": "drive.googleapis.com",
"gmail": "gmail.googleapis.com",
"docs": "docs.googleapis.com",
"sheets": "sheets.googleapis.com",
"slides": "slides.googleapis.com",
"forms": "forms.googleapis.com",
"tasks": "tasks.googleapis.com",
"chat": "chat.googleapis.com",
"customsearch": "customsearch.googleapis.com",
"search": "customsearch.googleapis.com",
}
def extract_api_info_from_error(error_details: str) -> Tuple[Optional[str], Optional[str]]:
"""
Extract API service and project ID from error details.
Returns:
Tuple of (api_service, project_id) or (None, None) if not found
"""
api_pattern = r'https://console\.developers\.google\.com/apis/api/([^/]+)/overview'
project_pattern = r'project[=\s]+([a-zA-Z0-9-]+)'
api_match = re.search(api_pattern, error_details)
project_match = re.search(project_pattern, error_details)
api_service = api_match.group(1) if api_match else None
project_id = project_match.group(1) if project_match else None
return api_service, project_id
def get_api_enablement_message(error_details: str, service_type: Optional[str] = None) -> str:
"""
Generate a helpful error message with direct API enablement link.
Args:
error_details: The error details string from the HttpError
service_type: Optional service type (e.g., "calendar", "gmail", or "Google Calendar")
Returns:
Formatted error message with enablement link
"""
api_service, project_id = extract_api_info_from_error(error_details)
if not api_service and service_type:
# Check internal service names first (e.g., "calendar", "gmail")
api_service = INTERNAL_SERVICE_TO_API.get(service_type)
if not api_service:
# Check display names (e.g., "Google Calendar")
api_service = SERVICE_NAME_TO_API.get(service_type)
if api_service and api_service in API_ENABLEMENT_LINKS:
enable_link = API_ENABLEMENT_LINKS[api_service]
service_display_name = next(
(name for name, api in SERVICE_NAME_TO_API.items() if api == api_service),
api_service
)
message = (
f"{service_display_name} API is not enabled for your project"
f"{f' ({project_id})' if project_id else ''}.\n\n"
f"Enable it here: {enable_link}\n\n"
f"After enabling, wait 1-2 minutes for the change to propagate, then try again."
f"IMPORTANT - LLM: share the link provided as a clickable hyperlink and instruct the user to enable the required API."
)
return message
return ""

View File

@@ -37,75 +37,75 @@ def create_comment_tools(app_name: str, file_id_param: str):
# Create functions without decorators first, then apply decorators with proper names
if file_id_param == "document_id":
@require_google_service("drive", "drive_read")
@handle_http_errors(read_func_name)
@handle_http_errors(read_func_name, service_type="drive")
async def read_comments(service, user_google_email: str, document_id: str) -> str:
"""Read all comments from a Google Document."""
return await _read_comments_impl(service, app_name, document_id)
@require_google_service("drive", "drive_file")
@handle_http_errors(create_func_name)
@handle_http_errors(create_func_name, service_type="drive")
async def create_comment(service, user_google_email: str, document_id: str, comment_content: str) -> str:
"""Create a new comment on a Google Document."""
return await _create_comment_impl(service, app_name, document_id, comment_content)
@require_google_service("drive", "drive_file")
@handle_http_errors(reply_func_name)
@handle_http_errors(reply_func_name, service_type="drive")
async def reply_to_comment(service, user_google_email: str, document_id: str, comment_id: str, reply_content: str) -> str:
"""Reply to a specific comment in a Google Document."""
return await _reply_to_comment_impl(service, app_name, document_id, comment_id, reply_content)
@require_google_service("drive", "drive_file")
@handle_http_errors(resolve_func_name)
@handle_http_errors(resolve_func_name, service_type="drive")
async def resolve_comment(service, user_google_email: str, document_id: str, comment_id: str) -> str:
"""Resolve a comment in a Google Document."""
return await _resolve_comment_impl(service, app_name, document_id, comment_id)
elif file_id_param == "spreadsheet_id":
@require_google_service("drive", "drive_read")
@handle_http_errors(read_func_name)
@handle_http_errors(read_func_name, service_type="drive")
async def read_comments(service, user_google_email: str, spreadsheet_id: str) -> str:
"""Read all comments from a Google Spreadsheet."""
return await _read_comments_impl(service, app_name, spreadsheet_id)
@require_google_service("drive", "drive_file")
@handle_http_errors(create_func_name)
@handle_http_errors(create_func_name, service_type="drive")
async def create_comment(service, user_google_email: str, spreadsheet_id: str, comment_content: str) -> str:
"""Create a new comment on a Google Spreadsheet."""
return await _create_comment_impl(service, app_name, spreadsheet_id, comment_content)
@require_google_service("drive", "drive_file")
@handle_http_errors(reply_func_name)
@handle_http_errors(reply_func_name, service_type="drive")
async def reply_to_comment(service, user_google_email: str, spreadsheet_id: str, comment_id: str, reply_content: str) -> str:
"""Reply to a specific comment in a Google Spreadsheet."""
return await _reply_to_comment_impl(service, app_name, spreadsheet_id, comment_id, reply_content)
@require_google_service("drive", "drive_file")
@handle_http_errors(resolve_func_name)
@handle_http_errors(resolve_func_name, service_type="drive")
async def resolve_comment(service, user_google_email: str, spreadsheet_id: str, comment_id: str) -> str:
"""Resolve a comment in a Google Spreadsheet."""
return await _resolve_comment_impl(service, app_name, spreadsheet_id, comment_id)
elif file_id_param == "presentation_id":
@require_google_service("drive", "drive_read")
@handle_http_errors(read_func_name)
@handle_http_errors(read_func_name, service_type="drive")
async def read_comments(service, user_google_email: str, presentation_id: str) -> str:
"""Read all comments from a Google Presentation."""
return await _read_comments_impl(service, app_name, presentation_id)
@require_google_service("drive", "drive_file")
@handle_http_errors(create_func_name)
@handle_http_errors(create_func_name, service_type="drive")
async def create_comment(service, user_google_email: str, presentation_id: str, comment_content: str) -> str:
"""Create a new comment on a Google Presentation."""
return await _create_comment_impl(service, app_name, presentation_id, comment_content)
@require_google_service("drive", "drive_file")
@handle_http_errors(reply_func_name)
@handle_http_errors(reply_func_name, service_type="drive")
async def reply_to_comment(service, user_google_email: str, presentation_id: str, comment_id: str, reply_content: str) -> str:
"""Reply to a specific comment in a Google Presentation."""
return await _reply_to_comment_impl(service, app_name, presentation_id, comment_id, reply_content)
@require_google_service("drive", "drive_file")
@handle_http_errors(resolve_func_name)
@handle_http_errors(resolve_func_name, service_type="drive")
async def resolve_comment(service, user_google_email: str, presentation_id: str, comment_id: str) -> str:
"""Resolve a comment in a Google Presentation."""
return await _resolve_comment_impl(service, app_name, presentation_id, comment_id)

View File

@@ -28,7 +28,6 @@ except ImportError as e:
# Import shared configuration
from auth.scopes import (
OAUTH_STATE_TO_SESSION_ID_MAP,
SCOPES,
USERINFO_EMAIL_SCOPE, # noqa: F401
OPENID_SCOPE, # noqa: F401
@@ -64,6 +63,8 @@ from auth.scopes import (
TASKS_SCOPE, # noqa: F401
TASKS_READONLY_SCOPE, # noqa: F401
TASKS_SCOPES, # noqa: F401
CUSTOM_SEARCH_SCOPE, # noqa: F401
CUSTOM_SEARCH_SCOPES, # noqa: F401
)
# Configure logging
@@ -80,7 +81,6 @@ _current_transport_mode = "stdio" # Default to stdio
# Basic MCP server instance
server = FastMCP(
name="google_workspace",
server_url=f"{WORKSPACE_MCP_BASE_URI}:{WORKSPACE_MCP_PORT}/mcp",
port=WORKSPACE_MCP_PORT,
host="0.0.0.0"
)
@@ -192,23 +192,16 @@ async def oauth2_callback(request: Request) -> HTMLResponse:
logger.info(f"OAuth callback: Received code (state: {state}). Attempting to exchange for tokens.")
mcp_session_id: Optional[str] = OAUTH_STATE_TO_SESSION_ID_MAP.pop(state, None)
if mcp_session_id:
logger.info(f"OAuth callback: Retrieved MCP session ID '{mcp_session_id}' for state '{state}'.")
else:
logger.warning(f"OAuth callback: No MCP session ID found for state '{state}'. Auth will not be tied to a specific session directly via this callback.")
# Exchange code for credentials. handle_auth_callback will save them.
# The user_id returned here is the Google-verified email.
verified_user_id, credentials = handle_auth_callback(
scopes=SCOPES, # Ensure all necessary scopes are requested
authorization_response=str(request.url),
redirect_uri=get_oauth_redirect_uri_for_current_mode(),
session_id=mcp_session_id # Pass session_id if available
session_id=None # Session ID tracking removed
)
log_session_part = f" (linked to session: {mcp_session_id})" if mcp_session_id else ""
logger.info(f"OAuth callback: Successfully authenticated user: {verified_user_id} (state: {state}){log_session_part}.")
logger.info(f"OAuth callback: Successfully authenticated user: {verified_user_id} (state: {state}).")
# Return success page using shared template
return create_success_response(verified_user_id)
@@ -222,14 +215,13 @@ async def oauth2_callback(request: Request) -> HTMLResponse:
@server.tool()
async def start_google_auth(
service_name: str,
user_google_email: str = USER_GOOGLE_EMAIL,
mcp_session_id: Optional[str] = Header(None, alias="Mcp-Session-Id")
user_google_email: str = USER_GOOGLE_EMAIL
) -> str:
"""
Initiates the Google OAuth 2.0 authentication flow for the specified user email and service.
This is the primary method to establish credentials when no valid session exists or when targeting a specific account for a particular service.
It generates an authorization URL that the LLM must present to the user.
The authentication attempt is linked to the current MCP session via `mcp_session_id`.
This initiates a new authentication flow for the specified user and service.
LLM Guidance:
- Use this tool when you need to authenticate a user for a specific Google service (e.g., "Google Calendar", "Google Docs", "Gmail", "Google Drive")
@@ -245,7 +237,6 @@ async def start_google_auth(
Args:
user_google_email (str): The user's full Google email address (e.g., 'example@gmail.com'). This is REQUIRED.
service_name (str): The name of the Google service for which authentication is being requested (e.g., "Google Calendar", "Google Docs"). This is REQUIRED.
mcp_session_id (Optional[str]): The active MCP session ID (automatically injected by FastMCP from the Mcp-Session-Id header). Links the OAuth flow state to the session.
Returns:
str: A detailed message for the LLM with the authorization URL and instructions to guide the user through the authentication process.
@@ -260,15 +251,18 @@ async def start_google_auth(
logger.error(f"[start_google_auth] {error_msg}")
raise Exception(error_msg)
logger.info(f"Tool 'start_google_auth' invoked for user_google_email: '{user_google_email}', service: '{service_name}', session: '{mcp_session_id}'.")
logger.info(f"Tool 'start_google_auth' invoked for user_google_email: '{user_google_email}', service: '{service_name}'.")
# Ensure OAuth callback is available for current transport mode
redirect_uri = get_oauth_redirect_uri_for_current_mode()
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.")
success, error_msg = ensure_oauth_callback_available(_current_transport_mode, WORKSPACE_MCP_PORT, WORKSPACE_MCP_BASE_URI)
if not success:
if error_msg:
raise Exception(f"Failed to start OAuth callback server: {error_msg}")
else:
raise Exception("Failed to start OAuth callback server. Please try again.")
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

View File

@@ -10,6 +10,7 @@ import functools
from typing import List, Optional
from googleapiclient.errors import HttpError
from .api_enablement import get_api_enablement_message
logger = logging.getLogger(__name__)
@@ -233,7 +234,7 @@ def extract_office_xml_text(file_bytes: bytes, mime_type: str) -> Optional[str]:
return None
def handle_http_errors(tool_name: str, is_read_only: bool = False):
def handle_http_errors(tool_name: str, is_read_only: bool = False, service_type: Optional[str] = None):
"""
A decorator to handle Google API HttpErrors and transient SSL errors in a standardized way.
@@ -247,6 +248,7 @@ def handle_http_errors(tool_name: str, is_read_only: bool = False):
tool_name (str): The name of the tool being decorated (e.g., 'list_calendars').
is_read_only (bool): If True, the operation is considered safe to retry on
transient network errors. Defaults to False.
service_type (str): Optional. The Google service type (e.g., 'calendar', 'gmail').
"""
def decorator(func):
@@ -275,12 +277,31 @@ def handle_http_errors(tool_name: str, is_read_only: bool = False):
) from e
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)
error_details = str(error)
# Check if this is an API not enabled error
if error.resp.status == 403 and "accessNotConfigured" in error_details:
enablement_msg = get_api_enablement_message(error_details, service_type)
if enablement_msg:
message = (
f"API error in {tool_name}: {enablement_msg}\n\n"
f"User: {user_google_email}"
)
else:
message = (
f"API error in {tool_name}: {error}. "
f"The required API is not enabled for your project. "
f"Please check the Google Cloud Console to enable it."
)
else:
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(f"API error in {tool_name}: {error}", exc_info=True)
raise Exception(message) from error
except TransientNetworkError:
# Re-raise without wrapping to preserve the specific error type