feat: enhance Google Drive account authentication handling

- Added checks for expired authentication in Google Drive file creation and deletion tools, returning appropriate error messages for re-authentication. - Updated the Google Drive tool metadata service to track account health and persist authentication status. - Enhanced UI components to display authentication errors and differentiate between valid and expired accounts, improving user experience during file operations.
2026-04-30 19:36:25 +02:00 · 2026-03-20 12:34:30 +05:30 · 2026-03-20 12:34:30 +05:30 · 75f0975674
commit 75f0975674
parent d21593ee71
7 changed files with 213 additions and 27 deletions
--- a/surfsense_backend/app/agents/new_chat/tools/google_drive/create_file.py
+++ b/surfsense_backend/app/agents/new_chat/tools/google_drive/create_file.py
@ -87,6 +87,15 @@ def create_create_google_drive_file_tool(
                logger.error(f"Failed to fetch creation context: {context['error']}")
                return {"status": "error", "message": context["error"]}

+            accounts = context.get("accounts", [])
+            if accounts and all(a.get("auth_expired") for a in accounts):
+                logger.warning("All Google Drive accounts have expired authentication")
+                return {
+                    "status": "auth_error",
+                    "message": "All connected Google Drive accounts need re-authentication. Please re-authenticate in your connector settings.",
+                    "connector_type": "google_drive",
+                }
+
            logger.info(
                f"Requesting approval for creating Google Drive file: name='{name}', type='{file_type}'"
            )
--- a/surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py
+++ b/surfsense_backend/app/agents/new_chat/tools/google_drive/trash_file.py
@ -76,6 +76,18 @@ def create_delete_google_drive_file_tool(
                logger.error(f"Failed to fetch trash context: {error_msg}")
                return {"status": "error", "message": error_msg}

+            account = context.get("account", {})
+            if account.get("auth_expired"):
+                logger.warning(
+                    "Google Drive account %s has expired authentication",
+                    account.get("id"),
+                )
+                return {
+                    "status": "auth_error",
+                    "message": "The Google Drive account for this file needs re-authentication. Please re-authenticate in your connector settings.",
+                    "connector_type": "google_drive",
+                }
+
            file = context["file"]
            file_id = file["file_id"]
            document_id = file.get("document_id")
--- a/surfsense_backend/app/services/google_drive/tool_metadata_service.py
+++ b/surfsense_backend/app/services/google_drive/tool_metadata_service.py
@ -1,15 +1,21 @@
+import logging
 from dataclasses import dataclass

 from sqlalchemy import and_, func
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.future import select
+from sqlalchemy.orm.attributes import flag_modified

+from app.connectors.google_drive.client import GoogleDriveClient
 from app.db import (
    Document,
    DocumentType,
    SearchSourceConnector,
    SearchSourceConnectorType,
 )
+from app.utils.google_credentials import build_composio_credentials
+
+logger = logging.getLogger(__name__)


@dataclass
@ -71,8 +77,17 @@ class GoogleDriveToolMetadataService:
                "error": "No Google Drive account connected",
            }

+        accounts_with_status = []
+        for acc in accounts:
+            acc_dict = acc.to_dict()
+            auth_expired = await self._check_account_health(acc.id)
+            acc_dict["auth_expired"] = auth_expired
+            if auth_expired:
+                await self._persist_auth_expired(acc.id)
+            accounts_with_status.append(acc_dict)
+
        return {
-            "accounts": [acc.to_dict() for acc in accounts],
+            "accounts": accounts_with_status,
            "supported_types": ["google_doc", "google_sheet"],
        }

@ -127,8 +142,14 @@ class GoogleDriveToolMetadataService:
        account = GoogleDriveAccount.from_connector(connector)
        file = GoogleDriveFile.from_document(document)

+        acc_dict = account.to_dict()
+        auth_expired = await self._check_account_health(connector.id)
+        acc_dict["auth_expired"] = auth_expired
+        if auth_expired:
+            await self._persist_auth_expired(connector.id)
+
        return {
-            "account": account.to_dict(),
+            "account": acc_dict,
            "file": file.to_dict(),
        }

@ -151,3 +172,67 @@ class GoogleDriveToolMetadataService:
        )
        connectors = result.scalars().all()
        return [GoogleDriveAccount.from_connector(c) for c in connectors]
