fix(loading): fix spans and draft chat state

Author: waleedlatif1

apps/sim/app/workspace/[workspaceId]/home/components/mothership-chat/mothership-chat.tsx

@@ -1,6 +1,6 @@
 'use client'
 
-import { useLayoutEffect, useMemo, useRef } from 'react'
+import { useLayoutEffect, useRef } from 'react'
 import { cn } from '@/lib/core/utils/cn'
 import { MessageActions } from '@/app/workspace/[workspaceId]/components'
 import { ChatMessageAttachments } from '@/app/workspace/[workspaceId]/home/components/chat-message-attachments'
@@ -111,17 +111,14 @@ export function MothershipChat({
   const { staged: stagedMessages, isStaging } = useProgressiveList(messages, stagingKey)
   const stagedMessageCount = stagedMessages.length
   const stagedOffset = messages.length - stagedMessages.length
-  const precedingUserContentByIndex = useMemo(() => {
-    const contentByIndex: Array<string | undefined> = []
-    let lastUserContent: string | undefined
-    for (const [index, message] of messages.entries()) {
-      contentByIndex[index] = lastUserContent
-      if (message.role === 'user') {
-        lastUserContent = message.content
-      }
+  const precedingUserContentByIndex: Array<string | undefined> = []
+  let lastUserContent: string | undefined
+  for (const [index, message] of messages.entries()) {
+    precedingUserContentByIndex[index] = lastUserContent
+    if (message.role === 'user') {
+      lastUserContent = message.content
     }
-    return contentByIndex
-  }, [messages])
+  }
   const initialScrollDoneRef = useRef(false)
   const userInputRef = useRef<UserInputHandle>(null)
 

apps/sim/app/workspace/[workspaceId]/home/components/user-input/user-input.tsx

@@ -153,12 +153,13 @@ export const UserInput = forwardRef<UserInputHandle, UserInputProps>(function Us
   const overlayRef = useRef<HTMLDivElement>(null)
   const plusMenuRef = useRef<PlusMenuHandle>(null)
 
-  const prevDefaultValueRef = useRef(defaultValue)
-  useEffect(() => {
-    if (defaultValue === prevDefaultValueRef.current) return
-    prevDefaultValueRef.current = defaultValue
-    if (defaultValue) setValue(defaultValue)
-  }, [defaultValue])
+  const [prevDefaultValue, setPrevDefaultValue] = useState(defaultValue)
+  if (defaultValue && defaultValue !== prevDefaultValue) {
+    setPrevDefaultValue(defaultValue)
+    setValue(defaultValue)
+  } else if (!defaultValue && prevDefaultValue) {
+    setPrevDefaultValue(defaultValue)
+  }
 
   const files = useFileAttachments({
     userId: userId || session?.user?.id,
@@ -206,6 +207,13 @@ export const UserInput = forwardRef<UserInputHandle, UserInputProps>(function Us
         }))
       )
     }
+    if (draft.text) {
+      const textarea = textareaRef.current
+      if (textarea) {
+        textarea.focus()
+        textarea.setSelectionRange(draft.text.length, draft.text.length)
+      }
+    }
   }, []) // eslint-disable-line react-hooks/exhaustive-deps -- intentional mount-only restore
 
   // Skip the initial save — restore's setState calls haven't propagated yet, so