merge
This commit is contained in:
104
core/api_enablement.py
Normal file
104
core/api_enablement.py
Normal 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 ""
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user