From 510f9150cb8747b75506323f1dfc673ead8a48d5 Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Fri, 20 Mar 2026 14:48:31 +0530
Subject: [PATCH] feat: improve error handling for Google Drive authentication

- Enhanced error handling in both Composio and Google Drive folder listing functions to check for expired authentication and mark connectors as 'auth_expired' when necessary.
- Added logging for failed persistence of authentication status and raised appropriate HTTP exceptions to prompt users for re-authentication.
- Streamlined the error checking process to include various authentication failure scenarios, improving overall robustness of the integration.
---
 .../app/routes/composio_routes.py             | 19 ++++++++++++
 .../google_drive_add_connector_route.py       | 31 +++++++++++++++++--
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/surfsense_backend/app/routes/composio_routes.py b/surfsense_backend/app/routes/composio_routes.py
index 25250e16a..cbcb96853 100644
--- a/surfsense_backend/app/routes/composio_routes.py
+++ b/surfsense_backend/app/routes/composio_routes.py
@@ -686,6 +686,25 @@ async def list_composio_drive_folders(
         items, error = await list_folder_contents(drive_client, parent_id=parent_id)
 
         if error:
+            error_lower = error.lower()
+            if (
+                "401" in error
+                or "invalid_grant" in error_lower
+                or "token has been expired or revoked" in error_lower
+                or "invalid credentials" in error_lower
+                or "authentication failed" in error_lower
+            ):
+                try:
+                    if connector and not connector.config.get("auth_expired"):
+                        connector.config = {**connector.config, "auth_expired": True}
+                        flag_modified(connector, "config")
+                        await session.commit()
+                        logger.info(f"Marked Composio connector {connector_id} as auth_expired")
+                except Exception:
+                    logger.warning(f"Failed to persist auth_expired for connector {connector_id}", exc_info=True)
+                raise HTTPException(
+                    status_code=400, detail="Google Drive authentication expired. Please re-authenticate."
+                )
             raise HTTPException(
                 status_code=500, detail=f"Failed to list folder contents: {error}"
             )
diff --git a/surfsense_backend/app/routes/google_drive_add_connector_route.py b/surfsense_backend/app/routes/google_drive_add_connector_route.py
index 508fcf63f..2fef58fec 100644
--- a/surfsense_backend/app/routes/google_drive_add_connector_route.py
+++ b/surfsense_backend/app/routes/google_drive_add_connector_route.py
@@ -503,11 +503,31 @@ async def list_google_drive_folders(
         items, error = await list_folder_contents(drive_client, parent_id=parent_id)
 
         if error:
+            error_lower = error.lower()
+            if (
+                "401" in error
+                or "invalid_grant" in error_lower
+                or "token has been expired or revoked" in error_lower
+                or "invalid credentials" in error_lower
+                or "authentication failed" in error_lower
+            ):
+                from sqlalchemy.orm.attributes import flag_modified
+
+                try:
+                    if connector and not connector.config.get("auth_expired"):
+                        connector.config = {**connector.config, "auth_expired": True}
+                        flag_modified(connector, "config")
+                        await session.commit()
+                        logger.info(f"Marked connector {connector_id} as auth_expired")
+                except Exception:
+                    logger.warning(f"Failed to persist auth_expired for connector {connector_id}", exc_info=True)
+                raise HTTPException(
+                    status_code=400, detail="Google Drive authentication expired. Please re-authenticate."
+                )
             raise HTTPException(
                 status_code=500, detail=f"Failed to list folder contents: {error}"
             )
 
-        # Count folders and files for better logging
         folder_count = sum(1 for item in items if item.get("isFolder", False))
         file_count = len(items) - folder_count
 
@@ -516,7 +536,6 @@ async def list_google_drive_folders(
             + (f" in folder {parent_id}" if parent_id else " in ROOT")
         )
 
-        # Log first few items for debugging
         if items:
             logger.info(f"First 3 items: {[item.get('name') for item in items[:3]]}")
 
@@ -527,7 +546,13 @@ async def list_google_drive_folders(
     except Exception as e:
         logger.error(f"Error listing Drive contents: {e!s}", exc_info=True)
         error_lower = str(e).lower()
-        if "invalid_grant" in error_lower or "token has been expired or revoked" in error_lower or "authentication failed" in error_lower:
+        if (
+            "401" in str(e)
+            or "invalid_grant" in error_lower
+            or "token has been expired or revoked" in error_lower
+            or "invalid credentials" in error_lower
+            or "authentication failed" in error_lower
+        ):
             from sqlalchemy.orm.attributes import flag_modified
 
             try: