diff --git a/autogpt_platform/backend/backend/api/rest_api.py b/autogpt_platform/backend/backend/api/rest_api.py index 9b8ac890fd..9c24f43ce6 100644 --- a/autogpt_platform/backend/backend/api/rest_api.py +++ b/autogpt_platform/backend/backend/api/rest_api.py @@ -46,6 +46,7 @@ from backend.integrations.providers import ProviderName from backend.monitoring.instrumentation import instrument_fastapi from backend.util import json from backend.util.cloud_storage import shutdown_cloud_storage_handler +from backend.util.workspace_storage import shutdown_workspace_storage from backend.util.exceptions import ( MissingConfigError, NotAuthorizedError, @@ -125,6 +126,11 @@ async def lifespan_context(app: fastapi.FastAPI): except Exception as e: logger.warning(f"Error shutting down cloud storage handler: {e}") + try: + await shutdown_workspace_storage() + except Exception as e: + logger.warning(f"Error shutting down workspace storage: {e}") + await backend.data.db.disconnect() diff --git a/autogpt_platform/backend/backend/util/workspace_storage.py b/autogpt_platform/backend/backend/util/workspace_storage.py index ef0ce7f20d..d1f3cd5dc7 100644 --- a/autogpt_platform/backend/backend/util/workspace_storage.py +++ b/autogpt_platform/backend/backend/util/workspace_storage.py @@ -455,6 +455,23 @@ async def get_workspace_storage() -> WorkspaceStorageBackend: return _workspace_storage +async def shutdown_workspace_storage() -> None: + """ + Properly shutdown the global workspace storage backend. + + Closes aiohttp sessions and other resources for GCS backend. + Should be called during application shutdown. + """ + global _workspace_storage + + if _workspace_storage is not None: + async with _storage_lock: + if _workspace_storage is not None: + if isinstance(_workspace_storage, GCSWorkspaceStorage): + await _workspace_storage.close() + _workspace_storage = None + + def compute_file_checksum(content: bytes) -> str: """Compute SHA256 checksum of file content.""" return hashlib.sha256(content).hexdigest()