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

import { UIMessage as VercelMessage } from '@ai-sdk/react'
import { memo, useMemo, useState } from 'react'
import { toast } from 'sonner'
import { cn, copyToClipboard } from 'ui'

import { DeleteMessageConfirmModal } from './DeleteMessageConfirmModal'
import { MessageActions } from './Message.Actions'
import type { AddToolApprovalResponse, MessageInfo } from './Message.Context'
import { MessageProvider, useMessageActionsContext, useMessageInfoContext } from './Message.Context'
import { MessageDisplay } from './Message.Display'

function AssistantMessage({ message }: { message: VercelMessage }) {
  const { onBranch, onCancelEdit, onRate } = useMessageActionsContext()
  const { id, variant, state, isLastMessage, readOnly, rating, isLoading } = useMessageInfoContext()

  const handleRate = (newRating: 'positive' | 'negative', reason?: string) => {
    onRate?.(id, newRating, reason)
  }

  const handleCopy = (onSuccess: () => void) => {
    const response = message.parts
      .filter((x) => x.type === 'text')
      .map((x) => x.text)
      .join('\n')
    copyToClipboard(response, onSuccess)
  }

  return (
    <MessageDisplay.Container
      className={cn(
        variant === 'warning' && 'bg-warning-200',
        state === 'predecessor-editing' && 'opacity-50 transition-opacity cursor-pointer'
      )}
      onClick={state === 'predecessor-editing' ? onCancelEdit : undefined}
    >
      <MessageDisplay.MainArea>
        <MessageDisplay.Content message={message} />
      </MessageDisplay.MainArea>
      {!readOnly && onRate && !isLoading && (
        <MessageActions alwaysShow={isLastMessage}>
          <MessageActions.Copy onClick={handleCopy} />
          <MessageActions.ThumbsUp
            onClick={() => handleRate('positive')}
            isActive={rating === 'positive'}
            disabled={!!rating}
          />
          <MessageActions.ThumbsDown
            onClick={(reason) => handleRate('negative', reason)}
            isActive={rating === 'negative'}
            disabled={!!rating}
          />
          <MessageActions.Branch onClick={() => onBranch(id)} />
        </MessageActions>
      )}
    </MessageDisplay.Container>
  )
}

function UserMessage({ message }: { message: VercelMessage }) {
  const { id, variant, state } = useMessageInfoContext()
  const { onCancelEdit, onEdit, onDelete } = useMessageActionsContext()
  const [showDeleteConfirmModal, setShowDeleteConfirmModal] = useState(false)

  return (
    <>
      <MessageDisplay.Container
        className={cn(
          'mt-6 text-foreground',
          variant === 'warning' && 'bg-warning-200',
          state === 'predecessor-editing' && 'opacity-50 transition-opacity cursor-pointer'
        )}
        onClick={state === 'predecessor-editing' ? onCancelEdit : undefined}
      >
        <MessageDisplay.MainArea className="w-full max-w-3xl mx-auto">
          <MessageDisplay.ProfileImage />
          <MessageDisplay.Content message={message} />
        </MessageDisplay.MainArea>
        <MessageActions>
          <MessageActions.Edit
            onClick={state === 'idle' ? () => onEdit(id) : onCancelEdit}
            tooltip={state === 'idle' ? 'Edit message' : 'Cancel editing'}
          />
          <MessageActions.Delete onClick={() => setShowDeleteConfirmModal(true)} />
        </MessageActions>
      </MessageDisplay.Container>
      <DeleteMessageConfirmModal
        visible={showDeleteConfirmModal}
        onConfirm={() => {
          onDelete(id)
          setShowDeleteConfirmModal(false)
          toast.success('Message deleted successfully')
        }}
        onCancel={() => setShowDeleteConfirmModal(false)}
      />
    </>
  )
}

interface MessageProps {
  id: string
  message: VercelMessage
  isLoading: boolean
  readOnly?: boolean
  variant?: 'default' | 'warning'
  addToolApprovalResponse?: AddToolApprovalResponse
  onDelete: (id: string) => void
  onEdit: (id: string) => void
  onBranch: (id: string) => void
  isAfterEditedMessage: boolean
  isBeingEdited: boolean
  onCancelEdit: () => void
  isLastMessage?: boolean
  onRate?: (id: string, rating: 'positive' | 'negative', reason?: string) => void
  rating?: 'positive' | 'negative' | null
}

export const Message = memo(function Message(props: MessageProps) {
  const message = props.message
  const { role } = message
  const isUserMessage = role === 'user'
  let messageState: MessageInfo['state'] = 'idle'
  if (props.isBeingEdited) messageState = 'editing'
  else if (props.isAfterEditedMessage) messageState = 'predecessor-editing'

  const messageInfo = useMemo<MessageInfo>(
    () => ({
      id: props.id,
      isLoading: props.isLoading,
      readOnly: props.readOnly,
      variant: props.variant,
      isUserMessage,
      state: messageState,
      isLastMessage: props.isLastMessage,
      rating: props.rating,
    }),
    [
      props.id,
      props.isLoading,
      props.readOnly,
      props.variant,
      isUserMessage,
      messageState,
      props.isLastMessage,
      props.rating,
    ]
  )

  const messageActions = useMemo(
    () => ({
      addToolApprovalResponse: props.addToolApprovalResponse,
      onDelete: props.onDelete,
      onEdit: props.onEdit,
      onBranch: props.onBranch,
      onCancelEdit: props.onCancelEdit,
      onRate: props.onRate,
    }),
    [
      props.addToolApprovalResponse,
      props.onDelete,
      props.onEdit,
      props.onBranch,
      props.onCancelEdit,
      props.onRate,
    ]
  )

  return (
    <MessageProvider messageInfo={messageInfo} messageActions={messageActions}>
      {isUserMessage ? <UserMessage message={message} /> : <AssistantMessage message={message} />}
    </MessageProvider>
  )
})