refactor oauth flow to simplify and leverage google built in resources

This commit is contained in:
Taylor Wilsdon
2025-05-06 12:17:22 -04:00
parent cc54fc1b8f
commit db9452f3ad
5 changed files with 227 additions and 399 deletions

View File

@@ -6,15 +6,17 @@ This module provides MCP tools for interacting with Google Calendar API.
import datetime
import logging
import asyncio
import os # Added for os.getenv
import sys
from typing import List, Optional
from typing import List, Optional, Dict # Added Dict
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
# Use functions directly from google_auth
from auth.google_auth import get_credentials, start_auth_flow, handle_auth_callback
from auth.oauth_manager import start_oauth_flow, check_auth_status, stop_oauth_flow
# Removed: from auth.oauth_manager import start_oauth_flow, check_auth_status, stop_oauth_flow
# Configure module logger
logger = logging.getLogger(__name__)
@@ -25,147 +27,137 @@ from core.server import server
# Define Google Calendar API Scopes
CALENDAR_READONLY_SCOPE = "https://www.googleapis.com/auth/calendar.readonly"
CALENDAR_EVENTS_SCOPE = "https://www.googleapis.com/auth/calendar.events"
USERINFO_EMAIL_SCOPE = "https://www.googleapis.com/auth/userinfo.email"
OPENID_SCOPE = "openid"
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# --- Configuration Placeholders (should ideally be loaded from a central app config) ---
_client_secrets_env = os.getenv("GOOGLE_CLIENT_SECRETS")
if _client_secrets_env:
CONFIG_CLIENT_SECRETS_PATH = _client_secrets_env # User provided, assume correct path
else:
# Default to client_secret.json in the same directory as this script (gcalendar_tools.py)
_current_dir = os.path.dirname(os.path.abspath(__file__))
CONFIG_CLIENT_SECRETS_PATH = os.path.join(_current_dir, "client_secret.json")
CONFIG_PORT = int(os.getenv("OAUTH_CALLBACK_PORT", 8080))
CONFIG_REDIRECT_URI = os.getenv("OAUTH_REDIRECT_URI", f"http://localhost:{CONFIG_PORT}/callback")
# ---
# Removed duplicate logging.basicConfig(level=logging.INFO)
# logger = logging.getLogger(__name__) # This is fine, but basicConfig should be in main
async def _initiate_auth_and_get_message(user_id: str, scopes: List[str]) -> str:
"""
Initiates the Google OAuth flow and returns a message for the user.
Handles the callback internally to exchange the code for tokens.
"""
logger.info(f"Initiating auth for user '{user_id}' with scopes: {scopes}")
# This inner function is called by OAuthCallbackServer (via start_auth_flow) with code and state
def _handle_redirect_for_token_exchange(received_code: str, received_state: str):
# This function runs in the OAuthCallbackServer's thread.
# It needs access to user_id, scopes, CONFIG_CLIENT_SECRETS_PATH, CONFIG_REDIRECT_URI
# These are available via closure from the _initiate_auth_and_get_message call.
current_user_id_for_flow = user_id # Capture user_id for this specific flow instance
flow_scopes = scopes # Capture scopes for this specific flow instance
logger.info(f"OAuth callback received for user '{current_user_id_for_flow}', state '{received_state}'. Exchanging code.")
try:
full_auth_response_url = f"{CONFIG_REDIRECT_URI}?code={received_code}&state={received_state}"
authenticated_user_email, credentials = handle_auth_callback(
client_secrets_path=CONFIG_CLIENT_SECRETS_PATH,
scopes=flow_scopes, # Crucial: these must be the scopes used for auth_url generation
authorization_response=full_auth_response_url,
redirect_uri=CONFIG_REDIRECT_URI
)
# Credentials are saved by handle_auth_callback
logger.info(f"Successfully exchanged token and saved credentials for {authenticated_user_email} (flow initiated for '{current_user_id_for_flow}').")
# Optionally, could signal completion if a wait mechanism was in place.
# For "auth-then-retry", this log is the primary confirmation.
except Exception as e:
logger.error(f"Error during token exchange for user '{current_user_id_for_flow}', state '{received_state}': {e}", exc_info=True)
# Optionally, could signal error if a wait mechanism was in place.
try:
# Ensure the callback function uses the specific user_id and scopes for *this* auth attempt
# by defining it within this scope or ensuring it has access to them.
# The current closure approach for _handle_redirect_for_token_exchange handles this.
auth_url, state = await asyncio.to_thread(
start_auth_flow, # This is now the function from auth.google_auth
client_secrets_path=CONFIG_CLIENT_SECRETS_PATH,
scopes=scopes,
redirect_uri=CONFIG_REDIRECT_URI,
auto_handle_callback=True, # Server starts, browser opens
callback_function=_handle_redirect_for_token_exchange, # Receives (code, state)
port=CONFIG_PORT
)
logger.info(f"Auth flow started for user '{user_id}'. State: {state}. Advise user to visit: {auth_url}")
return (
f"ACTION REQUIRED for user '{user_id}':\n"
f"1. Please visit this URL to authorize access: {auth_url}\n"
f"2. A browser window should open automatically. Complete the authorization.\n"
f"3. After successful authorization, please **RETRY** your original command.\n\n"
f"(OAuth callback server is listening on port {CONFIG_PORT} for the redirect)."
)
except Exception as e:
logger.error(f"Failed to start the OAuth flow for user '{user_id}': {e}", exc_info=True)
return f"Error: Could not initiate authentication for user '{user_id}'. {str(e)}"
# --- Tool Implementations ---
@server.tool()
async def start_auth(user_id: str) -> str:
"""
Start the Google OAuth authentication process with automatic callback handling.
This tool provides a smoother authentication experience by automatically
opening a browser window and handling the callback process.
Args:
user_id: The unique identifier (e.g., email address) for the user.
Returns:
Instructions for completing the authentication.
Starts the Google OAuth authentication process.
The user will be prompted to visit a URL and then retry their command.
This tool is useful for pre-authentication or if other tools fail due to auth.
"""
logger.info(f"Starting OAuth authentication flow for user: {user_id}")
logger.info(f"Tool 'start_auth' invoked for user: {user_id}")
# Define desired scopes for general authentication, including userinfo
auth_scopes = list(set([
CALENDAR_READONLY_SCOPE, # Default for viewing
USERINFO_EMAIL_SCOPE,
OPENID_SCOPE
]))
return await _initiate_auth_and_get_message(user_id, auth_scopes)
# Use the Calendar readonly scope by default
# Request calendar scope, user email scope, AND openid scope
scopes = [CALENDAR_READONLY_SCOPE, "https://www.googleapis.com/auth/userinfo.email", "openid"]
# Removed auth_status tool
# Removed complete_auth tool
try:
# Start the OAuth flow with automatic callback handling
# Run synchronous function in a thread
result = await asyncio.to_thread(start_oauth_flow, user_id, scopes)
return result
except Exception as e:
logger.error(f"Error starting authentication flow: {e}")
return f"Failed to start authentication: {e}"
@server.tool() # Added decorator
async def auth_status(user_id: str) -> str:
"""
Check the status of an ongoing authentication process.
Args:
user_id: The unique identifier (e.g., email address) for the user.
Returns:
A status message about the authentication process.
"""
logger.info(f"Checking authentication status for user: {user_id}")
try:
# Check current status
# Run synchronous function in a thread
result = await asyncio.to_thread(check_auth_status, user_id)
return result
except Exception as e:
logger.error(f"Error checking authentication status: {e}")
return f"Failed to check authentication status: {e}"
@server.tool() # Added decorator
async def complete_auth(user_id: str, authorization_code: str) -> str:
"""
Completes the OAuth flow by exchanging the authorization code for credentials.
Args:
user_id: The unique identifier (e.g., email address) for the user.
authorization_code: The authorization code received from Google OAuth.
Returns:
A string indicating success or failure.
"""
logger.info(f"Attempting to complete authentication for user: {user_id}")
try:
# Get the scopes used during the initial auth request
scopes = [CALENDAR_READONLY_SCOPE] # Default to readonly scope
# Construct the full callback URL
redirect_uri = "http://localhost:8080/callback"
full_callback_url = f"{redirect_uri}?code={authorization_code}"
# Use handle_auth_callback to exchange the code for credentials
# Run synchronous function in a thread
user_email, credentials = await asyncio.to_thread(
handle_auth_callback,
client_secrets_path='client_secret.json',
scopes=scopes,
authorization_response=full_callback_url,
redirect_uri=redirect_uri
)
# Verify the user_id matches the authenticated email
if user_email.lower() != user_id.lower():
logger.warning(f"User ID mismatch: provided {user_id}, authenticated as {user_email}")
return (f"Warning: You authenticated as {user_email}, but requested credentials for {user_id}. "
f"Using authenticated email {user_email} for credentials.")
logger.info(f"Successfully completed authentication for user: {user_email}")
return f"Authentication successful! You can now use the Google Calendar tools with user: {user_email}"
except Exception as e:
logger.error(f"Error completing authentication: {e}", exc_info=True)
return f"Failed to complete authentication: {e}"
@server.tool() # Added decorator
@server.tool()
async def list_calendars(user_id: str) -> str:
"""
Lists the Google Calendars the user has access to.
Args:
user_id: The unique identifier (e.g., email address) for the user.
Returns:
A string listing the calendars or an authentication prompt.
If not authenticated, prompts the user to authenticate and retry.
"""
logger.info(f"Attempting to list calendars for user: {user_id}")
scopes = [CALENDAR_READONLY_SCOPE]
logger.debug(f"Calling get_credentials with user_id: {user_id}, scopes: {scopes}")
required_scopes = [CALENDAR_READONLY_SCOPE]
try:
# Run synchronous function in a thread
credentials = await asyncio.to_thread(get_credentials, user_id, scopes, client_secrets_path='client_secret.json')
credentials = await asyncio.to_thread(
get_credentials,
user_id,
required_scopes,
client_secrets_path=CONFIG_CLIENT_SECRETS_PATH # Use config
)
logger.debug(f"get_credentials returned: {credentials}")
except Exception as e:
logger.error(f"Error getting credentials: {e}")
return f"Failed to get credentials: {e}"
logger.error(f"Error getting credentials for {user_id}: {e}", exc_info=True)
return f"Failed to get credentials: {e}. You might need to authenticate using the 'start_auth' tool."
if not credentials or not credentials.valid:
logger.warning(f"Missing or invalid credentials for user: {user_id}")
try:
# Use the automatic flow for better user experience
# Run synchronous function in a thread
result = await asyncio.to_thread(start_oauth_flow, user_id, scopes)
return result
except Exception as e:
logger.error(f"Failed to start auth flow: {e}")
return f"Failed to start authentication flow: {e}"
logger.warning(f"Missing or invalid credentials for user '{user_id}' for list_calendars. Initiating auth.")
return await _initiate_auth_and_get_message(user_id, required_scopes)
try:
service = build('calendar', 'v3', credentials=credentials)
logger.info(f"Successfully built calendar service for user: {user_id}")
calendar_list = service.calendarList().list().execute()
calendar_list = await asyncio.to_thread(service.calendarList().list().execute)
items = calendar_list.get('items', [])
if not items:
@@ -180,14 +172,14 @@ async def list_calendars(user_id: str) -> str:
return output.strip()
except HttpError as error:
logger.error(f"An API error occurred for user {user_id}: {error}")
logger.error(f"An API error occurred for user {user_id} listing calendars: {error}", exc_info=True)
# TODO: Check error details for specific auth issues (e.g., revoked token)
return f"An API error occurred: {error}. You might need to re-authenticate."
return f"An API error occurred: {error}. You might need to re-authenticate using 'start_auth'."
except Exception as e:
logger.exception(f"An unexpected error occurred while listing calendars for {user_id}: {e}")
return f"An unexpected error occurred: {e}"
@server.tool() # Added decorator
@server.tool()
async def get_events(
user_id: str,
calendar_id: str = 'primary',
@@ -197,63 +189,52 @@ async def get_events(
) -> str:
"""
Lists events from a specified Google Calendar within a given time range.
Args:
user_id: The unique identifier (e.g., email address) for the user.
calendar_id: The ID of the calendar to fetch events from. Defaults to 'primary'.
time_min: The start time for the event query (RFC3339 format, e.g., '2025-04-27T10:00:00-04:00').
Defaults to the current time if not provided.
time_max: The end time for the event query (RFC3339 format). Optional.
max_results: The maximum number of events to return. Defaults to 25.
Returns:
A string listing the events or an authentication prompt.
If not authenticated, prompts the user to authenticate and retry.
"""
logger.info(f"Attempting to get events for user: {user_id}, calendar: {calendar_id}")
scopes = [CALENDAR_READONLY_SCOPE]
try:
# Run synchronous function in a thread
credentials = await asyncio.to_thread(get_credentials, user_id, scopes, client_secrets_path='client_secret.json')
logger.debug(f"get_credentials returned: {credentials}")
required_scopes = [CALENDAR_READONLY_SCOPE]
if not credentials or not credentials.valid:
logger.warning(f"Missing or invalid credentials for user: {user_id}")
try:
# Use the automatic flow for better user experience
# Run synchronous function in a thread
result = await asyncio.to_thread(start_oauth_flow, user_id, scopes)
return result
except Exception as e:
logger.error(f"Failed to start auth flow: {e}")
return f"Failed to start authentication flow: {e}"
try:
credentials = await asyncio.to_thread(
get_credentials,
user_id,
required_scopes,
client_secrets_path=CONFIG_CLIENT_SECRETS_PATH # Use config
)
logger.debug(f"get_credentials returned: {credentials}")
except Exception as e:
logger.error(f"Error getting credentials: {e}")
return f"Failed to get credentials: {e}"
logger.error(f"Error getting credentials for {user_id}: {e}", exc_info=True)
return f"Failed to get credentials: {e}. You might need to authenticate using the 'start_auth' tool."
if not credentials or not credentials.valid:
logger.warning(f"Missing or invalid credentials for user '{user_id}' for get_events. Initiating auth.")
return await _initiate_auth_and_get_message(user_id, required_scopes)
try:
service = build('calendar', 'v3', credentials=credentials)
logger.info(f"Successfully built calendar service for user: {user_id}")
# Set default time_min to now if not provided
if time_min is None:
now = datetime.datetime.utcnow().isoformat() + 'Z' # 'Z' indicates UTC time
now = datetime.datetime.utcnow().isoformat() + 'Z'
time_min = now
logger.info(f"Defaulting time_min to current time: {time_min}")
events_result = service.events().list(
calendarId=calendar_id,
timeMin=time_min,
timeMax=time_max,
maxResults=max_results,
singleEvents=True,
orderBy='startTime'
).execute()
events_result = await asyncio.to_thread(
service.events().list(
calendarId=calendar_id,
timeMin=time_min,
timeMax=time_max,
maxResults=max_results,
singleEvents=True,
orderBy='startTime'
).execute
)
events = events_result.get('items', [])
if not events:
return f"No upcoming events found in calendar '{calendar_id}'."
return f"No upcoming events found in calendar '{calendar_id}' for the specified period."
output = f"Here are the upcoming events for calendar '{calendar_id}':\n"
output = f"Events for calendar '{calendar_id}':\n"
for event in events:
start = event['start'].get('dateTime', event['start'].get('date'))
end = event['end'].get('dateTime', event['end'].get('date'))
@@ -264,20 +245,18 @@ async def get_events(
output += f" Start: {start}\n"
output += f" End: {end}\n"
output += f" Location: {location}\n"
logger.info(f"Successfully retrieved {len(events)} events for user: {user_id}, calendar: {calendar_id}")
return output.strip()
except HttpError as error:
logger.error(f"An API error occurred for user {user_id} getting events: {error}")
# TODO: Check error details for specific auth issues
return f"An API error occurred while fetching events: {error}. You might need to re-authenticate."
logger.error(f"An API error occurred for user {user_id} getting events: {error}", exc_info=True)
return f"An API error occurred while fetching events: {error}. You might need to re-authenticate using 'start_auth'."
except Exception as e:
logger.exception(f"An unexpected error occurred while getting events for {user_id}: {e}")
return f"An unexpected error occurred: {e}"
@server.tool() # Added decorator
@server.tool()
async def create_event(
user_id: str,
summary: str,
@@ -287,47 +266,30 @@ async def create_event(
description: Optional[str] = None,
location: Optional[str] = None,
attendees: Optional[List[str]] = None,
timezone: Optional[str] = None, # e.g., "America/New_York"
timezone: Optional[str] = None,
) -> str:
"""
Creates a new event in a specified Google Calendar.
Args:
user_id: The unique identifier (e.g., email address) for the user.
summary: The title of the event.
start_time: The start time of the event (RFC3339 format, e.g., '2025-04-28T10:00:00-04:00').
end_time: The end time of the event (RFC3339 format, e.g., '2025-04-28T11:00:00-04:00').
calendar_id: The ID of the calendar to create the event in. Defaults to 'primary'.
description: An optional description for the event.
location: An optional location for the event.
attendees: An optional list of email addresses for attendees.
timezone: The timezone for the event start/end times (e.g., "America/New_York").
If not provided, the calendar's default timezone might be used.
Returns:
A confirmation message or an authentication prompt.
If not authenticated, prompts the user to authenticate and retry.
"""
logger.info(f"Attempting to create event for user: {user_id}, calendar: {calendar_id}")
# Request write scope for creating events
scopes = [CALENDAR_EVENTS_SCOPE]
try:
# Run synchronous function in a thread
credentials = await asyncio.to_thread(get_credentials, user_id, scopes, client_secrets_path='client_secret.json')
logger.debug(f"get_credentials returned: {credentials}")
required_scopes = [CALENDAR_EVENTS_SCOPE] # Write scope needed
if not credentials or not credentials.valid:
logger.warning(f"Missing or invalid credentials for user: {user_id}")
try:
# Use the automatic flow for better user experience
# Run synchronous function in a thread
result = await asyncio.to_thread(start_oauth_flow, user_id, scopes)
return result
except Exception as e:
logger.error(f"Failed to start auth flow: {e}")
return f"Failed to start authentication flow: {e}"
try:
credentials = await asyncio.to_thread(
get_credentials,
user_id,
required_scopes,
client_secrets_path=CONFIG_CLIENT_SECRETS_PATH # Use config
)
logger.debug(f"get_credentials returned: {credentials}")
except Exception as e:
logger.error(f"Error getting credentials: {e}")
return f"Failed to get credentials: {e}"
logger.error(f"Error getting credentials for {user_id}: {e}", exc_info=True)
return f"Failed to get credentials: {e}. You might need to authenticate using the 'start_auth' tool."
if not credentials or not credentials.valid:
logger.warning(f"Missing or invalid credentials for user '{user_id}' for create_event. Initiating auth.")
return await _initiate_auth_and_get_message(user_id, required_scopes)
try:
service = build('calendar', 'v3', credentials=credentials)
@@ -335,32 +297,35 @@ async def create_event(
event_body = {
'summary': summary,
'location': location,
'description': description,
'start': {'dateTime': start_time, 'timeZone': timezone},
'end': {'dateTime': end_time, 'timeZone': timezone},
'attendees': [{'email': email} for email in attendees] if attendees else [],
'start': {'dateTime': start_time},
'end': {'dateTime': end_time},
}
# Remove None values from the event body
event_body = {k: v for k, v in event_body.items() if v is not None}
if 'attendees' in event_body and not event_body['attendees']:
del event_body['attendees'] # Don't send empty attendees list
if location:
event_body['location'] = location
if description:
event_body['description'] = description
if timezone: # Apply timezone to start and end if provided
event_body['start']['timeZone'] = timezone
event_body['end']['timeZone'] = timezone
if attendees:
event_body['attendees'] = [{'email': email} for email in attendees]
logger.debug(f"Creating event with body: {event_body}")
created_event = service.events().insert(
calendarId=calendar_id,
body=event_body
).execute()
created_event = await asyncio.to_thread(
service.events().insert(
calendarId=calendar_id,
body=event_body
).execute
)
event_link = created_event.get('htmlLink')
logger.info(f"Successfully created event for user: {user_id}, event ID: {created_event['id']}")
return f"Event created successfully! View it here: {event_link}"
except HttpError as error:
logger.error(f"An API error occurred for user {user_id} creating event: {error}")
# TODO: Check error details for specific auth issues
return f"An API error occurred while creating the event: {error}. You might need to re-authenticate."
logger.error(f"An API error occurred for user {user_id} creating event: {error}", exc_info=True)
return f"An API error occurred while creating the event: {error}. You might need to re-authenticate using 'start_auth'."
except Exception as e:
logger.exception(f"An unexpected error occurred while creating event for {user_id}: {e}")
return f"An unexpected error occurred: {e}"