Merge pull request #5358 from ethanova/allow-assistant-message-edits

web-app/package.json

@@ -45,6 +45,7 @@
     "fzf": "^0.5.2",
     "i18next": "^25.0.1",
     "katex": "^0.16.22",
+    "lodash.clonedeep": "^4.5.0",
     "lodash.debounce": "^4.0.8",
     "lucide-react": "^0.522.0",
     "motion": "^12.10.5",
@@ -77,6 +78,7 @@
     "@eslint/js": "^9.22.0",
     "@tanstack/router-plugin": "^1.116.1",
     "@types/culori": "^2.1.1",
+    "@types/lodash.clonedeep": "^4",
     "@types/lodash.debounce": "^4",
     "@types/node": "^22.14.1",
     "@types/react": "^19.0.10",

web-app/src/containers/ThreadContent.tsx

@@ -26,7 +26,6 @@ import {
 } from '@/components/ui/dialog'
 import { Button } from '@/components/ui/button'
 import { Textarea } from '@/components/ui/textarea'
-import { toast } from 'sonner'
 import {
   Tooltip,
   TooltipContent,
@@ -75,6 +74,69 @@ const CopyButton = ({ text }: { text: string }) => {
   )
 }
 
+const EditDialog = ({
+  message,
+  setMessage,
+}: {
+  message: string
+  setMessage: (message: string) => void
+}) => {
+  const { t } = useTranslation()
+  const [draft, setDraft] = useState(message)
+
+  const handleSave = () => {
+    if (draft !== message) {
+      setMessage(draft)
+    }
+  }
+
+  return (
+    <Dialog>
+      <DialogTrigger>
+        <Tooltip>
+          <TooltipTrigger asChild>
+            <div className="flex outline-0 items-center gap-1 hover:text-accent transition-colors cursor-pointer group relative">
+              <IconPencil size={16} />
+            </div>
+          </TooltipTrigger>
+          <TooltipContent>
+            <p>{t('edit')}</p>
+          </TooltipContent>
+        </Tooltip>
+      </DialogTrigger>
+      <DialogContent className="w-3/4 h-3/4">
+        <DialogHeader>
+          <DialogTitle>{t('common:dialogs.editMessage.title')}</DialogTitle>
+          <Textarea
+            value={draft}
+            onChange={(e) => setDraft(e.target.value)}
+            className="mt-2 resize-none h-full w-full"
+            onKeyDown={(e) => {
+              // Prevent key from being captured by parent components
+              e.stopPropagation()
+            }}
+          />
+          <DialogFooter className="mt-2 flex items-center">
+            <DialogClose asChild>
+              <Button variant="link" size="sm" className="hover:no-underline">
+                Cancel
+              </Button>
+            </DialogClose>
+            <DialogClose asChild>
+              <Button
+                disabled={draft === message || !draft}
+                onClick={handleSave}
+              >
+                Save
+              </Button>
+            </DialogClose>
+          </DialogFooter>
+        </DialogHeader>
+      </DialogContent>
+    </Dialog>
+  )
+}
+
 // Use memo to prevent unnecessary re-renders, but allow re-renders when props change
 export const ThreadContent = memo(
   (
@@ -85,9 +147,9 @@ export const ThreadContent = memo(
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       streamTools?: any
       contextOverflowModal?: React.ReactNode | null
+      updateMessage?: (item: ThreadMessage, message: string) => void
     }
   ) => {
-    const [message, setMessage] = useState(item.content?.[0]?.text?.value || '')
     const { t } = useTranslation()
 
     // Use useMemo to stabilize the components prop
@@ -166,23 +228,6 @@ export const ThreadContent = memo(
       }
     }, [deleteMessage, getMessages, item])
 
-    const editMessage = useCallback(
-      (messageId: string) => {
-        const threadMessages = getMessages(item.thread_id)
-
-        const index = threadMessages.findIndex((msg) => msg.id === messageId)
-        if (index === -1) return
-
-        // Delete all messages after the edited message
-        for (let i = threadMessages.length - 1; i >= index; i--) {
-          deleteMessage(threadMessages[i].thread_id, threadMessages[i].id)
-        }
-
-        sendMessage(message)
-      },
-      [deleteMessage, getMessages, item.thread_id, message, sendMessage]
-    )
-
     const isToolCalls =
       item.metadata &&
       'tool_calls' in item.metadata &&
