feat: improve Composio Drive folder listing and authentication handling

- Simplified the folder and file listing process in the Composio Drive integration by utilizing the GoogleDriveClient for fetching contents.
- Enhanced error handling for authentication issues, marking connectors as 'auth_expired' when necessary and prompting users to re-authenticate.
- Updated UI components to display authentication status and provide re-authentication options, improving user experience during folder selection.

surfsense_backend/app/routes/composio_routes.py

@@ -636,31 +636,23 @@ async def list_composio_drive_folders(
     user: User = Depends(current_active_user),
 ):
     """
-    List folders AND files in user's Google Drive via Composio with hierarchical support.
+    List folders AND files in user's Google Drive via Composio.
 
-    This is called at index time from the manage connector page to display
+    Uses the same GoogleDriveClient / list_folder_contents path as the native
-    the complete file system (folders and files). Only folders are selectable.
+    connector, with Composio-sourced credentials.  This means auth errors
-
+    propagate identically (Google returns 401 → exception → auth_expired flag).
-    Args:
-        connector_id: ID of the Composio Google Drive connector
-        parent_id: Optional parent folder ID to list contents (None for root)
-
-    Returns:
-        JSON with list of items: {
-            "items": [
-                {"id": str, "name": str, "mimeType": str, "isFolder": bool, ...},
-                ...
-            ]
-        }
     """
+    from app.connectors.google_drive import GoogleDriveClient, list_folder_contents
+    from app.utils.google_credentials import build_composio_credentials
+
     if not ComposioService.is_enabled():
         raise HTTPException(
             status_code=503,
             detail="Composio integration is not enabled.",
         )
 
+    connector = None
     try:
-        # Get connector and verify ownership
         result = await session.execute(
             select(SearchSourceConnector).filter(
                 SearchSourceConnector.id == connector_id,
@@ -677,7 +669,6 @@ async def list_composio_drive_folders(
                 detail="Composio Google Drive connector not found or access denied",
             )
 
-        # Get Composio connected account ID from config
         composio_connected_account_id = connector.config.get(
             "composio_connected_account_id"
         )
@@ -687,63 +678,20 @@ async def list_composio_drive_folders(
                 detail="Composio connected account not found. Please reconnect the connector.",
             )
 
-        # Initialize Composio service and fetch files
+        credentials = build_composio_credentials(composio_connected_account_id)
-        service = ComposioService()
+        drive_client = GoogleDriveClient(
-        entity_id = f"surfsense_{user.id}"
+            session, connector_id, credentials=credentials
-
-        # Fetch files/folders from Composio Google Drive
-        files, _next_token, error = await service.get_drive_files(
-            connected_account_id=composio_connected_account_id,
-            entity_id=entity_id,
-            folder_id=parent_id,
-            page_size=100,
         )
 
+        items, error = await list_folder_contents(drive_client, parent_id=parent_id)
+
         if error:
-            logger.error(f"Failed to list Composio Drive files: {error}")
             raise HTTPException(
                 status_code=500, detail=f"Failed to list folder contents: {error}"
             )
 
-        # Transform files to match the expected format with isFolder field
+        folder_count = sum(1 for item in items if item.get("isFolder", False))
-        items = []
+        file_count = len(items) - folder_count
-        for file_info in files:
-            file_id = file_info.get("id", "") or file_info.get("fileId", "")
-            file_name = (
-                file_info.get("name", "") or file_info.get("fileName", "") or "Untitled"
-            )
-            mime_type = file_info.get("mimeType", "") or file_info.get("mime_type", "")
-
-            if not file_id:
-                continue
-
-            is_folder = mime_type == "application/vnd.google-apps.folder"
-
-            items.append(
-                {
-                    "id": file_id,
-                    "name": file_name,
-                    "mimeType": mime_type,
-                    "isFolder": is_folder,
-                    "parents": file_info.get("parents", []),
-                    "size": file_info.get("size"),
-                    "iconLink": file_info.get("iconLink"),
-                }
-            )
-
-        # Sort: folders first, then files, both alphabetically
-        folders = sorted(
-            [item for item in items if item["isFolder"]],
-            key=lambda x: x["name"].lower(),
-        )
-        files_list = sorted(
-            [item for item in items if not item["isFolder"]],
-            key=lambda x: x["name"].lower(),
-        )
-        items = folders + files_list
-
-        folder_count = len(folders)
-        file_count = len(files_list)
 
         logger.info(
             f"Listed {len(items)} total items ({folder_count} folders, {file_count} files) for Composio connector {connector_id}"
@@ -757,7 +705,13 @@ async def list_composio_drive_folders(
     except Exception as e:
         logger.error(f"Error listing Composio Drive contents: {e!s}", exc_info=True)
         error_lower = str(e).lower()
-        if "401" in str(e) or "invalid credentials" in error_lower or "autherror" in error_lower or "authentication failed" in error_lower:
+        if (
+            "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
+            or "401" in str(e)
+        ):
             try:
                 if connector and not connector.config.get("auth_expired"):
                     connector.config = {**connector.config, "auth_expired": True}

surfsense_web/components/assistant-ui/connector-popup/connector-configs/components/composio-drive-config.tsx

@@ -1,5 +1,7 @@
 "use client";
 
+import { cn } from "@/lib/utils";
+import { useAtomValue } from "jotai";
 import {
 	ChevronDown,
 	ChevronRight,
@@ -9,11 +11,15 @@ import {
 	FolderClosed,
 	Image,
 	Presentation,
+	RefreshCw,
 	X,
 } from "lucide-react";
 import type { FC } from "react";
-import { useEffect, useState } from "react";
+import { useCallback, useEffect, useState } from "react";
+import { toast } from "sonner";
+import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
 import { ComposioDriveFolderTree } from "@/components/connectors/composio-drive-folder-tree";
+import { Button } from "@/components/ui/button";
 import { Label } from "@/components/ui/label";
 import {
 	Select,
@@ -23,13 +29,8 @@ import {
 	SelectValue,
 } from "@/components/ui/select";
 import { Switch } from "@/components/ui/switch";
-import type { SearchSourceConnector } from "@/contracts/types/connector.types";
+import { authenticatedFetch } from "@/lib/auth-utils";
-
+import type { ConnectorConfigProps } from "../index";
-interface ComposioDriveConfigProps {
-	connector: SearchSourceConnector;
-	onConfigChange?: (config: Record<string, unknown>) => void;
-	onNameChange?: (name: string) => void;
-}
 
 interface SelectedFolder {
 	id: string;
@@ -88,10 +89,13 @@ function getFileIconFromName(fileName: string, className: string = "size-3.5 shr
 	return <File className={`${className} text-gray-500`} />;
 }
 
-export const ComposioDriveConfig: FC<ComposioDriveConfigProps> = ({
+const COMPOSIO_REAUTH_ENDPOINT = "/api/v1/auth/composio/connector/reauth";
+
+export const ComposioDriveConfig: FC<ConnectorConfigProps> = ({
 	connector,
 	onConfigChange,
 }) => {
+	const searchSpaceId = useAtomValue(activeSearchSpaceIdAtom);
 	const isIndexable = connector.config?.is_indexable as boolean;
 
 	const existingFolders =
@@ -103,6 +107,43 @@ export const ComposioDriveConfig: FC<ComposioDriveConfigProps> = ({
 	const [selectedFolders, setSelectedFolders] = useState<SelectedFolder[]>(existingFolders);
 	const [selectedFiles, setSelectedFiles] = useState<SelectedFolder[]>(existingFiles);
 	const [indexingOptions, setIndexingOptions] = useState<IndexingOptions>(existingIndexingOptions);
+	const [reauthing, setReauthing] = useState(false);
+	const [authError, setAuthError] = useState(false);
+
+	const isAuthExpired = connector.config?.auth_expired === true || authError;
+
+	const handleReauth = useCallback(async () => {
+		if (!searchSpaceId) return;
+		setReauthing(true);
+		try {
+			const backendUrl = process.env.NEXT_PUBLIC_FASTAPI_BACKEND_URL || "http://localhost:8000";
+			const url = new URL(`${backendUrl}${COMPOSIO_REAUTH_ENDPOINT}`);
+			url.searchParams.set("connector_id", String(connector.id));
+			url.searchParams.set("space_id", String(searchSpaceId));
+			url.searchParams.set("return_url", window.location.pathname);
+			const response = await authenticatedFetch(url.toString());
+			if (!response.ok) {
+				const data = await response.json().catch(() => ({}));
+				toast.error(data.detail ?? "Failed to initiate re-authentication.");
+				return;
+			}
+			const data = await response.json();
+			if (data.auth_url) {
+				window.location.href = data.auth_url;
+			} else if (data.success) {
+				toast.success("Authentication refreshed successfully.");
+				setAuthError(false);
+			}
+		} catch {
+			toast.error("Failed to initiate re-authentication.");
+		} finally {
+			setReauthing(false);
+		}
+	}, [searchSpaceId, connector.id]);
+
+	const handleAuthError = useCallback(() => {
+		setAuthError(true);
+	}, []);
 
 	const [isEditMode] = useState(() => existingFolders.length > 0 || existingFiles.length > 0);
 	const [isFolderTreeOpen, setIsFolderTreeOpen] = useState(!isEditMode);
@@ -235,39 +276,58 @@ export const ComposioDriveConfig: FC<ComposioDriveConfigProps> = ({
 					</div>
 				)}
 
-				{isEditMode ? (
+			{isAuthExpired && (
-					<div className="space-y-2">
+				<div className="flex items-center gap-2">
-						<button
+					<p className="text-xs text-amber-600 dark:text-amber-500">
-							type="button"
+						Your Google Drive authentication has expired.
-							onClick={() => setIsFolderTreeOpen(!isFolderTreeOpen)}
+					</p>
-							className="flex items-center gap-2 text-xs sm:text-sm text-muted-foreground hover:text-foreground transition-colors w-fit"
+					<Button
-						>
+						size="sm"
-							{isFolderTreeOpen ? (
+						className="h-7 text-[11px] px-3 rounded-lg font-medium bg-amber-600 hover:bg-amber-700 text-white border-0 shadow-xs shrink-0"
-								<ChevronDown className="size-4" />
+						onClick={handleReauth}
-							) : (
+						disabled={reauthing}
-								<ChevronRight className="size-4" />
+					>
-							)}
+						<RefreshCw className={cn("size-3.5", reauthing && "animate-spin")} />
-							Change Selection
+						Re-authenticate
-						</button>
+					</Button>
-						{isFolderTreeOpen && (
+				</div>
-							<ComposioDriveFolderTree
+			)}
-								connectorId={connector.id}
+
-								selectedFolders={selectedFolders}
+			{isEditMode ? (
-								onSelectFolders={handleSelectFolders}
+				<div className="space-y-2">
-								selectedFiles={selectedFiles}
+					<button
-								onSelectFiles={handleSelectFiles}
+						type="button"
-							/>
+						onClick={() => setIsFolderTreeOpen(!isFolderTreeOpen)}
+						className="flex items-center gap-2 text-xs sm:text-sm text-muted-foreground hover:text-foreground transition-colors w-fit"
+					>
+						Change Selection
+						{isFolderTreeOpen ? (
+							<ChevronDown className="size-4" />
+						) : (
+							<ChevronRight className="size-4" />
 						)}
-					</div>
+					</button>
-				) : (
+					{isFolderTreeOpen && (
-					<ComposioDriveFolderTree
+						<ComposioDriveFolderTree
-						connectorId={connector.id}
+							connectorId={connector.id}
-						selectedFolders={selectedFolders}
+							selectedFolders={selectedFolders}
-						onSelectFolders={handleSelectFolders}
+							onSelectFolders={handleSelectFolders}
-						selectedFiles={selectedFiles}
+							selectedFiles={selectedFiles}
-						onSelectFiles={handleSelectFiles}
+							onSelectFiles={handleSelectFiles}
-					/>
+							onAuthError={handleAuthError}
-				)}
+						/>
+					)}
+				</div>
+			) : (
+				<ComposioDriveFolderTree
+					connectorId={connector.id}
+					selectedFolders={selectedFolders}
+					onSelectFolders={handleSelectFolders}
+					selectedFiles={selectedFiles}
+					onSelectFiles={handleSelectFiles}
+					onAuthError={handleAuthError}
+				/>
+			)}
 			</div>
 
 			{/* Indexing Options */}

surfsense_web/components/assistant-ui/connector-popup/connector-configs/components/google-drive-config.tsx

@@ -1,5 +1,6 @@
 "use client";
 
+import { cn } from "@/lib/utils";
 import { useAtomValue } from "jotai";
 import {
 
@@ -294,11 +295,7 @@ export const GoogleDriveConfig: FC<ConnectorConfigProps> = ({ connector, onConfi
 								onClick={handleReauth}
 								disabled={reauthing}
 							>
-								{reauthing ? (
+								<RefreshCw className={cn("size-3.5", reauthing && "animate-spin")} />
-									<Spinner size="xs" />
-								) : (
-									<RefreshCw className="size-3.5" />
-								)}
 								Re-authenticate
 							</Button>
 						</div>

surfsense_web/components/connectors/composio-drive-folder-tree.tsx

@@ -12,7 +12,7 @@ import {
 	Image,
 	Presentation,
 } from "lucide-react";
-import { useState } from "react";
+import { useEffect, useState } from "react";
 import { Checkbox } from "@/components/ui/checkbox";
 import { ScrollArea } from "@/components/ui/scroll-area";
 import { Spinner } from "@/components/ui/spinner";
@@ -48,6 +48,7 @@ interface ComposioDriveFolderTreeProps {
 	onSelectFolders: (folders: SelectedFolder[]) => void;
 	selectedFiles?: SelectedFolder[];
 	onSelectFiles?: (files: SelectedFolder[]) => void;
+	onAuthError?: (message: string) => void;
 }
 
 // Helper to get appropriate icon for file type
@@ -73,13 +74,23 @@ export function ComposioDriveFolderTree({
 	onSelectFolders,
 	selectedFiles = [],
 	onSelectFiles = () => {},
+	onAuthError,
 }: ComposioDriveFolderTreeProps) {
 	const [itemStates, setItemStates] = useState<Map<string, ItemTreeNode>>(new Map());
 
-	const { data: rootData, isLoading: isLoadingRoot } = useComposioDriveFolders({
+	const { data: rootData, isLoading: isLoadingRoot, error: rootError } = useComposioDriveFolders({
 		connectorId,
 	});
 
+	useEffect(() => {
+		if (rootError && onAuthError) {
+			const msg = rootError instanceof Error ? rootError.message : String(rootError);
+			if (msg.toLowerCase().includes("authentication expired") || msg.toLowerCase().includes("re-authenticate")) {
+				onAuthError(msg);
+			}
+		}
+	}, [rootError, onAuthError]);
+
 	const rootItems = rootData?.items || [];
 
 	const isFolderSelected = (folderId: string): boolean => {
@@ -352,11 +363,19 @@ export function ComposioDriveFolderTree({
 						{!isLoadingRoot && rootItems.map((item) => renderItem(item, 0))}
 					</div>
 
-					{!isLoadingRoot && rootItems.length === 0 && (
+				{!isLoadingRoot && rootError && (
-						<div className="text-center text-xs sm:text-sm text-muted-foreground py-4 sm:py-8">
+					<div className="text-center text-xs sm:text-sm text-amber-600 dark:text-amber-500 py-4 sm:py-8">
-							No files or folders found in your Google Drive
+						{(rootError instanceof Error ? rootError.message : String(rootError)).includes("authentication expired")
-						</div>
+							? "Google Drive authentication has expired. Please re-authenticate above."
-					)}
+							: "Failed to load Google Drive contents."}
+					</div>
+				)}
+
+				{!isLoadingRoot && !rootError && rootItems.length === 0 && (
+					<div className="text-center text-xs sm:text-sm text-muted-foreground py-4 sm:py-8">
+						No files or folders found in your Google Drive
+					</div>
+				)}
 				</div>
 			</ScrollArea>
 		</div>