fix: Fix optimistic update not working on reload for first HumanMessage

Move optimistic update logic from useStream to SessionDetailPageClient to ensure
consistent behavior. Previously, the isFirstMessage flag was reset on reload,
preventing optimistic updates for the first message after page reload.

Changes:
- Remove isFirstMessage ref from useStream.ts
- Remove optimistic update logic from start function
- Expose setMessages from useStream hook
- Remove senderName param from useStream (moved to parent component)
- Implement handleSendMessage in SessionDetailPageClient with optimistic update
- Rollback optimistic message on error

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

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,