@@ -209,61 +254,14 @@ export const ThreadContent = memo(
               </div>
             </div>
             <div className="flex items-center justify-end gap-2 text-main-view-fg/60 text-xs mt-2">
-              <Dialog>
+              <EditDialog
-                <DialogTrigger>
+                message={item.content?.[0]?.text.value}
-                  <Tooltip>
+                setMessage={(message) => {
-                    <TooltipTrigger asChild>
+                  if (item.updateMessage) {
-                      <div className="flex outline-0 items-center gap-1 hover:text-accent transition-colors cursor-pointer group relative">
+                    item.updateMessage(item, message)
-                        <IconPencil size={16} />
+                  }
-                      </div>
-                    </TooltipTrigger>
-                    <TooltipContent>
-                      <p>{t('edit')}</p>
-                    </TooltipContent>
-                  </Tooltip>
-                </DialogTrigger>
-                <DialogContent>
-                  <DialogHeader>
-                    <DialogTitle>{t('common:dialogs.editMessage.title')}</DialogTitle>
-                    <Textarea
-                      value={message}
-                      onChange={(e) => {
-                        setMessage(e.target.value)
-                      }}
-                      className="mt-2 resize-none"
-                      onKeyDown={(e) => {
-                        // Prevent key from being captured by parent components
-                        e.stopPropagation()
                 }}
               />
-                    <DialogFooter className="mt-2 flex items-center">
-                      <DialogClose asChild>
-                        <Button
-                          variant="link"
-                          size="sm"
-                          className="hover:no-underline"
-                        >
-                          Cancel
-                        </Button>
-                      </DialogClose>
-                      <DialogClose asChild>
-                        <Button
-                          disabled={!message}
-                          onClick={() => {
-                            editMessage(item.id)
-                            toast.success(t('common:toast.editMessage.title'), {
-                              id: 'edit-message',
-                              description: t('common:toast.editMessage.description'),
-                            })
-                          }}
-                        >
-                          Save
-                        </Button>
-                      </DialogClose>
-                    </DialogFooter>
-                  </DialogHeader>
-                </DialogContent>
-              </Dialog>
               <Tooltip>
                 <TooltipTrigger asChild>
                   <button
@@ -360,6 +358,12 @@ export const ThreadContent = memo(
                         'hidden'
                     )}
                   >
+                    <EditDialog
+                      message={item.content?.[0]?.text.value}
+                      setMessage={(message) =>
+                        item.updateMessage && item.updateMessage(item, message)
+                      }
+                    />
                     <CopyButton text={item.content?.[0]?.text.value || ''} />
                     <Tooltip>
                       <TooltipTrigger asChild>
@@ -391,7 +395,9 @@ export const ThreadContent = memo(
                       </DialogTrigger>
                       <DialogContent>
                         <DialogHeader>
-                          <DialogTitle>{t('common:dialogs.messageMetadata.title')}</DialogTitle>
+                          <DialogTitle>
+                            {t('common:dialogs.messageMetadata.title')}
+                          </DialogTitle>
                           <div className="space-y-2">
                             <div className="border border-main-view-fg/10 rounded-md overflow-hidden">
                               <CodeEditor

web-app/src/routes/threads/$threadId.tsx

@@ -2,8 +2,10 @@ import { useEffect, useMemo, useRef, useState } from 'react'
 import { createFileRoute, useParams } from '@tanstack/react-router'
 import { UIEventHandler } from 'react'
 import debounce from 'lodash.debounce'
+import cloneDeep from 'lodash.clonedeep'
 import { cn } from '@/lib/utils'
 import { ArrowDown } from 'lucide-react'
+import { Play } from 'lucide-react'
 
 import HeaderPage from '@/containers/HeaderPage'
 import { useThreads } from '@/hooks/useThreads'
@@ -18,7 +20,9 @@ import { useAppState } from '@/hooks/useAppState'
 import DropdownAssistant from '@/containers/DropdownAssistant'
 import { useAssistant } from '@/hooks/useAssistant'
 import { useAppearance } from '@/hooks/useAppearance'
+import { ContentType, ThreadMessage } from '@janhq/core'
 import { useTranslation } from '@/i18n/react-i18next-compat'
+import { useChat } from '@/hooks/useChat'
 import { useSmallScreen } from '@/hooks/useMediaQuery'
 
 // as route.threadsDetail
@@ -38,6 +42,7 @@ function ThreadDetail() {
   const { setMessages } = useMessages()
   const { streamingContent } = useAppState()
   const { appMainViewBgColor, chatWidth } = useAppearance()
+  const { sendMessage } = useChat()
   const isSmallScreen = useSmallScreen()
 
   const { messages } = useMessages(
@@ -180,6 +185,26 @@ function ThreadDetail() {
     lastScrollTopRef.current = scrollTop
   }
 
+  const updateMessage = (item: ThreadMessage, message: string) => {
+    const newMessages: ThreadMessage[] = messages.map((m) => {
+      if (m.id === item.id) {
+        const msg: ThreadMessage = cloneDeep(m)
+        msg.content = [
+          {
+            type: ContentType.Text,
+            text: {
+              value: message,
+              annotations: m.content[0].text?.annotations ?? [],
+            },
+          },
+        ]
+        return msg
+      }
+      return m
+    })
+    setMessages(threadId, newMessages)
+  }
+
   // Use a shorter debounce time for more responsive scrolling
   const debouncedScroll = debounce(handleDOMScroll)
 
@@ -193,10 +218,22 @@ function ThreadDetail() {
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [])
 
+  // used when there is a sent/added user message and no assistant message (error or manual deletion)
+  const generateAIResponse = () => {
+    const latestUserMessage = messages[messages.length - 1]
+    if (latestUserMessage?.content?.[0]?.text?.value) {
+      sendMessage(latestUserMessage.content[0].text.value, false)
+    }
+  }
+
   const threadModel = useMemo(() => thread?.model, [thread])
 
   if (!messages || !threadModel) return null
 
+  const showScrollToBottomBtn = !isAtBottom && hasScrollbar
+  const showGenerateAIResponseBtn =
+    messages[messages.length - 1]?.role === 'user' && !streamingContent
+
   return (
     <div className="flex flex-col h-full">
       <HeaderPage>
@@ -243,6 +280,7 @@ function ThreadDetail() {
                           ))
                       }
                       index={index}
+                      updateMessage={updateMessage}
                     />
                   </div>
                 )
@@ -266,9 +304,11 @@ function ThreadDetail() {
               appMainViewBgColor.a === 1
                 ? 'from-main-view/20 bg-gradient-to-b to-main-view backdrop-blur'
                 : 'bg-transparent',
-              !isAtBottom && hasScrollbar && 'visibility-visible opacity-100'
+              (showScrollToBottomBtn || showGenerateAIResponseBtn) &&
+                'visibility-visible opacity-100'
             )}
           >
+            {showScrollToBottomBtn && (
               <div
                 className="bg-main-view-fg/10 px-4 border border-main-view-fg/5 flex items-center justify-center rounded-xl gap-x-2 cursor-pointer pointer-events-auto"
                 onClick={() => {
@@ -279,6 +319,16 @@ function ThreadDetail() {
                 <p className="text-xs">{t('scrollToBottom')}</p>
                 <ArrowDown size={12} />
               </div>
+            )}
+            {showGenerateAIResponseBtn && (
+              <div
+                className="bg-main-view-fg/10 px-4 border border-main-view-fg/5 flex items-center justify-center rounded-xl gap-x-2 cursor-pointer pointer-events-auto"
+                onClick={generateAIResponse}
+              >
+                <p className="text-xs">{t('Generate AI Response')}</p>
+                <Play size={12} />
+              </div>
+            )}
           </div>
           <ChatInput model={threadModel} />
         </div>