mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-01 11:56:25 +02:00
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.
This commit is contained in:
parent
d21593ee71
commit
75f0975674
7 changed files with 213 additions and 27 deletions
|
|
@ -87,6 +87,15 @@ def create_create_google_drive_file_tool(
|
||||||
logger.error(f"Failed to fetch creation context: {context['error']}")
|
logger.error(f"Failed to fetch creation context: {context['error']}")
|
||||||
return {"status": "error", "message": 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(
|
logger.info(
|
||||||
f"Requesting approval for creating Google Drive file: name='{name}', type='{file_type}'"
|
f"Requesting approval for creating Google Drive file: name='{name}', type='{file_type}'"
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -76,6 +76,18 @@ def create_delete_google_drive_file_tool(
|
||||||
logger.error(f"Failed to fetch trash context: {error_msg}")
|
logger.error(f"Failed to fetch trash context: {error_msg}")
|
||||||
return {"status": "error", "message": 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 = context["file"]
|
||||||
file_id = file["file_id"]
|
file_id = file["file_id"]
|
||||||
document_id = file.get("document_id")
|
document_id = file.get("document_id")
|
||||||
|
|
|
||||||
|
|
@ -1,15 +1,21 @@
|
||||||
|
import logging
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
|
||||||
from sqlalchemy import and_, func
|
from sqlalchemy import and_, func
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
from sqlalchemy.future import select
|
from sqlalchemy.future import select
|
||||||
|
from sqlalchemy.orm.attributes import flag_modified
|
||||||
|
|
||||||
|
from app.connectors.google_drive.client import GoogleDriveClient
|
||||||
from app.db import (
|
from app.db import (
|
||||||
Document,
|
Document,
|
||||||
DocumentType,
|
DocumentType,
|
||||||
SearchSourceConnector,
|
SearchSourceConnector,
|
||||||
SearchSourceConnectorType,
|
SearchSourceConnectorType,
|
||||||
)
|
)
|
||||||
|
from app.utils.google_credentials import build_composio_credentials
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|
@ -71,8 +77,17 @@ class GoogleDriveToolMetadataService:
|
||||||
"error": "No Google Drive account connected",
|
"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 {
|
return {
|
||||||
"accounts": [acc.to_dict() for acc in accounts],
|
"accounts": accounts_with_status,
|
||||||
"supported_types": ["google_doc", "google_sheet"],
|
"supported_types": ["google_doc", "google_sheet"],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -127,8 +142,14 @@ class GoogleDriveToolMetadataService:
|
||||||
account = GoogleDriveAccount.from_connector(connector)
|
account = GoogleDriveAccount.from_connector(connector)
|
||||||
file = GoogleDriveFile.from_document(document)
|
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 {
|
return {
|
||||||
"account": account.to_dict(),
|
"account": acc_dict,
|
||||||
"file": file.to_dict(),
|
"file": file.to_dict(),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -151,3 +172,67 @@ class GoogleDriveToolMetadataService:
|
||||||
)
|
)
|
||||||
connectors = result.scalars().all()
|
connectors = result.scalars().all()
|
||||||
return [GoogleDriveAccount.from_connector(c) for c in connectors]
|
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,
|
||||||
|
)
|
||||||
|
|
|
||||||
|
|
@ -30,6 +30,7 @@ import { openHitlEditPanelAtom } from "@/atoms/chat/hitl-edit-panel.atom";
|
||||||
interface GoogleDriveAccount {
|
interface GoogleDriveAccount {
|
||||||
id: number;
|
id: number;
|
||||||
name: string;
|
name: string;
|
||||||
|
auth_expired?: boolean;
|
||||||
}
|
}
|
||||||
|
|
||||||
interface InterruptResult {
|
interface InterruptResult {
|
||||||
|
|
@ -69,11 +70,18 @@ interface InsufficientPermissionsResult {
|
||||||
message: string;
|
message: string;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
interface AuthErrorResult {
|
||||||
|
status: "auth_error";
|
||||||
|
message: string;
|
||||||
|
connector_type?: string;
|
||||||
|
}
|
||||||
|
|
||||||
type CreateGoogleDriveFileResult =
|
type CreateGoogleDriveFileResult =
|
||||||
| InterruptResult
|
| InterruptResult
|
||||||
| SuccessResult
|
| SuccessResult
|
||||||
| ErrorResult
|
| ErrorResult
|
||||||
| InsufficientPermissionsResult;
|
| InsufficientPermissionsResult
|
||||||
|
| AuthErrorResult;
|
||||||
|
|
||||||
function isInterruptResult(result: unknown): result is InterruptResult {
|
function isInterruptResult(result: unknown): result is InterruptResult {
|
||||||
return (
|
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> = {
|
const FILE_TYPE_LABELS: Record<string, string> = {
|
||||||
google_doc: "Google Doc",
|
google_doc: "Google Doc",
|
||||||
google_sheet: "Google Sheet",
|
google_sheet: "Google Sheet",
|
||||||
|
|
@ -127,11 +144,13 @@ function ApprovalCard({
|
||||||
const openHitlEditPanel = useSetAtom(openHitlEditPanelAtom);
|
const openHitlEditPanel = useSetAtom(openHitlEditPanelAtom);
|
||||||
|
|
||||||
const accounts = interruptData.context?.accounts ?? [];
|
const accounts = interruptData.context?.accounts ?? [];
|
||||||
|
const validAccounts = accounts.filter(a => !a.auth_expired);
|
||||||
|
const expiredAccounts = accounts.filter(a => a.auth_expired);
|
||||||
|
|
||||||
const defaultAccountId = useMemo(() => {
|
const defaultAccountId = useMemo(() => {
|
||||||
if (accounts.length === 1) return String(accounts[0].id);
|
if (validAccounts.length === 1) return String(validAccounts[0].id);
|
||||||
return "";
|
return "";
|
||||||
}, [accounts]);
|
}, [validAccounts]);
|
||||||
|
|
||||||
const [selectedAccountId, setSelectedAccountId] = useState<string>(defaultAccountId);
|
const [selectedAccountId, setSelectedAccountId] = useState<string>(defaultAccountId);
|
||||||
const [selectedFileType, setSelectedFileType] = useState<string>(args.file_type ?? "google_doc");
|
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>
|
<p className="text-sm text-destructive">{interruptData.context.error}</p>
|
||||||
) : (
|
) : (
|
||||||
<>
|
<>
|
||||||
{accounts.length > 0 && (
|
{accounts.length > 0 && (
|
||||||
<div className="space-y-2">
|
<div className="space-y-2">
|
||||||
<p className="text-xs font-medium text-muted-foreground">
|
<p className="text-xs font-medium text-muted-foreground">
|
||||||
Google Drive Account <span className="text-destructive">*</span>
|
Google Drive Account <span className="text-destructive">*</span>
|
||||||
</p>
|
</p>
|
||||||
<Select value={selectedAccountId} onValueChange={setSelectedAccountId}>
|
<Select value={selectedAccountId} onValueChange={setSelectedAccountId}>
|
||||||
<SelectTrigger className="w-full">
|
<SelectTrigger className="w-full">
|
||||||
<SelectValue placeholder="Select an account" />
|
<SelectValue placeholder="Select an account" />
|
||||||
</SelectTrigger>
|
</SelectTrigger>
|
||||||
<SelectContent>
|
<SelectContent>
|
||||||
{accounts.map((account) => (
|
{validAccounts.map((account) => (
|
||||||
<SelectItem key={account.id} value={String(account.id)}>
|
<SelectItem key={account.id} value={String(account.id)}>
|
||||||
{account.name}
|
{account.name}
|
||||||
</SelectItem>
|
</SelectItem>
|
||||||
))}
|
))}
|
||||||
</SelectContent>
|
{expiredAccounts.map((a) => (
|
||||||
</Select>
|
<div
|
||||||
</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">
|
<div className="space-y-2">
|
||||||
<p className="text-xs font-medium text-muted-foreground">
|
<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 }) {
|
function SuccessCard({ result }: { result: SuccessResult }) {
|
||||||
return (
|
return (
|
||||||
<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30">
|
<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;
|
return null;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
|
||||||
|
|
||||||
if (isInsufficientPermissionsResult(result))
|
if (isInsufficientPermissionsResult(result))
|
||||||
return <InsufficientPermissionsCard result={result} />;
|
return <InsufficientPermissionsCard result={result} />;
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -19,6 +19,7 @@ import { authenticatedFetch } from "@/lib/auth-utils";
|
||||||
interface GoogleDriveAccount {
|
interface GoogleDriveAccount {
|
||||||
id: number;
|
id: number;
|
||||||
name: string;
|
name: string;
|
||||||
|
auth_expired?: boolean;
|
||||||
}
|
}
|
||||||
|
|
||||||
interface GoogleDriveFile {
|
interface GoogleDriveFile {
|
||||||
|
|
@ -76,13 +77,20 @@ interface InsufficientPermissionsResult {
|
||||||
message: string;
|
message: string;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
interface AuthErrorResult {
|
||||||
|
status: "auth_error";
|
||||||
|
message: string;
|
||||||
|
connector_type?: string;
|
||||||
|
}
|
||||||
|
|
||||||
type DeleteGoogleDriveFileResult =
|
type DeleteGoogleDriveFileResult =
|
||||||
| InterruptResult
|
| InterruptResult
|
||||||
| SuccessResult
|
| SuccessResult
|
||||||
| WarningResult
|
| WarningResult
|
||||||
| ErrorResult
|
| ErrorResult
|
||||||
| NotFoundResult
|
| NotFoundResult
|
||||||
| InsufficientPermissionsResult;
|
| InsufficientPermissionsResult
|
||||||
|
| AuthErrorResult;
|
||||||
|
|
||||||
function isInterruptResult(result: unknown): result is InterruptResult {
|
function isInterruptResult(result: unknown): result is InterruptResult {
|
||||||
return (
|
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> = {
|
const MIME_TYPE_LABELS: Record<string, string> = {
|
||||||
"application/vnd.google-apps.document": "Google Doc",
|
"application/vnd.google-apps.document": "Google Doc",
|
||||||
"application/vnd.google-apps.spreadsheet": "Google Sheet",
|
"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 }) {
|
function WarningCard({ result }: { result: WarningResult }) {
|
||||||
return (
|
return (
|
||||||
<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30">
|
<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;
|
return null;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
|
||||||
|
|
||||||
if (isInsufficientPermissionsResult(result))
|
if (isInsufficientPermissionsResult(result))
|
||||||
return <InsufficientPermissionsCard result={result} />;
|
return <InsufficientPermissionsCard result={result} />;
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30">
|
||||||
<div className="px-5 pt-5 pb-4">
|
<div className="px-5 pt-5 pb-4">
|
||||||
<p className="text-sm font-semibold text-destructive">
|
<p className="text-sm font-semibold text-destructive">
|
||||||
Linear authentication expired
|
All Linear accounts expired
|
||||||
</p>
|
</p>
|
||||||
</div>
|
</div>
|
||||||
<div className="mx-5 h-px bg-border/50" />
|
<div className="mx-5 h-px bg-border/50" />
|
||||||
|
|
|
||||||
|
|
@ -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="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30">
|
||||||
<div className="px-5 pt-5 pb-4">
|
<div className="px-5 pt-5 pb-4">
|
||||||
<p className="text-sm font-semibold text-destructive">
|
<p className="text-sm font-semibold text-destructive">
|
||||||
Notion authentication expired
|
All Notion accounts expired
|
||||||
</p>
|
</p>
|
||||||
</div>
|
</div>
|
||||||
<div className="mx-5 h-px bg-border/50" />
|
<div className="mx-5 h-px bg-border/50" />
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue