feat: add update functionality for Gmail drafts

- Introduced a new tool to update existing Gmail drafts, allowing users to modify draft content, recipients, and subject lines.
- Updated the Gmail tools registry to include the new update_gmail_draft tool.
- Enhanced the GmailKBSyncService to support draft ID handling during synchronization.
- Added UI components for the update draft functionality in the web application, improving user interaction with Gmail drafts.

surfsense_backend/app/agents/new_chat/tools/gmail/__init__.py

@@ -7,9 +7,13 @@ from app.agents.new_chat.tools.gmail.send_email import (
 from app.agents.new_chat.tools.gmail.trash_email import (
     create_trash_gmail_email_tool,
 )
+from app.agents.new_chat.tools.gmail.update_draft import (
+    create_update_gmail_draft_tool,
+)
 
 __all__ = [
     "create_create_gmail_draft_tool",
     "create_send_gmail_email_tool",
     "create_trash_gmail_email_tool",
+    "create_update_gmail_draft_tool",
 ]

surfsense_backend/app/agents/new_chat/tools/gmail/create_draft.py

@@ -309,6 +309,7 @@ def create_create_gmail_draft_tool(
                     connector_id=actual_connector_id,
                     search_space_id=search_space_id,
                     user_id=user_id,
+                    draft_id=created.get("id"),
                 )
                 if kb_result["status"] == "success":
                     kb_message_suffix = " Your knowledge base has also been updated."

surfsense_backend/app/agents/new_chat/tools/gmail/update_draft.py

@@ -0,0 +1,420 @@
+import asyncio
+import base64
+import logging
+from datetime import datetime
+from email.mime.text import MIMEText
+from typing import Any
+
+from langchain_core.tools import tool
+from langgraph.types import interrupt
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.services.gmail import GmailToolMetadataService
+
+logger = logging.getLogger(__name__)
+
+
+def create_update_gmail_draft_tool(
+    db_session: AsyncSession | None = None,
+    search_space_id: int | None = None,
+    user_id: str | None = None,
+):
+    @tool
+    async def update_gmail_draft(
+        draft_subject_or_id: str,
+        body: str,
+        to: str | None = None,
+        subject: str | None = None,
+        cc: str | None = None,
+        bcc: str | None = None,
+    ) -> dict[str, Any]:
+        """Update an existing Gmail draft.
+
+        Use when the user asks to modify, edit, or add content to an existing
+        email draft. This replaces the draft content with the new version.
+
+        IMPORTANT: This tool is ONLY for Gmail drafts, NOT for Notion pages,
+        calendar events, or any other content type.
+
+        Args:
+            draft_subject_or_id: The exact subject line of the draft to update
+                (as it appears in Gmail drafts).
+            body: The full updated body content for the draft.
+            to: Optional new recipient email address (keeps original if omitted).
+            subject: Optional new subject line (keeps original if omitted).
+            cc: Optional CC recipient(s), comma-separated.
+            bcc: Optional BCC recipient(s), comma-separated.
+
+        Returns:
+            Dictionary with:
+            - status: "success", "rejected", "not_found", or "error"
+            - draft_id: Gmail draft ID (if success)
+            - message: Result message
+
+            IMPORTANT:
+            - If status is "rejected", the user explicitly declined the action.
+              Respond with a brief acknowledgment and do NOT retry or suggest alternatives.
+            - If status is "not_found", relay the exact message to the user and ask them
+              to verify the draft subject or check if it has been indexed.
+            - If status is "insufficient_permissions", the connector lacks the required OAuth scope.
+              Inform the user they need to re-authenticate and do NOT retry the action.
+
+        Examples:
+            - "Update the Kurseong Plan draft with the new itinerary details"
+            - "Edit my draft about the project proposal and change the recipient"
+        """
+        logger.info(
+            f"update_gmail_draft called: draft_subject_or_id='{draft_subject_or_id}'"
+        )
+
+        if db_session is None or search_space_id is None or user_id is None:
+            return {
+                "status": "error",
+                "message": "Gmail tool not properly configured. Please contact support.",
+            }
+
+        try:
+            metadata_service = GmailToolMetadataService(db_session)
+            context = await metadata_service.get_update_context(
+                search_space_id, user_id, draft_subject_or_id
+            )
+
+            if "error" in context:
+                error_msg = context["error"]
+                if "not found" in error_msg.lower():
+                    logger.warning(f"Draft not found: {error_msg}")
+                    return {"status": "not_found", "message": error_msg}
+                logger.error(f"Failed to fetch update context: {error_msg}")
+                return {"status": "error", "message": error_msg}
+
+            account = context.get("account", {})
+            if account.get("auth_expired"):
+                logger.warning(
+                    "Gmail account %s has expired authentication",
+                    account.get("id"),
+                )
+                return {
+                    "status": "auth_error",
+                    "message": "The Gmail account for this draft needs re-authentication. Please re-authenticate in your connector settings.",
+                    "connector_type": "gmail",
+                }
+
+            email = context["email"]
+            message_id = email["message_id"]
+            document_id = email.get("document_id")
+            connector_id_from_context = account["id"]
+            draft_id_from_context = context.get("draft_id")
+
+            original_subject = email.get("subject", draft_subject_or_id)
+            final_subject_default = subject if subject else original_subject
+            final_to_default = to if to else ""
+
+            logger.info(
+                f"Requesting approval for updating Gmail draft: '{original_subject}' "
+                f"(message_id={message_id}, draft_id={draft_id_from_context})"
+            )
+            approval = interrupt(
+                {
+                    "type": "gmail_draft_update",
+                    "action": {
+                        "tool": "update_gmail_draft",
+                        "params": {
+                            "message_id": message_id,
+                            "draft_id": draft_id_from_context,
+                            "to": final_to_default,
+                            "subject": final_subject_default,
+                            "body": body,
+                            "cc": cc,
+                            "bcc": bcc,
+                            "connector_id": connector_id_from_context,
+                        },
+                    },
+                    "context": context,
+                }
+            )
+
+            decisions_raw = (
+                approval.get("decisions", []) if isinstance(approval, dict) else []
+            )
+            decisions = (
+                decisions_raw if isinstance(decisions_raw, list) else [decisions_raw]
+            )
+            decisions = [d for d in decisions if isinstance(d, dict)]
+            if not decisions:
+                logger.warning("No approval decision received")
+                return {"status": "error", "message": "No approval decision received"}
+
+            decision = decisions[0]
+            decision_type = decision.get("type") or decision.get("decision_type")
+            logger.info(f"User decision: {decision_type}")
+
+            if decision_type == "reject":
+                return {
+                    "status": "rejected",
+                    "message": "User declined. The draft was not updated. Do not ask again or suggest alternatives.",
+                }
+
+            final_params: dict[str, Any] = {}
+            edited_action = decision.get("edited_action")
+            if isinstance(edited_action, dict):
+                edited_args = edited_action.get("args")
+                if isinstance(edited_args, dict):
+                    final_params = edited_args
+            elif isinstance(decision.get("args"), dict):
+                final_params = decision["args"]
+
+            final_to = final_params.get("to", final_to_default)
+            final_subject = final_params.get("subject", final_subject_default)
+            final_body = final_params.get("body", body)
+            final_cc = final_params.get("cc", cc)
+            final_bcc = final_params.get("bcc", bcc)
+            final_connector_id = final_params.get(
+                "connector_id", connector_id_from_context
+            )
+            final_draft_id = final_params.get("draft_id", draft_id_from_context)
+
+            if not final_connector_id:
+                return {
+                    "status": "error",
+                    "message": "No connector found for this draft.",
+                }
+
+            from sqlalchemy.future import select
+
+            from app.db import SearchSourceConnector, SearchSourceConnectorType
+
+            _GMAIL_TYPES = [
+                SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR,
+                SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR,
+            ]
+
+            result = await db_session.execute(
+                select(SearchSourceConnector).filter(
+                    SearchSourceConnector.id == final_connector_id,
+                    SearchSourceConnector.search_space_id == search_space_id,
+                    SearchSourceConnector.user_id == user_id,
+                    SearchSourceConnector.connector_type.in_(_GMAIL_TYPES),
+                )
+            )
+            connector = result.scalars().first()
+            if not connector:
+                return {
+                    "status": "error",
+                    "message": "Selected Gmail connector is invalid or has been disconnected.",
+                }
+
+            logger.info(
+                f"Updating Gmail draft: subject='{final_subject}', connector={final_connector_id}"
+            )
+
+            if connector.connector_type == SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR:
+                from app.utils.google_credentials import build_composio_credentials
+
+                cca_id = connector.config.get("composio_connected_account_id")
+                if cca_id:
+                    creds = build_composio_credentials(cca_id)
+                else:
+                    return {
+                        "status": "error",
+                        "message": "Composio connected account ID not found for this Gmail connector.",
+                    }
+            else:
+                from google.oauth2.credentials import Credentials
+
+                from app.config import config
+                from app.utils.oauth_security import TokenEncryption
+
+                config_data = dict(connector.config)
+                token_encrypted = config_data.get("_token_encrypted", False)
+                if token_encrypted and config.SECRET_KEY:
+                    token_encryption = TokenEncryption(config.SECRET_KEY)
+                    if config_data.get("token"):
+                        config_data["token"] = token_encryption.decrypt_token(
+                            config_data["token"]
+                        )
+                    if config_data.get("refresh_token"):
+                        config_data["refresh_token"] = token_encryption.decrypt_token(
+                            config_data["refresh_token"]
+                        )
+                    if config_data.get("client_secret"):
+                        config_data["client_secret"] = token_encryption.decrypt_token(
+                            config_data["client_secret"]
+                        )
+
+                exp = config_data.get("expiry", "")
+                if exp:
+                    exp = exp.replace("Z", "")
+
+                creds = Credentials(
+                    token=config_data.get("token"),
+                    refresh_token=config_data.get("refresh_token"),
+                    token_uri=config_data.get("token_uri"),
+                    client_id=config_data.get("client_id"),
+                    client_secret=config_data.get("client_secret"),
+                    scopes=config_data.get("scopes", []),
+                    expiry=datetime.fromisoformat(exp) if exp else None,
+                )
+
+            from googleapiclient.discovery import build
+
+            gmail_service = build("gmail", "v1", credentials=creds)
+
+            # Resolve draft_id if not already available
+            if not final_draft_id:
+                logger.info(
+                    f"draft_id not in metadata, looking up via drafts.list for message_id={message_id}"
+                )
+                final_draft_id = await _find_draft_id_by_message(
+                    gmail_service, message_id
+                )
+
+            if not final_draft_id:
+                return {
+                    "status": "error",
+                    "message": (
+                        "Could not find this draft in Gmail. "
+                        "It may have already been sent or deleted."
+                    ),
+                }
+
+            message = MIMEText(final_body)
+            if final_to:
+                message["to"] = final_to
+            message["subject"] = final_subject
+            if final_cc:
+                message["cc"] = final_cc
+            if final_bcc:
+                message["bcc"] = final_bcc
+            raw = base64.urlsafe_b64encode(message.as_bytes()).decode()
+
+            try:
+                updated = await asyncio.get_event_loop().run_in_executor(
+                    None,
+                    lambda: gmail_service.users()
+                    .drafts()
+                    .update(
+                        userId="me",
+                        id=final_draft_id,
+                        body={"message": {"raw": raw}},
+                    )
+                    .execute(),
+                )
+            except Exception as api_err:
+                from googleapiclient.errors import HttpError
+
+                if isinstance(api_err, HttpError) and api_err.resp.status == 403:
+                    logger.warning(
+                        f"Insufficient permissions for connector {connector.id}: {api_err}"
+                    )
+                    try:
+                        from sqlalchemy.orm.attributes import flag_modified
+
+                        if not connector.config.get("auth_expired"):
+                            connector.config = {
+                                **connector.config,
+                                "auth_expired": True,
+                            }
+                            flag_modified(connector, "config")
+                            await db_session.commit()
+                    except Exception:
+                        logger.warning(
+                            "Failed to persist auth_expired for connector %s",
+                            connector.id,
+                            exc_info=True,
+                        )
+                    return {
+                        "status": "insufficient_permissions",
+                        "connector_id": connector.id,
+                        "message": "This Gmail account needs additional permissions. Please re-authenticate in connector settings.",
+                    }
+                if isinstance(api_err, HttpError) and api_err.resp.status == 404:
+                    return {
+                        "status": "error",
+                        "message": "Draft no longer exists in Gmail. It may have been sent or deleted.",
+                    }
+                raise
+
+            logger.info(f"Gmail draft updated: id={updated.get('id')}")
+
+            kb_message_suffix = ""
+            if document_id:
+                try:
+                    from sqlalchemy.future import select as sa_select
+                    from sqlalchemy.orm.attributes import flag_modified
+
+                    from app.db import Document
+
+                    doc_result = await db_session.execute(
+                        sa_select(Document).filter(Document.id == document_id)
+                    )
+                    document = doc_result.scalars().first()
+                    if document:
+                        document.source_markdown = final_body
+                        document.title = final_subject
+                        meta = dict(document.document_metadata or {})
+                        meta["subject"] = final_subject
+                        meta["draft_id"] = updated.get("id", final_draft_id)
+                        updated_msg = updated.get("message", {})
+                        if updated_msg.get("id"):
+                            meta["message_id"] = updated_msg["id"]
+                        document.document_metadata = meta
+                        flag_modified(document, "document_metadata")
+                        await db_session.commit()
+                        kb_message_suffix = " Your knowledge base has also been updated."
+                        logger.info(
+                            f"KB document {document_id} updated for draft {final_draft_id}"
+                        )
+                    else:
+                        kb_message_suffix = " This draft will be fully updated in your knowledge base in the next scheduled sync."
+                except Exception as kb_err:
+                    logger.warning(f"KB update after draft edit failed: {kb_err}")
+                    await db_session.rollback()
+                    kb_message_suffix = " This draft will be fully updated in your knowledge base in the next scheduled sync."
+
+            return {
+                "status": "success",
+                "draft_id": updated.get("id"),
+                "message": f"Successfully updated Gmail draft with subject '{final_subject}'.{kb_message_suffix}",
+            }
+
+        except Exception as e:
+            from langgraph.errors import GraphInterrupt
+
+            if isinstance(e, GraphInterrupt):
+                raise
+
+            logger.error(f"Error updating Gmail draft: {e}", exc_info=True)
+            return {
+                "status": "error",
+                "message": "Something went wrong while updating the draft. Please try again.",
+            }
+
+    return update_gmail_draft
+
+
+async def _find_draft_id_by_message(gmail_service: Any, message_id: str) -> str | None:
+    """Look up a draft's ID by its message ID via the Gmail API."""
+    try:
+        page_token = None
+        while True:
+            kwargs: dict[str, Any] = {"userId": "me", "maxResults": 100}
+            if page_token:
+                kwargs["pageToken"] = page_token
+
+            response = await asyncio.get_event_loop().run_in_executor(
+                None,
+                lambda: gmail_service.users().drafts().list(**kwargs).execute(),
+            )
+
+            for draft in response.get("drafts", []):
+                if draft.get("message", {}).get("id") == message_id:
+                    return draft["id"]
+
+            page_token = response.get("nextPageToken")
+            if not page_token:
+                break
+
+        return None
+    except Exception as e:
+        logger.warning(f"Failed to look up draft by message_id: {e}")
+        return None

surfsense_backend/app/agents/new_chat/tools/registry.py

@@ -51,6 +51,7 @@ from .gmail import (
     create_create_gmail_draft_tool,
     create_send_gmail_email_tool,
     create_trash_gmail_email_tool,
+    create_update_gmail_draft_tool,
 )
 from .google_calendar import (
     create_create_calendar_event_tool,
@@ -381,7 +382,7 @@ BUILTIN_TOOLS: list[ToolDefinition] = [
         requires=["db_session", "search_space_id", "user_id"],
     ),
     # =========================================================================
-    # GMAIL TOOLS - create drafts, send emails, trash emails
+    # GMAIL TOOLS - create drafts, update drafts, send emails, trash emails
     # Auto-disabled when no Gmail connector is configured
     # =========================================================================
     ToolDefinition(
@@ -414,6 +415,16 @@ BUILTIN_TOOLS: list[ToolDefinition] = [
         ),
         requires=["db_session", "search_space_id", "user_id"],
     ),
+    ToolDefinition(
+        name="update_gmail_draft",
+        description="Update an existing Gmail draft",
+        factory=lambda deps: create_update_gmail_draft_tool(
+            db_session=deps["db_session"],
+            search_space_id=deps["search_space_id"],
+            user_id=deps["user_id"],
+        ),
+        requires=["db_session", "search_space_id", "user_id"],
+    ),
 ]
 
 

surfsense_backend/app/services/gmail/kb_sync_service.py

@@ -31,6 +31,7 @@ class GmailKBSyncService:
         connector_id: int,
         search_space_id: int,
         user_id: str,
+        draft_id: str | None = None,
     ) -> dict:
         from app.tasks.connector_indexers.base import (
             check_document_by_unique_identifier,
@@ -103,18 +104,22 @@ class GmailKBSyncService:
             chunks = await create_document_chunks(indexable_content)
             now_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
+            doc_metadata = {
+                "message_id": message_id,
+                "thread_id": thread_id,
+                "subject": subject,
+                "sender": sender,
+                "date": date_str,
+                "connector_id": connector_id,
+                "indexed_at": now_str,
+            }
+            if draft_id:
+                doc_metadata["draft_id"] = draft_id
+
             document = Document(
                 title=subject,
                 document_type=DocumentType.GOOGLE_GMAIL_CONNECTOR,
-                document_metadata={
+                document_metadata=doc_metadata,
-                    "message_id": message_id,
-                    "thread_id": thread_id,
-                    "subject": subject,
-                    "sender": sender,
-                    "date": date_str,
-                    "connector_id": connector_id,
-                    "indexed_at": now_str,
-                },
                 content=summary_content,
                 content_hash=content_hash,
                 unique_identifier_hash=unique_hash,

surfsense_backend/app/services/gmail/tool_metadata_service.py

@@ -237,6 +237,44 @@ class GmailToolMetadataService:
 
         return {"accounts": accounts_with_status}
 
+    async def get_update_context(
+        self, search_space_id: int, user_id: str, email_ref: str
+    ) -> dict:
+        document, connector = await self._resolve_email(
+            search_space_id, user_id, email_ref
+        )
+
+        if not document or not connector:
+            return {
+                "error": (
+                    f"Draft '{email_ref}' not found in your indexed Gmail messages. "
+                    "This could mean: (1) the draft doesn't exist, "
+                    "(2) it hasn't been indexed yet, "
+                    "or (3) the subject is different. "
+                    "Please check the exact draft subject in Gmail."
+                )
+            }
+
+        account = GmailAccount.from_connector(connector)
+        message = GmailMessage.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)
+
+        result = {
+            "account": acc_dict,
+            "email": message.to_dict(),
+        }
+
+        meta = document.document_metadata or {}
+        if meta.get("draft_id"):
+            result["draft_id"] = meta["draft_id"]
+
+        return result
+
     async def get_trash_context(
         self, search_space_id: int, user_id: str, email_ref: str
     ) -> dict:

surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx

@@ -50,6 +50,7 @@ import {
 	CreateGmailDraftToolUI,
 	SendGmailEmailToolUI,
 	TrashGmailEmailToolUI,
+	UpdateGmailDraftToolUI,
 } from "@/components/tool-ui/gmail";
 import {
 	CreateGoogleDriveFileToolUI,
@@ -1696,6 +1697,7 @@ export default function NewChatPage() {
 			<UpdateCalendarEventToolUI />
 			<DeleteCalendarEventToolUI />
 			<CreateGmailDraftToolUI />
+			<UpdateGmailDraftToolUI />
 			<SendGmailEmailToolUI />
 			<TrashGmailEmailToolUI />
 			<SandboxExecuteToolUI />

surfsense_web/components/tool-ui/gmail/index.ts

@@ -1,3 +1,4 @@
 export { CreateGmailDraftToolUI } from "./create-draft";
 export { SendGmailEmailToolUI } from "./send-email";
 export { TrashGmailEmailToolUI } from "./trash-email";
+export { UpdateGmailDraftToolUI } from "./update-draft";

surfsense_web/components/tool-ui/gmail/update-draft.tsx

@@ -0,0 +1,600 @@
+"use client";
+
+import { makeAssistantToolUI } from "@assistant-ui/react";
+import {
+	CornerDownLeftIcon,
+	MailIcon,
+	Pen,
+	TriangleAlertIcon,
+	UserIcon,
+	UsersIcon,
+} from "lucide-react";
+import { useCallback, useEffect, useState } from "react";
+import { Button } from "@/components/ui/button";
+import { PlateEditor } from "@/components/editor/plate-editor";
+import { TextShimmerLoader } from "@/components/prompt-kit/loader";
+import { useSetAtom } from "jotai";
+import { openHitlEditPanelAtom } from "@/atoms/chat/hitl-edit-panel.atom";
+import type { ExtraField } from "@/atoms/chat/hitl-edit-panel.atom";
+
+interface GmailAccount {
+	id: number;
+	name: string;
+	email: string;
+	auth_expired?: boolean;
+}
+
+interface GmailMessage {
+	message_id: string;
+	thread_id?: string;
+	subject: string;
+	sender: string;
+	date: string;
+	connector_id: number;
+	document_id: number;
+}
+
+interface InterruptResult {
+	__interrupt__: true;
+	__decided__?: "approve" | "reject" | "edit";
+	action_requests: Array<{
+		name: string;
+		args: Record<string, unknown>;
+	}>;
+	review_configs: Array<{
+		action_name: string;
+		allowed_decisions: Array<"approve" | "edit" | "reject">;
+	}>;
+	context?: {
+		account?: GmailAccount;
+		email?: GmailMessage;
+		draft_id?: string;
+		error?: string;
+	};
+}
+
+interface SuccessResult {
+	status: "success";
+	draft_id?: string;
+	message?: string;
+}
+
+interface ErrorResult {
+	status: "error";
+	message: string;
+}
+
+interface NotFoundResult {
+	status: "not_found";
+	message: string;
+}
+
+interface AuthErrorResult {
+	status: "auth_error";
+	message: string;
+	connector_type?: string;
+}
+
+interface InsufficientPermissionsResult {
+	status: "insufficient_permissions";
+	connector_id: number;
+	message: string;
+}
+
+type UpdateGmailDraftResult =
+	| InterruptResult
+	| SuccessResult
+	| ErrorResult
+	| NotFoundResult
+	| InsufficientPermissionsResult
+	| AuthErrorResult;
+
+function isInterruptResult(result: unknown): result is InterruptResult {
+	return (
+		typeof result === "object" &&
+		result !== null &&
+		"__interrupt__" in result &&
+		(result as InterruptResult).__interrupt__ === true
+	);
+}
+
+function isErrorResult(result: unknown): result is ErrorResult {
+	return (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as ErrorResult).status === "error"
+	);
+}
+
+function isNotFoundResult(result: unknown): result is NotFoundResult {
+	return (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as NotFoundResult).status === "not_found"
+	);
+}
+
+function isAuthErrorResult(result: unknown): result is AuthErrorResult {
+	return (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as AuthErrorResult).status === "auth_error"
+	);
+}
+
+function isInsufficientPermissionsResult(
+	result: unknown,
+): result is InsufficientPermissionsResult {
+	return (
+		typeof result === "object" &&
+		result !== null &&
+		"status" in result &&
+		(result as InsufficientPermissionsResult).status ===
+			"insufficient_permissions"
+	);
+}
+
+function ApprovalCard({
+	args,
+	interruptData,
+	onDecision,
+}: {
+	args: {
+		draft_subject_or_id: string;
+		body: string;
+		to?: string;
+		subject?: string;
+		cc?: string;
+		bcc?: string;
+	};
+	interruptData: InterruptResult;
+	onDecision: (decision: {
+		type: "approve" | "reject" | "edit";
+		message?: string;
+		edited_action?: { name: string; args: Record<string, unknown> };
+	}) => void;
+}) {
+	const [decided, setDecided] = useState<
+		"approve" | "reject" | "edit" | null
+	>(interruptData.__decided__ ?? null);
+	const [wasAlreadyDecided] = useState(
+		() => interruptData.__decided__ != null,
+	);
+	const [isPanelOpen, setIsPanelOpen] = useState(false);
+	const openHitlEditPanel = useSetAtom(openHitlEditPanelAtom);
+	const [pendingEdits, setPendingEdits] = useState<{
+		subject: string;
+		body: string;
+		to: string;
+		cc: string;
+		bcc: string;
+	} | null>(null);
+
+	const account = interruptData.context?.account;
+	const email = interruptData.context?.email;
+	const draftId = interruptData.context?.draft_id;
+
+	const reviewConfig = interruptData.review_configs[0];
+	const allowedDecisions = reviewConfig?.allowed_decisions ?? [
+		"approve",
+		"reject",
+	];
+	const canEdit = allowedDecisions.includes("edit");
+
+	const currentSubject =
+		pendingEdits?.subject ??
+		args.subject ??
+		email?.subject ??
+		args.draft_subject_or_id;
+	const currentBody = pendingEdits?.body ?? args.body;
+	const currentTo = pendingEdits?.to ?? args.to ?? "";
+	const currentCc = pendingEdits?.cc ?? args.cc ?? "";
+	const currentBcc = pendingEdits?.bcc ?? args.bcc ?? "";
+
+	const handleApprove = useCallback(() => {
+		if (decided || isPanelOpen) return;
+		if (!allowedDecisions.includes("approve")) return;
+		const isEdited = pendingEdits !== null;
+		setDecided(isEdited ? "edit" : "approve");
+		onDecision({
+			type: isEdited ? "edit" : "approve",
+			edited_action: {
+				name: interruptData.action_requests[0].name,
+				args: {
+					message_id: email?.message_id,
+					draft_id: draftId,
+					to: currentTo,
+					subject: currentSubject,
+					body: currentBody,
+					cc: currentCc,
+					bcc: currentBcc,
+					connector_id: email?.connector_id ?? account?.id,
+				},
+			},
+		});
+	}, [
+		decided,
+		isPanelOpen,
+		allowedDecisions,
+		onDecision,
+		interruptData,
+		email,
+		account?.id,
+		draftId,
+		pendingEdits,
+		currentSubject,
+		currentBody,
+		currentTo,
+		currentCc,
+		currentBcc,
+	]);
+
+	useEffect(() => {
+		const handler = (e: KeyboardEvent) => {
+			if (e.key === "Enter" && !e.shiftKey && !e.ctrlKey && !e.metaKey) {
+				handleApprove();
+			}
+		};
+		window.addEventListener("keydown", handler);
+		return () => window.removeEventListener("keydown", handler);
+	}, [handleApprove]);
+
+	return (
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 transition-all duration-300">
+			{/* Header */}
+			<div className="flex items-start justify-between px-5 pt-5 pb-4 select-none">
+				<div className="flex items-center gap-2">
+					<div>
+						<p className="text-sm font-semibold text-foreground">
+							{decided === "reject"
+								? "Draft Update Rejected"
+								: decided === "approve" || decided === "edit"
+									? "Draft Update Approved"
+									: "Update Gmail Draft"}
+						</p>
+						{decided === "approve" || decided === "edit" ? (
+							wasAlreadyDecided ? (
+								<p className="text-xs text-muted-foreground mt-0.5">
+									{decided === "edit"
+										? "Draft updated with your changes"
+										: "Draft updated"}
+								</p>
+							) : (
+								<TextShimmerLoader
+									text={
+										decided === "edit"
+											? "Updating draft with your changes"
+											: "Updating draft"
+									}
+									size="sm"
+								/>
+							)
+						) : (
+							<p className="text-xs text-muted-foreground mt-0.5">
+								{decided === "reject"
+									? "Draft update was cancelled"
+									: "Requires your approval to proceed"}
+							</p>
+						)}
+					</div>
+				</div>
+				{!decided && canEdit && (
+					<Button
+						size="sm"
+						variant="ghost"
+						className="rounded-lg text-muted-foreground -mt-1 -mr-2"
+						onClick={() => {
+							setIsPanelOpen(true);
+							const extraFields: ExtraField[] = [
+								{
+									key: "to",
+									label: "To",
+									type: "emails",
+									value: currentTo,
+								},
+								{
+									key: "cc",
+									label: "CC",
+									type: "emails",
+									value: currentCc,
+								},
+								{
+									key: "bcc",
+									label: "BCC",
+									type: "emails",
+									value: currentBcc,
+								},
+							];
+							openHitlEditPanel({
+								title: currentSubject,
+								content: currentBody,
+								toolName: "Gmail Draft",
+								extraFields,
+								onSave: (
+									newTitle,
+									newContent,
+									extraFieldValues,
+								) => {
+									setIsPanelOpen(false);
+									const extras = extraFieldValues ?? {};
+									setPendingEdits({
+										subject: newTitle,
+										body: newContent,
+										to: extras.to ?? currentTo,
+										cc: extras.cc ?? currentCc,
+										bcc: extras.bcc ?? currentBcc,
+									});
+								},
+							});
+						}}
+					>
+						<Pen className="size-3.5" />
+						Edit
+					</Button>
+				)}
+			</div>
+
+			{/* Context — account and current draft info */}
+			{!decided && interruptData.context && (
+				<>
+					<div className="mx-5 h-px bg-border/50" />
+					<div className="px-5 py-4 space-y-4 select-none">
+						{interruptData.context.error ? (
+							<p className="text-sm text-destructive">
+								{interruptData.context.error}
+							</p>
+						) : (
+							<>
+								{account && (
+									<div className="space-y-2">
+										<p className="text-xs font-medium text-muted-foreground">
+											Gmail Account
+										</p>
+										<div className="w-full rounded-md border border-input bg-muted/50 px-3 py-2 text-sm">
+											{account.name}
+										</div>
+									</div>
+								)}
+
+								{email && (
+									<div className="space-y-2">
+										<p className="text-xs font-medium text-muted-foreground">
+											Draft to Update
+										</p>
+										<div className="w-full rounded-md border border-input bg-muted/50 px-3 py-2 text-sm space-y-1">
+											<div className="flex items-center gap-1.5">
+												<MailIcon className="size-3 shrink-0 text-muted-foreground" />
+												<span className="font-medium">
+													{email.subject}
+												</span>
+											</div>
+										</div>
+									</div>
+								)}
+							</>
+						)}
+					</div>
+				</>
+			)}
+
+			{/* Email headers + body preview */}
+			<div className="mx-5 h-px bg-border/50" />
+			<div className="px-5 pt-3 pb-2 space-y-1.5 select-none">
+				{currentTo && (
+					<div className="flex items-center gap-1.5 text-xs text-muted-foreground">
+						<UserIcon className="size-3 shrink-0" />
+						<span>To: {currentTo}</span>
+					</div>
+				)}
+				{currentCc && currentCc.trim() !== "" && (
+					<div className="flex items-center gap-1.5 text-xs text-muted-foreground">
+						<UsersIcon className="size-3 shrink-0" />
+						<span>CC: {currentCc}</span>
+					</div>
+				)}
+				{currentBcc && currentBcc.trim() !== "" && (
+					<div className="flex items-center gap-1.5 text-xs text-muted-foreground">
+						<UsersIcon className="size-3 shrink-0" />
+						<span>BCC: {currentBcc}</span>
+					</div>
+				)}
+			</div>
+
+			<div className="px-5 pt-1">
+				{currentSubject != null && (
+					<p className="text-sm font-medium text-foreground">
+						{currentSubject}
+					</p>
+				)}
+				{currentBody != null && (
+					<div
+						className="mt-2 max-h-[7rem] overflow-hidden text-sm"
+						style={{
+							maskImage:
+								"linear-gradient(to bottom, black 50%, transparent 100%)",
+							WebkitMaskImage:
+								"linear-gradient(to bottom, black 50%, transparent 100%)",
+						}}
+					>
+						<PlateEditor
+							markdown={String(currentBody)}
+							readOnly
+							preset="readonly"
+							editorVariant="none"
+							className="h-auto [&_[data-slate-editor]]:!min-h-0 [&_[data-slate-editor]>*:first-child]:!mt-0"
+						/>
+					</div>
+				)}
+			</div>
+
+			{/* Action buttons */}
+			{!decided && (
+				<>
+					<div className="mx-5 h-px bg-border/50" />
+					<div className="px-5 py-4 flex items-center gap-2 select-none">
+						{allowedDecisions.includes("approve") && (
+							<Button
+								size="sm"
+								className="rounded-lg gap-1.5"
+								onClick={handleApprove}
+							>
+								Approve
+								<CornerDownLeftIcon className="size-3 opacity-60" />
+							</Button>
+						)}
+						{allowedDecisions.includes("reject") && (
+							<Button
+								size="sm"
+								variant="ghost"
+								className="rounded-lg text-muted-foreground"
+								onClick={() => {
+									setDecided("reject");
+									onDecision({
+										type: "reject",
+										message: "User rejected the action.",
+									});
+								}}
+							>
+								Reject
+							</Button>
+						)}
+					</div>
+				</>
+			)}
+		</div>
+	);
+}
+
+function ErrorCard({ result }: { result: ErrorResult }) {
+	return (
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+			<div className="px-5 pt-5 pb-4">
+				<p className="text-sm font-semibold text-destructive">
+					Failed to update Gmail draft
+				</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 AuthErrorCard({ result }: { result: AuthErrorResult }) {
+	return (
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+			<div className="px-5 pt-5 pb-4">
+				<p className="text-sm font-semibold text-destructive">
+					Gmail 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 InsufficientPermissionsCard({
+	result,
+}: { result: InsufficientPermissionsResult }) {
+	return (
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border bg-muted/30 select-none">
+			<div className="px-5 pt-5 pb-4">
+				<p className="text-sm font-semibold text-destructive">
+					Additional Gmail permissions required
+				</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 NotFoundCard({ result }: { result: NotFoundResult }) {
+	return (
+		<div className="my-4 max-w-lg overflow-hidden rounded-2xl border border-amber-500/50 bg-muted/30 select-none">
+			<div className="px-5 pt-5 pb-4">
+				<div className="flex items-center gap-2">
+					<TriangleAlertIcon className="size-4 text-amber-500 shrink-0" />
+					<p className="text-sm font-semibold text-amber-600 dark:text-amber-400">
+						Draft not found
+					</p>
+				</div>
+			</div>
+			<div className="mx-5 h-px bg-amber-500/30" />
+			<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 select-none">
+			<div className="px-5 pt-5 pb-4">
+				<p className="text-sm font-semibold text-foreground">
+					{result.message || "Gmail draft updated successfully"}
+				</p>
+			</div>
+		</div>
+	);
+}
+
+export const UpdateGmailDraftToolUI = makeAssistantToolUI<
+	{
+		draft_subject_or_id: string;
+		body: string;
+		to?: string;
+		subject?: string;
+		cc?: string;
+		bcc?: string;
+	},
+	UpdateGmailDraftResult
+>({
+	toolName: "update_gmail_draft",
+	render: function UpdateGmailDraftUI({ args, result }) {
+		if (!result) return null;
+
+		if (isInterruptResult(result)) {
+			return (
+				<ApprovalCard
+					args={args}
+					interruptData={result}
+					onDecision={(decision) => {
+						window.dispatchEvent(
+							new CustomEvent("hitl-decision", {
+								detail: { decisions: [decision] },
+							}),
+						);
+					}}
+				/>
+			);
+		}
+
+		if (
+			typeof result === "object" &&
+			result !== null &&
+			"status" in result &&
+			(result as { status: string }).status === "rejected"
+		) {
+			return null;
+		}
+
+		if (isAuthErrorResult(result)) return <AuthErrorCard result={result} />;
+		if (isInsufficientPermissionsResult(result))
+			return <InsufficientPermissionsCard result={result} />;
+		if (isNotFoundResult(result)) return <NotFoundCard result={result} />;
+		if (isErrorResult(result)) return <ErrorCard result={result} />;
+
+		return <SuccessCard result={result as SuccessResult} />;
+	},
+});