+
+    async def _check_account_health(self, connector_id: int) -> bool:
+        """Check if a Google Drive connector's credentials are still valid.
+
+        Uses a lightweight ``files.list(pageSize=1)`` call to verify access.
+
+        Returns True if the credentials are expired/invalid, False if healthy.
+        """
+        try:
+            result = await self._db_session.execute(
+                select(SearchSourceConnector).where(
+                    SearchSourceConnector.id == connector_id
+                )
+            )
+            connector = result.scalar_one_or_none()
+            if not connector:
+                return True
+
+            pre_built_creds = None
+            if (
+                connector.connector_type
+                == SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR
+            ):
+                cca_id = connector.config.get("composio_connected_account_id")
+                if cca_id:
+                    pre_built_creds = build_composio_credentials(cca_id)
+
+            client = GoogleDriveClient(
+                session=self._db_session,
+                connector_id=connector_id,
+                credentials=pre_built_creds,
+            )
+            await client.list_files(
+                query="trashed = false", page_size=1, fields="files(id)"
+            )
+            return False
+        except Exception as e:
+            logger.warning(
+                "Google Drive connector %s health check failed: %s",
+                connector_id,
+                e,
+            )
+            return True
+
+    async def _persist_auth_expired(self, connector_id: int) -> None:
+        """Persist ``auth_expired: True`` to the connector config if not already set."""
+        try:
+            result = await self._db_session.execute(
+                select(SearchSourceConnector).where(
+                    SearchSourceConnector.id == connector_id
+                )
+            )
+            db_connector = result.scalar_one_or_none()
+            if db_connector and not db_connector.config.get("auth_expired"):
+                db_connector.config = {**db_connector.config, "auth_expired": True}
+                flag_modified(db_connector, "config")
+                await self._db_session.commit()
+                await self._db_session.refresh(db_connector)
+        except Exception:
+            logger.warning(
+                "Failed to persist auth_expired for connector %s",
+                connector_id,
+                exc_info=True,
+            )
--- a/surfsense_web/components/tool-ui/google-drive/create-file.tsx
+++ b/surfsense_web/components/tool-ui/google-drive/create-file.tsx
@ -30,6 +30,7 @@ import { openHitlEditPanelAtom } from "@/atoms/chat/hitl-edit-panel.atom";
 interface GoogleDriveAccount {
 	id: number;
 	name: string;
+	auth_expired?: boolean;
 }

 interface InterruptResult {
@ -69,11 +70,18 @@ interface InsufficientPermissionsResult {
 	message: string;
 }

+interface AuthErrorResult {
+	status: "auth_error";
+	message: string;
+	connector_type?: string;
+}
+
 type CreateGoogleDriveFileResult =
 	| InterruptResult
 	| SuccessResult
 	| ErrorResult
-	| InsufficientPermissionsResult;
+	| InsufficientPermissionsResult
+	| AuthErrorResult;

 function isInterruptResult(result: unknown): result is InterruptResult {
 	return (
@ -102,6 +110,15 @@ function isInsufficientPermissionsResult(result: unknown): result is Insufficien
 	);
 }

+function isAuthErrorResult(result: unknown): result is AuthErrorResult {
+	return (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as AuthErrorResult).status === "auth_error"
+	);
+}
+
 const FILE_TYPE_LABELS: Record<string, string> = {
 	google_doc: "Google Doc",
 	google_sheet: "Google Sheet",
@ -127,11 +144,13 @@ function ApprovalCard({
 	const openHitlEditPanel = useSetAtom(openHitlEditPanelAtom);

 	const accounts = interruptData.context?.accounts ?? [];
+	const validAccounts = accounts.filter(a => !a.auth_expired);
+	const expiredAccounts = accounts.filter(a => a.auth_expired);

 	const defaultAccountId = useMemo(() => {
-		if (accounts.length === 1) return String(accounts[0].id);
+		if (validAccounts.length === 1) return String(validAccounts[0].id);
 		return "";
-	}, [accounts]);
+	}, [validAccounts]);

 	const [selectedAccountId, setSelectedAccountId] = useState<string>(defaultAccountId);
 	const [selectedFileType, setSelectedFileType] = useState<string>(args.file_type ?? "google_doc");
@ -245,25 +264,33 @@ function ApprovalCard({
 							<p className="text-sm text-destructive">{interruptData.context.error}</p>
 						) : (
 							<>
-								{accounts.length > 0 && (
-									<div className="space-y-2">
-										<p className="text-xs font-medium text-muted-foreground">
-											Google Drive Account <span className="text-destructive">*</span>
-										</p>
-										<Select value={selectedAccountId} onValueChange={setSelectedAccountId}>
-											<SelectTrigger className="w-full">
-												<SelectValue placeholder="Select an account" />
-											</SelectTrigger>
-											<SelectContent>
-												{accounts.map((account) => (
-													<SelectItem key={account.id} value={String(account.id)}>
-														{account.name}
-													</SelectItem>
-												))}
-											</SelectContent>
-										</Select>
-									</div>
-								)}
+							{accounts.length > 0 && (
+								<div className="space-y-2">
+									<p className="text-xs font-medium text-muted-foreground">
+										Google Drive Account <span className="text-destructive">*</span>
+									</p>
+									<Select value={selectedAccountId} onValueChange={setSelectedAccountId}>
+										<SelectTrigger className="w-full">
+											<SelectValue placeholder="Select an account" />
+										</SelectTrigger>
+										<SelectContent>
+											{validAccounts.map((account) => (
+												<SelectItem key={account.id} value={String(account.id)}>
+													{account.name}
+												</SelectItem>
+											))}
+											{expiredAccounts.map((a) => (
+												<div
+													key={a.id}
+													className="relative flex w-full cursor-default items-center gap-2 rounded-sm py-1.5 px-2 text-sm select-none opacity-50 pointer-events-none"
+												>
+													{a.name} (expired, retry after re-auth)
+												</div>
+											))}
+										</SelectContent>
+									</Select>
+								</div>
+							)}

 								<div className="space-y-2">
 									<p className="text-xs font-medium text-muted-foreground">
@ -435,6 +462,22 @@ function ErrorCard({ result }: { result: ErrorResult }) {
 	);
 }

+function AuthErrorCard({ result }: { result: AuthErrorResult }) {
+	return (
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30">
+			<div className="px-5 pt-5 pb-4">
+				<p className="text-sm font-semibold text-destructive">
+					Google Drive authentication expired
+				</p>
+			</div>
+			<div className="mx-5 h-px bg-border/50" />
+			<div className="px-5 py-4">
+				<p className="text-sm text-muted-foreground">{result.message}</p>
+			</div>
+		</div>
+	);
+}
+
 function SuccessCard({ result }: { result: SuccessResult }) {
 	return (
 		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30">
@ -506,6 +549,8 @@ export const CreateGoogleDriveFileToolUI = makeAssistantToolUI<
 			return null;
 		}

+		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+
 		if (isInsufficientPermissionsResult(result))
 			return <InsufficientPermissionsCard result={result} />;

--- a/surfsense_web/components/tool-ui/google-drive/trash-file.tsx
+++ b/surfsense_web/components/tool-ui/google-drive/trash-file.tsx
@ -19,6 +19,7 @@ import { authenticatedFetch } from "@/lib/auth-utils";
 interface GoogleDriveAccount {
 	id: number;
 	name: string;
+	auth_expired?: boolean;
 }

 interface GoogleDriveFile {
@ -76,13 +77,20 @@ interface InsufficientPermissionsResult {
 	message: string;
 }

+interface AuthErrorResult {
+	status: "auth_error";
+	message: string;
+	connector_type?: string;
+}
+
 type DeleteGoogleDriveFileResult =
 	| InterruptResult
 	| SuccessResult
 	| WarningResult
 	| ErrorResult
 	| NotFoundResult
-	| InsufficientPermissionsResult;
+	| InsufficientPermissionsResult
+	| AuthErrorResult;

 function isInterruptResult(result: unknown): result is InterruptResult {
 	return (
@ -131,6 +139,15 @@ function isInsufficientPermissionsResult(result: unknown): result is Insufficien
 	);
 }

+function isAuthErrorResult(result: unknown): result is AuthErrorResult {
+	return (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as AuthErrorResult).status === "auth_error"
+	);
+}
+
 const MIME_TYPE_LABELS: Record<string, string> = {
 	"application/vnd.google-apps.document": "Google Doc",
 	"application/vnd.google-apps.spreadsheet": "Google Sheet",
@ -363,6 +380,22 @@ function InsufficientPermissionsCard({ result }: { result: InsufficientPermissio
 	);
 }

+function AuthErrorCard({ result }: { result: AuthErrorResult }) {
+	return (
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30">
+			<div className="px-5 pt-5 pb-4">
+				<p className="text-sm font-semibold text-destructive">
+					Google Drive authentication expired
+				</p>
+			</div>
+			<div className="mx-5 h-px bg-border/50" />
+			<div className="px-5 py-4">
+				<p className="text-sm text-muted-foreground">{result.message}</p>
+			</div>
+		</div>
+	);
+}
+
 function WarningCard({ result }: { result: WarningResult }) {
 	return (
 		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30">
@ -464,6 +497,8 @@ export const DeleteGoogleDriveFileToolUI = makeAssistantToolUI<
 			return null;
 		}

+		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+
 		if (isInsufficientPermissionsResult(result))
 			return <InsufficientPermissionsCard result={result} />;

--- a/surfsense_web/components/tool-ui/linear/create-linear-issue.tsx
+++ b/surfsense_web/components/tool-ui/linear/create-linear-issue.tsx
@ -507,7 +507,7 @@ function AuthErrorCard({ result }: { result: AuthErrorResult }) {
 		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30">
 			<div className="px-5 pt-5 pb-4">
 				<p className="text-sm font-semibold text-destructive">
-					Linear authentication expired
+					All Linear accounts expired
 				</p>
 			</div>
 			<div className="mx-5 h-px bg-border/50" />
--- a/surfsense_web/components/tool-ui/notion/delete-notion-page.tsx
+++ b/surfsense_web/components/tool-ui/notion/delete-notion-page.tsx
@ -282,7 +282,7 @@ function AuthErrorCard({ result }: { result: AuthErrorResult }) {
 		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30">
 			<div className="px-5 pt-5 pb-4">
 				<p className="text-sm font-semibold text-destructive">
-					Notion authentication expired
+					All Notion accounts expired
 				</p>
 			</div>
 			<div className="mx-5 h-px bg-border/50" />