Add message deletion functionality to AI Assistant (supabase#37698)

* add delete chat message functionality

* Implement message editing functionality in AIAssistant

* add functionality to delete messages in AI Assistant chat in assistant state

* remove duplicated `text-xs` class

* add a 'group' class for improved styling and refining action button visibility on hover

* Refine action button visibility in AIAssistantPanel to show only for user messages on hover, improving user experience.

* remove `"` from the message to edit

* Enhance editing message display in AIAssistantPanel by updating styling for improved visibility and user experience.

* Refactor AIAssistantPanel state management by simplifying original message content initialization and clarifying comments for better code readability.

* clean up action button rendering in AIAssistantPanel message component

* adjust the chat message layout (space)

* simplify the editing mode message

* show the edit button even when user is already in the edit mode

* Update AssistantChatForm styling for editing mode, positioning the editing message above the input field without changing the layout

* use `ButtonTooltip` for edit button to show tooltip

* simplify the condiiton to show the edit buton

* add the delete action button

* Refactor action button rendering in AIAssistantPanel to improve clarity and maintainability. Ensure buttons are only shown for user messages on hover, and simplify the conditional rendering logic.

* Remove commented-out code for cleaner implementation.

* improve unused onDeleteAfter prop and consolidating delete logic into a single `deleteMessageFromHere` function for improved clarity

* abort streaming when a user deletes  a message

* delete unnecessary `originalMessageContent` state variable

* delete `confirmEdit` and move its logic into `sendMessageToAssistant`

* add `strokeWidth={1.5}` to `Pencil` icon

* add `opacity-50` to any messages that are after the message being edited, and it cancels the edit when those messages are clicked

* add a pulse animation to the user’s message being edited

* add more space between mesages and `AssistantChatForm`

* add framer-motion animation to `Editing message` element

* update strokeWidth of Trash2 icon in Message component for better visibility

* refactor Message component to make onDelete prop required

* add `DeleteMessageConfirmModal` for message deletion confirmation in Message component

* refactor message deletion logic in AIAssistant component to streamline editing process

* cancel editing when a user clicks the edit button again or the edit button in subsequent messages

* Add success toast notification for message deletion

* add Pencil icon to AIAssistantPanel component and fixed duplicated tailwind class

* Update AIAssistantPanel to use the Pencil icon

* Refactor AIAssistant to set editing message content from text parts

* Remove unused `reload` prop from AIAssistant component

* `setIsSticky(false)` when a user deletes all the messages or clear the chat to remove the `ArrowDown` icon scroll down button

* refactor: remove deleteMessagesBefore function from AiAssistantState

* refactor: simplify AssistantChatForm layout by removing unnecessary `z-index`

* enable the input field when a user wants to edit the message while the response is streaming

* Smol fixes

* remove conosle log

* Add state for resubmitting in AI Assistant panel

* Remove classes that hide the focus styles + fix focusing when editing message to ensure cursor is at the end of the text

---------

Co-authored-by: Joshen Lim <[email protected]>

Author: taishikato

apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx

@@ -42,6 +42,7 @@ const MemoizedMessage = memo(
     message,
     isLoading,
     onResults,
+    onDelete,
     onEdit,
     isAfterEditedMessage,
     isBeingEdited,
@@ -58,6 +59,7 @@ const MemoizedMessage = memo(
       resultId?: string
       results: any[]
     }) => void
+    onDelete: (id: string) => void
     onEdit: (id: string) => void
     isAfterEditedMessage: boolean
     isBeingEdited: boolean
@@ -70,6 +72,7 @@ const MemoizedMessage = memo(
         readOnly={message.role === 'user'}
         isLoading={isLoading}
         onResults={onResults}
+        onDelete={onDelete}
         onEdit={onEdit}
         isAfterEditedMessage={isAfterEditedMessage}
         isBeingEdited={isBeingEdited}
@@ -105,7 +108,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   )
 
   const inputRef = useRef<HTMLTextAreaElement>(null)
-  const { ref: scrollContainerRef, isSticky, scrollToEnd } = useAutoScroll()
+  const { ref: scrollContainerRef, isSticky, scrollToEnd, setIsSticky } = useAutoScroll()
 
   const { aiOptInLevel, isHipaaProjectDisallowed } = useOrgAiOptInLevel()
   const showMetadataWarning =
@@ -119,6 +122,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   const [value, setValue] = useState<string>(snap.initialInput || '')
   const [isConfirmOptInModalOpen, setIsConfirmOptInModalOpen] = useState(false)
   const [editingMessageId, setEditingMessageId] = useState<string | null>(null)
+  const [isResubmitting, setIsResubmitting] = useState(false)
 
   const { data: check, isSuccess } = useCheckOpenAIKeyQuery()
   const isApiKeySet = IS_PLATFORM || !!check?.hasKey
@@ -232,7 +236,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
     onFinish: handleChatFinish,
   })
 
-  const isChatLoading = chatStatus === 'submitted' || chatStatus === 'streaming'
+  const isChatLoading = chatStatus === 'submitted' || chatStatus === 'streaming' || isResubmitting
 
   const updateMessage = useCallback(
     ({
@@ -249,6 +253,22 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
     [snap]
   )
 
+  const deleteMessageFromHere = useCallback(
+    (messageId: string) => {
+      // Find the message index in current chatMessages
+      const messageIndex = chatMessages.findIndex((msg) => msg.id === messageId)
+      if (messageIndex === -1) return
+
+      if (isChatLoading) stop()
+
+      snap.deleteMessagesAfter(messageId, { includeSelf: true })
+
+      const updatedMessages = chatMessages.slice(0, messageIndex)
+      setMessages(updatedMessages)
+    },
+    [snap, setMessages, chatMessages, isChatLoading, stop]
+  )
+
   const editMessage = useCallback(
     (messageId: string) => {
       const messageIndex = chatMessages.findIndex((msg) => msg.id === messageId)
@@ -266,9 +286,16 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
           .join('') ?? ''
       setValue(textContent)
 
-      if (inputRef.current) {
-        inputRef.current.focus()
-      }
+      setTimeout(() => {
+        if (inputRef.current) {
+          inputRef?.current?.focus()
+
+          // [Joshen] This is just to make the cursor go to the end of the text when focusing
+          const val = inputRef.current.value
+          inputRef.current.value = ''
+          inputRef.current.value = val
+        }
+      }, 100)
     },
     [chatMessages, setValue]
   )
@@ -292,14 +319,23 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
             message={message}
             isLoading={isChatLoading && message.id === chatMessages[chatMessages.length - 1].id}
             onResults={updateMessage}
+            onDelete={deleteMessageFromHere}
             onEdit={editMessage}
             isAfterEditedMessage={isAfterEditedMessage}
             isBeingEdited={isBeingEdited}
             onCancelEdit={cancelEdit}
           />
         )
       }),
-    [chatMessages, isChatLoading, updateMessage, editMessage, editingMessageId, cancelEdit]
+    [
+      chatMessages,
+      isChatLoading,
+      updateMessage,
+      deleteMessageFromHere,
+      editMessage,
+      cancelEdit,
+      editingMessageId,
+    ]
   )
 
   const hasMessages = chatMessages.length > 0
@@ -308,12 +344,9 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   const sendMessageToAssistant = (finalContent: string) => {
     if (editingMessageId) {
       // Handling when the user is in edit mode
-      const messageIndex = chatMessages.findIndex((msg) => msg.id === editingMessageId)
-      if (messageIndex === -1) return
-
-      snap.deleteMessagesAfter(editingMessageId, { includeSelf: true })
-      const updatedMessages = chatMessages.slice(0, messageIndex)
-      setMessages(updatedMessages)
+      // delete the message(s) from the chat just like the delete button
+      setIsResubmitting(true)
+      deleteMessageFromHere(editingMessageId)
       setEditingMessageId(null)
     }
 
@@ -355,6 +388,14 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
     setEditingMessageId(null)
   }
 
+  useEffect(() => {
+    // Keep "Thinking" visible while stopping and resubmitting during edit
+    // Only clear once the new response actually starts streaming (or errors)
+    if (isResubmitting && (chatStatus === 'streaming' || !!error)) {
+      setIsResubmitting(false)
+    }
+  }, [isResubmitting, chatStatus, error])
+
   // Update scroll behavior for new messages
   useEffect(() => {
     if (!isChatLoading) {
@@ -381,6 +422,10 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [snap.open, isInSQLEditor, snippetContent])
 
+  useEffect(() => {
+    if (chatMessages.length === 0 && !isSticky) setIsSticky(true)
+  }, [chatMessages, isSticky, setIsSticky])
+
   return (
     <ErrorBoundary
       message="Something went wrong with the AI Assistant"
@@ -699,7 +744,8 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
                 'z-20 [&>form>textarea]:text-base [&>form>textarea]:md:text-sm [&>form>textarea]:border-1 [&>form>textarea]:rounded-md [&>form>textarea]:!outline-none [&>form>textarea]:!ring-offset-0 [&>form>textarea]:!ring-0'
               )}
               loading={isChatLoading}
-              disabled={!isApiKeySet || disablePrompts || isChatLoading}
+              isEditing={!!editingMessageId}
+              disabled={!isApiKeySet || disablePrompts || (isChatLoading && !editingMessageId)}
               placeholder={
                 hasMessages
                   ? 'Ask a follow up question...'

apps/studio/components/ui/AIAssistantPanel/AssistantChatForm.tsx

@@ -41,6 +41,8 @@ export interface FormProps {
   snippetsClassName?: string
   /* Additional class name for the form wrapper */
   className?: string
+  /* If currently editing an existing message */
+  isEditing?: boolean
 }
 
 const AssistantChatFormComponent = forwardRef<HTMLFormElement, FormProps>(
@@ -59,6 +61,7 @@ const AssistantChatFormComponent = forwardRef<HTMLFormElement, FormProps>(
       snippetsClassName,
       includeSnippetsInMessage = false,
       className,
+      isEditing = false,
       ...props
     },
     ref
@@ -68,7 +71,7 @@ const AssistantChatFormComponent = forwardRef<HTMLFormElement, FormProps>(
 
     const handleSubmit = (event?: FormEvent<HTMLFormElement>) => {
       if (event) event.preventDefault()
-      if (!value || loading) return
+      if (!value || (loading && !isEditing)) return
 
       let finalMessage = value
       if (includeSnippetsInMessage && sqlSnippets && sqlSnippets.length > 0) {
@@ -107,8 +110,8 @@ const AssistantChatFormComponent = forwardRef<HTMLFormElement, FormProps>(
             />
           )}
           <ExpandingTextArea
+            autoFocus={!isMobile}
             ref={textAreaRef}
-            autoFocus={isMobile}
             disabled={disabled}
             className={cn(
               'text-sm pr-10 max-h-64',

apps/studio/components/ui/AIAssistantPanel/DeleteMessageConfirmModal.tsx

@@ -0,0 +1,54 @@
+import {
+  Button,
+  Dialog,
+  DialogContent,
+  DialogFooter,
+  DialogHeader,
+  DialogSection,
+  DialogSectionSeparator,
+  DialogTitle,
+} from 'ui'
+
+type DeleteMessageConfirmModalProps = {
+  visible: boolean
+  onConfirm: () => void
+  onCancel: () => void
+}
+
+export const DeleteMessageConfirmModal = ({
+  visible,
+  onConfirm,
+  onCancel,
+}: DeleteMessageConfirmModalProps) => {
+  const onOpenChange = (open: boolean) => {
+    if (!open) onCancel()
+  }
+
+  return (
+    <Dialog open={visible} onOpenChange={onOpenChange}>
+      <DialogContent size="small">
+        <DialogHeader padding="small">
+          <DialogTitle>Delete Message</DialogTitle>
+        </DialogHeader>
+
+        <DialogSectionSeparator />
+
+        <DialogSection padding="small">
+          <p className="text-sm text-foreground-light">
+            Are you sure you want to delete this message and all subsequent messages? This action
+            cannot be undone.
+          </p>
+        </DialogSection>
+
+        <DialogFooter padding="small">
+          <Button type="default" onClick={onCancel}>
+            Cancel
+          </Button>
+          <Button type="danger" onClick={onConfirm}>
+            Delete
+          </Button>
+        </DialogFooter>
+      </DialogContent>
+    </Dialog>
+  )
+}

apps/studio/components/ui/AIAssistantPanel/Message.tsx

@@ -1,16 +1,18 @@
 import { UIMessage as VercelMessage } from '@ai-sdk/react'
-import { Loader2, Pencil } from 'lucide-react'
-import { createContext, PropsWithChildren, ReactNode, useMemo } from 'react'
+import { Loader2, Pencil, Trash2 } from 'lucide-react'
+import { createContext, PropsWithChildren, ReactNode, useMemo, useState } from 'react'
 import ReactMarkdown from 'react-markdown'
 import { Components } from 'react-markdown/lib/ast-to-react'
 import remarkGfm from 'remark-gfm'
+import { toast } from 'sonner'
 
 import { ProfileImage } from 'components/ui/ProfileImage'
 import { useProfile } from 'lib/profile'
 import { cn, markdownComponents, WarningIcon } from 'ui'
 import { ButtonTooltip } from '../ButtonTooltip'
 import { EdgeFunctionBlock } from '../EdgeFunctionBlock/EdgeFunctionBlock'
 import { DisplayBlockRenderer } from './DisplayBlockRenderer'
+import { DeleteMessageConfirmModal } from './DeleteMessageConfirmModal'
 import {
   Heading3,
   Hyperlink,
@@ -51,6 +53,7 @@ interface MessageProps {
     resultId?: string
     results: any[]
   }) => void
+  onDelete: (id: string) => void
   onEdit: (id: string) => void
   isAfterEditedMessage: boolean
   isBeingEdited: boolean
@@ -65,12 +68,14 @@ export const Message = function Message({
   action = null,
   variant = 'default',
   onResults,
+  onDelete,
   onEdit,
   isAfterEditedMessage = false,
   isBeingEdited = false,
   onCancelEdit,
 }: PropsWithChildren<MessageProps>) {
   const { profile } = useProfile()
+  const [showDeleteConfirmModal, setShowDeleteConfirmModal] = useState(false)
   const allMarkdownComponents: Partial<Components> = useMemo(
     () => ({
       ...markdownComponents,
@@ -237,30 +242,54 @@ export const Message = function Message({
               <span className="text-foreground-lighter italic">Assistant is thinking...</span>
             )}
 
-            {/* Action button - only show for user messages on hover */}
-            <div className="opacity-0 group-hover:opacity-100 transition-opacity">
+            {/* Action buttons - only show for user messages on hover */}
+            <div className="opacity-0 group-hover:opacity-100 transition-opacity flex items-center gap-1">
               {message.role === 'user' && (
-                <ButtonTooltip
-                  type="text"
-                  icon={<Pencil size={14} strokeWidth={1.5} />}
-                  onClick={isBeingEdited || isAfterEditedMessage ? onCancelEdit : () => onEdit(id)}
-                  className="text-foreground-light hover:text-foreground p-1 rounded"
-                  aria-label={
-                    isBeingEdited || isAfterEditedMessage ? 'Cancel editing' : 'Edit message'
-                  }
-                  tooltip={{
-                    content: {
-                      side: 'bottom',
-                      text:
-                        isBeingEdited || isAfterEditedMessage ? 'Cancel editing' : 'Edit message',
-                    },
-                  }}
-                />
+                <>
+                  <ButtonTooltip
+                    type="text"
+                    icon={<Pencil size={14} strokeWidth={1.5} />}
+                    onClick={
+                      isBeingEdited || isAfterEditedMessage ? onCancelEdit : () => onEdit(id)
+                    }
+                    className="text-foreground-light hover:text-foreground p-1 rounded"
+                    aria-label={
+                      isBeingEdited || isAfterEditedMessage ? 'Cancel editing' : 'Edit message'
+                    }
+                    tooltip={{
+                      content: {
+                        side: 'bottom',
+                        text:
+                          isBeingEdited || isAfterEditedMessage ? 'Cancel editing' : 'Edit message',
+                      },
+                    }}
+                  />
+
+                  <ButtonTooltip
+                    type="text"
+                    icon={<Trash2 size={14} strokeWidth={1.5} />}
+                    tooltip={{ content: { side: 'bottom', text: 'Delete message' } }}
+                    onClick={() => setShowDeleteConfirmModal(true)}
+                    className="text-foreground-light hover:text-foreground p-1 rounded"
+                    title="Delete message"
+                    aria-label="Delete message"
+                  />
+                </>
               )}
             </div>
           </div>
         </div>
       </div>
+
+      <DeleteMessageConfirmModal
+        visible={showDeleteConfirmModal}
+        onConfirm={() => {
+          onDelete(id)
+          setShowDeleteConfirmModal(false)
+          toast.success('Message deleted successfully')
+        }}
+        onCancel={() => setShowDeleteConfirmModal(false)}
+      />
     </MessageContext.Provider>
   )
 }

apps/studio/components/ui/AIAssistantPanel/hooks.ts

@@ -69,5 +69,5 @@ export function useAutoScroll({ enabled = true }: UseAutoScrollProps = {}) {
     }
   }, [container, enabled, scrollToEnd])
 
-  return { ref, isSticky, scrollToEnd }
+  return { ref, isSticky, scrollToEnd, setIsSticky }
 }