Merge pull request #407 from chetan/fix/service-mem-leak

fix/service mem leak
This commit is contained in:
Taylor Wilsdon
2026-01-30 12:45:16 -05:00
committed by GitHub
4 changed files with 133 additions and 114 deletions

View File

@@ -784,6 +784,9 @@ def get_user_info(credentials: Credentials) -> Optional[Dict[str, Any]]:
except Exception as e:
logger.error(f"Unexpected error fetching user info: {e}")
return None
finally:
if service:
service.close()
# --- Centralized Google Service Authentication ---

View File

@@ -15,6 +15,7 @@ from dataclasses import dataclass
from fastmcp.server.auth import AccessToken
from google.oauth2.credentials import Credentials
from auth.oauth_config import is_external_oauth21_provider
logger = logging.getLogger(__name__)
@@ -743,7 +744,8 @@ def ensure_session_from_access_token(
else:
store_expiry = credentials.expiry
if email:
# Skip session storage for external OAuth 2.1 to prevent memory leak from ephemeral tokens
if email and not is_external_oauth21_provider():
try:
store = get_oauth21_session_store()
store.store_session(

View File

@@ -4,6 +4,7 @@ import logging
import re
from functools import wraps
from typing import Dict, List, Optional, Any, Callable, Union, Tuple
from contextlib import ExitStack
from google.auth.exceptions import RefreshError
from googleapiclient.discovery import build
@@ -687,6 +688,9 @@ def require_google_service(
e, actual_user_email, service_name
)
raise GoogleAuthenticationError(error_message)
finally:
if service:
service.close()
# Set the wrapper's signature to the one without 'service'
wrapper.__signature__ = wrapper_sig
@@ -757,6 +761,7 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
)
# Authenticate all services
with ExitStack() as stack:
for config in service_configs:
service_type = config["service_type"]
scopes = config["scopes"]
@@ -804,6 +809,7 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
# Inject service with specified parameter name
kwargs[param_name] = service
stack.callback(service.close)
except GoogleAuthenticationError as e:
logger.error(

View File

@@ -671,9 +671,14 @@ async def create_event(
event_body["attachments"] = []
drive_service = None
try:
drive_service = service._http and build("drive", "v3", http=service._http)
try:
drive_service = service._http and build(
"drive", "v3", http=service._http
)
except Exception as e:
logger.warning(f"Could not build Drive service for MIME type lookup: {e}")
logger.warning(
f"Could not build Drive service for MIME type lookup: {e}"
)
for att in attachments:
file_id = None
if att.startswith("https://"):
@@ -726,6 +731,9 @@ async def create_event(
"mimeType": mime_type,
}
)
finally:
if drive_service:
drive_service.close()
created_event = await asyncio.to_thread(
lambda: service.events()
.insert(