fully working mcp session id handling

This commit is contained in:
Taylor Wilsdon
2025-05-22 17:02:00 -04:00
parent 07930d6004
commit ab029b2b72
4 changed files with 192 additions and 21 deletions

View File

@@ -13,10 +13,13 @@ from google_auth_oauthlib.flow import Flow, InstalledAppFlow
from google.auth.transport.requests import Request
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from config.google_config import OAUTH_STATE_TO_SESSION_ID_MAP, SCOPES
from mcp import types
# Import our session ID getter
from core.streamable_http import get_current_session_id
# Configure logging
logging.basicConfig(level=logging.INFO)
@@ -72,7 +75,7 @@ def save_credentials_to_file(user_google_email: str, credentials: Credentials, b
def save_credentials_to_session(session_id: str, credentials: Credentials):
"""Saves user credentials to the in-memory session cache."""
_SESSION_CREDENTIALS_CACHE[session_id] = credentials
logger.info(f"Credentials saved to session cache for session_id: {session_id}")
logger.debug(f"Credentials saved to session cache for session_id: {session_id}")
def load_credentials_from_file(user_google_email: str, base_dir: str = DEFAULT_CREDENTIALS_DIR) -> Optional[Credentials]:
"""Loads user credentials from a file."""
@@ -92,7 +95,7 @@ def load_credentials_from_file(user_google_email: str, base_dir: str = DEFAULT_C
client_secret=creds_data.get('client_secret'),
scopes=creds_data.get('scopes')
)
logger.info(f"Credentials loaded for user {user_google_email} from {creds_path}")
logger.debug(f"Credentials loaded for user {user_google_email} from {creds_path}")
return credentials
except (IOError, json.JSONDecodeError, KeyError) as e:
logger.error(f"Error loading or parsing credentials for user {user_google_email} from {creds_path}: {e}")
@@ -102,9 +105,9 @@ def load_credentials_from_session(session_id: str) -> Optional[Credentials]:
"""Loads user credentials from the in-memory session cache."""
credentials = _SESSION_CREDENTIALS_CACHE.get(session_id)
if credentials:
logger.info(f"Credentials loaded from session cache for session_id: {session_id}")
logger.debug(f"Credentials loaded from session cache for session_id: {session_id}")
else:
logger.info(f"No credentials found in session cache for session_id: {session_id}")
logger.debug(f"No credentials found in session cache for session_id: {session_id}")
return credentials
def load_client_secrets(client_secrets_path: str) -> Dict[str, Any]:
@@ -297,36 +300,43 @@ def get_credentials(
"""
credentials: Optional[Credentials] = None
loaded_from_session = False
# Try to get the current session ID if not explicitly provided
if not session_id:
current_session_id = get_current_session_id()
if current_session_id:
session_id = current_session_id
logger.info(f"[get_credentials] No session_id provided, using current session ID: '{session_id}'")
logger.info(f"[get_credentials] Called for user_google_email: '{user_google_email}', session_id: '{session_id}', required_scopes: {required_scopes}")
logger.debug(f"[get_credentials] Called for user_google_email: '{user_google_email}', session_id: '{session_id}', required_scopes: {required_scopes}")
if session_id:
credentials = load_credentials_from_session(session_id)
if credentials:
logger.info(f"[get_credentials] Loaded credentials from session for session_id '{session_id}'.")
logger.debug(f"[get_credentials] Loaded credentials from session for session_id '{session_id}'.")
loaded_from_session = True
if not credentials and user_google_email:
logger.info(f"[get_credentials] No session credentials, trying file for user_google_email '{user_google_email}'.")
logger.debug(f"[get_credentials] No session credentials, trying file for user_google_email '{user_google_email}'.")
credentials = load_credentials_from_file(user_google_email, credentials_base_dir)
if credentials and session_id:
logger.info(f"[get_credentials] Loaded from file for user '{user_google_email}', caching to session '{session_id}'.")
logger.debug(f"[get_credentials] Loaded from file for user '{user_google_email}', caching to session '{session_id}'.")
save_credentials_to_session(session_id, credentials) # Cache for current session
if not credentials:
logger.info(f"[get_credentials] No credentials found for user '{user_google_email}' or session '{session_id}'.")
return None
logger.info(f"[get_credentials] Credentials found. Scopes: {credentials.scopes}, Valid: {credentials.valid}, Expired: {credentials.expired}")
logger.debug(f"[get_credentials] Credentials found. Scopes: {credentials.scopes}, Valid: {credentials.valid}, Expired: {credentials.expired}")
if not all(scope in credentials.scopes for scope in required_scopes):
logger.warning(f"[get_credentials] Credentials lack required scopes. Need: {required_scopes}, Have: {credentials.scopes}. User: '{user_google_email}', Session: '{session_id}'")
return None # Re-authentication needed for scopes
logger.info(f"[get_credentials] Credentials have sufficient scopes. User: '{user_google_email}', Session: '{session_id}'")
logger.debug(f"[get_credentials] Credentials have sufficient scopes. User: '{user_google_email}', Session: '{session_id}'")
if credentials.valid:
logger.info(f"[get_credentials] Credentials are valid. User: '{user_google_email}', Session: '{session_id}'")
logger.debug(f"[get_credentials] Credentials are valid. User: '{user_google_email}', Session: '{session_id}'")
return credentials
elif credentials.expired and credentials.refresh_token:
logger.info(f"[get_credentials] Credentials expired. Attempting refresh. User: '{user_google_email}', Session: '{session_id}'")
@@ -334,7 +344,7 @@ def get_credentials(
logger.error("[get_credentials] Client secrets path required for refresh but not provided.")
return None
try:
logger.info(f"[get_credentials] Refreshing token using client_secrets_path: {client_secrets_path}")
logger.debug(f"[get_credentials] Refreshing token using client_secrets_path: {client_secrets_path}")
# client_config = load_client_secrets(client_secrets_path) # Not strictly needed if creds have client_id/secret
credentials.refresh(Request())
logger.info(f"[get_credentials] Credentials refreshed successfully. User: '{user_google_email}', Session: '{session_id}'")