Merge pull request #3747 from liam-hq/fix_human_message

Fix optimistic update not working on reload for first HumanMessage

Author: sasamuku

frontend/apps/app/components/SessionDetailPage/SessionDetailPageClient.tsx

@@ -1,6 +1,7 @@
 'use client'
 
 import {
+  HumanMessage,
   mapStoredMessagesToChatMessages,
   type StoredMessage,
 } from '@langchain/core/messages'
@@ -128,11 +129,11 @@ export const SessionDetailPageClient: FC<Props> = ({
     (artifact !== null || selectedVersion !== null) && activeTab
 
   const chatMessages = mapStoredMessagesToChatMessages(initialMessages)
-  const { isStreaming, messages, start, replay, error } = useStream({
-    initialMessages: chatMessages,
-    designSessionId,
-    senderName,
-  })
+  const { isStreaming, messages, setMessages, start, replay, error } =
+    useStream({
+      initialMessages: chatMessages,
+      designSessionId,
+    })
 
   const handleLayoutChange = useCallback((sizes: number[]) => {
     setCookieJson(PANEL_LAYOUT_COOKIE_NAME, sizes, {
@@ -146,6 +147,31 @@ export const SessionDetailPageClient: FC<Props> = ({
   // Track if initial workflow has been triggered to prevent multiple executions
   const hasTriggeredInitialWorkflow = useRef(false)
 
+  const handleSendMessage = useCallback(
+    async (content: string) => {
+      const tempId = `optimistic-${crypto.randomUUID()}`
+      const optimisticMessage = new HumanMessage({
+        content,
+        id: tempId,
+        additional_kwargs: {
+          userName: senderName,
+        },
+      })
+      setMessages((prev) => [...prev, optimisticMessage])
+
+      const result = await start({
+        userInput: content,
+        designSessionId,
+        isDeepModelingEnabled,
+      })
+
+      if (result.isErr()) {
+        setMessages((prev) => prev.filter((msg) => msg.id !== tempId))
+      }
+    },
+    [setMessages, start, senderName, designSessionId, isDeepModelingEnabled],
+  )
+
   // Auto-trigger workflow on page load if there's an unanswered user message
   useEffect(() => {
     const triggerInitialWorkflow = async () => {
@@ -201,13 +227,7 @@ export const SessionDetailPageClient: FC<Props> = ({
                 schemaData={displayedSchema}
                 messages={messages}
                 isWorkflowRunning={isStreaming}
-                onSendMessage={(content: string) =>
-                  start({
-                    userInput: content,
-                    designSessionId,
-                    isDeepModelingEnabled,
-                  })
-                }
+                onSendMessage={handleSendMessage}
                 onNavigate={setActiveTab}
                 error={combinedError}
               />

frontend/apps/app/components/SessionDetailPage/hooks/useStream/useStream.ts

@@ -3,7 +3,6 @@
 import {
   type BaseMessage,
   coerceMessageLikeToMessage,
-  HumanMessage,
 } from '@langchain/core/messages'
 import { MessageTupleManager, SSE_EVENTS } from '@liam-hq/agent/client'
 import { err, ok, type Result } from 'neverthrow'
@@ -63,18 +62,11 @@ const extractStreamErrorMessage = (rawData: unknown): string => {
 type Props = {
   designSessionId: string
   initialMessages: BaseMessage[]
-  senderName: string
 }
-export const useStream = ({
-  designSessionId,
-  initialMessages,
-  senderName,
-}: Props) => {
+export const useStream = ({ designSessionId, initialMessages }: Props) => {
   const messageManagerRef = useRef(new MessageTupleManager())
   const storedMessage = useSessionStorageOnce(designSessionId)
 
-  const isFirstMessage = useRef(true)
-
   const processedInitialMessages = useMemo(() => {
     if (storedMessage) {
       return [storedMessage, ...initialMessages]
@@ -268,30 +260,12 @@ export const useStream = ({
       abortRef.current?.abort()
       retryCountRef.current = 0
 
-      let tempId: string | undefined
-      if (!isFirstMessage.current) {
-        tempId = `optimistic-${crypto.randomUUID()}`
-        const optimisticMessage = new HumanMessage({
-          content: params.userInput,
-          id: tempId,
-          additional_kwargs: {
-            userName: senderName,
-          },
-        })
-        setMessages((prev) => [...prev, optimisticMessage])
-      } else {
-        isFirstMessage.current = false
-      }
-
       // Set workflow in progress flag
       setWorkflowInProgress(params.designSessionId)
 
       const result = await runStreamAttempt('/api/chat/stream', params)
 
       if (result.isErr()) {
-        if (tempId) {
-          setMessages((prev) => prev.filter((msg) => msg.id !== tempId))
-        }
         return err(result.error)
       }
 
@@ -304,11 +278,12 @@ export const useStream = ({
         isDeepModelingEnabled: params.isDeepModelingEnabled,
       })
     },
-    [replay, runStreamAttempt, senderName],
+    [replay, runStreamAttempt],
   )
 
   return {
     messages,
+    setMessages,
     isStreaming,
     error,
     start,