From 623a09fd7e469ecd8255f97a052eed35c5bc4dac Mon Sep 17 00:00:00 2001
From: Shantur Rathore <i@shantur.com>
Date: Sun, 19 Apr 2026 19:56:48 +0100
Subject: [PATCH] fix(ui): stabilize long reply hold during streaming

---
 .../ui/src/components/message-section.tsx     | 11 ++-
 .../ui/src/components/virtual-follow-list.tsx | 77 +++++++++----------
 packages/ui/src/stores/session-events.ts      |  3 +-
 3 files changed, 47 insertions(+), 44 deletions(-)

diff --git a/packages/ui/src/components/message-section.tsx b/packages/ui/src/components/message-section.tsx
index 3fbdd6c9..442e021a 100644
--- a/packages/ui/src/components/message-section.tsx
+++ b/packages/ui/src/components/message-section.tsx
@@ -638,18 +638,25 @@ export default function MessageSection(props: MessageSectionProps) {
   const autoPinHoldTargetKey = createMemo(() => {
     if (!holdLongAssistantRepliesEnabled()) return null
     const messageId = lastVisibleMessageId()
-    return isAssistantTextMessage(messageId) ? messageId : null
+    return isStreamingAssistantTextMessage(messageId) ? messageId : null
   })
 
   function toggleHoldLongAssistantReplies() {
     updatePreferences({ holdLongAssistantReplies: !holdLongAssistantRepliesEnabled() })
   }
 
-  function isAssistantTextMessage(messageId: string | null | undefined) {
+  function isStreamingAssistantTextMessage(messageId: string | null | undefined) {
     if (!messageId) return false
     const resolvedStore = store()
     const record = resolvedStore.getMessage(messageId)
     if (!record || record.role !== "assistant") return false
+    if (record.status !== "streaming") return false
+
+    const info = resolvedStore.getMessageInfo(messageId)
+    if (!info) return false
+    const timeInfo = info?.time as { end?: number } | undefined
+    const isStreaming = timeInfo?.end === undefined || timeInfo.end === 0
+    if (!isStreaming) return false
 
     const { orderedParts } = buildRecordDisplayData(props.instanceId, record)
     return orderedParts.some((part) => {
diff --git a/packages/ui/src/components/virtual-follow-list.tsx b/packages/ui/src/components/virtual-follow-list.tsx
index 0c0bbe35..070e0b92 100644
--- a/packages/ui/src/components/virtual-follow-list.tsx
+++ b/packages/ui/src/components/virtual-follow-list.tsx
@@ -161,8 +161,9 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
   const [showScrollTopButton, setShowScrollTopButton] = createSignal(false)
   const [showScrollBottomButton, setShowScrollBottomButton] = createSignal(false)
   const [activeKey, setActiveKey] = createSignal<string | null>(null)
-  const [heldItemCount, setHeldItemCount] = createSignal<number | null>(null)
-  const effectiveSuspendAutoPinToBottom = () => externalSuspendAutoPinToBottom() || heldItemCount() !== null
+  const [activeHoldTargetKey, setActiveHoldTargetKey] = createSignal<string | null>(null)
+  const [didTriggerHoldForCurrentTarget, setDidTriggerHoldForCurrentTarget] = createSignal(false)
+  const effectiveSuspendAutoPinToBottom = () => externalSuspendAutoPinToBottom() || activeHoldTargetKey() !== null
 
   const scrollButtonsCount = createMemo(() => (showScrollTopButton() ? 1 : 0) + (showScrollBottomButton() ? 1 : 0))
   const itemElements = new Map<string, HTMLDivElement>()
@@ -196,6 +197,17 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
     return performance.now() <= userScrollIntentUntil
   }
 
+  function clearAutoPinHold(options?: { resumeBottom?: boolean }) {
+    if (activeHoldTargetKey() === null) return
+    setActiveHoldTargetKey(null)
+    if (options?.resumeBottom && autoScroll()) {
+      requestAnimationFrame(() => {
+        if (!autoScroll() || activeHoldTargetKey() !== null) return
+        scrollToBottom(false)
+      })
+    }
+  }
+
   function attachScrollIntentListeners(element: HTMLDivElement | undefined) {
     if (detachScrollIntentListeners) {
       detachScrollIntentListeners()
@@ -257,7 +269,7 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
     // scrollbar). If follow mode stays enabled, the next render notification
     // snaps the list straight back to bottom. A real upward viewport move away
     // from bottom should always break follow unless a hold target is active.
-    if (wasPinnedAtBottom && scrolledUp && autoScroll() && !atBottom && heldItemCount() === null) {
+    if (wasPinnedAtBottom && scrolledUp && autoScroll() && !atBottom && activeHoldTargetKey() === null) {
       setAutoScroll(false)
       lastObservedPinnedAtBottom = false
       return
@@ -265,9 +277,7 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
 
     // Sync autoScroll state based on scroll position if it was a user scroll
     if (hasUserScrollIntent()) {
-      if (atBottom && heldItemCount() !== null) {
-        setHeldItemCount(null)
-      }
+      clearAutoPinHold()
       if (atBottom && !autoScroll()) {
         setAutoScroll(true)
       } else if (!atBottom && autoScroll()) {
@@ -303,7 +313,6 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
       }
     }
     updateScrollButtons()
-    updateAutoPinHold()
     props.onScroll?.()
 
     // Find active key (roughly the first visible item)
@@ -335,25 +344,14 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
 
   function updateAutoPinHold() {
     const element = scrollElement()
-    const itemCount = props.items().length
-    const heldCount = heldItemCount()
     if (!element) return
 
-    if (heldCount !== null) {
-      if (itemCount > heldCount) {
-        setHeldItemCount(null)
-        if (autoScroll()) {
-          requestAnimationFrame(() => {
-            if (!autoScroll()) return
-            scrollToBottom(false)
-          })
-        }
-        return
-      }
+    const targetKey = holdTargetKey()
+    const heldKey = activeHoldTargetKey()
 
-      if (itemCount < heldCount) {
-        setHeldItemCount(null)
-        return
+    if (heldKey !== null) {
+      if (targetKey !== heldKey) {
+        clearAutoPinHold({ resumeBottom: true })
       }
 
       return
@@ -361,9 +359,8 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
 
     if (!autoScroll()) return
     if (externalSuspendAutoPinToBottom()) return
-
-    const targetKey = holdTargetKey()
     if (!targetKey) return
+    if (didTriggerHoldForCurrentTarget()) return
 
     const itemWrapper = itemElements.get(targetKey)
     if (!itemWrapper) return
@@ -379,7 +376,8 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
       if (Math.abs(alignDelta) > 1) {
         element.scrollTop = Math.max(0, element.scrollTop + alignDelta)
       }
-      setHeldItemCount(itemCount)
+      setActiveHoldTargetKey(targetKey)
+      setDidTriggerHoldForCurrentTarget(true)
     }
   }
 
@@ -395,7 +393,7 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
     },
     notifyContentRendered: () => {
       updateAutoPinHold()
-      if (heldItemCount() !== null) return
+      if (activeHoldTargetKey() !== null) return
       if (autoScroll() && !effectiveSuspendAutoPinToBottom()) {
         scrollToBottom(true)
       }
@@ -411,14 +409,23 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
 
   createEffect(on(() => props.resetKey?.(), () => {
     itemElements.clear()
-    setHeldItemCount(null)
+    setActiveHoldTargetKey(null)
+    setDidTriggerHoldForCurrentTarget(false)
     lastObservedScrollOffset = 0
     lastObservedPinnedAtBottom = false
   }))
 
+  createEffect(on(holdTargetKey, (nextTargetKey, prevTargetKey) => {
+    if (nextTargetKey !== prevTargetKey && didTriggerHoldForCurrentTarget()) {
+      setDidTriggerHoldForCurrentTarget(false)
+    }
+    if (activeHoldTargetKey() === null) return
+    if (nextTargetKey === activeHoldTargetKey()) return
+    clearAutoPinHold({ resumeBottom: true })
+  }, { defer: true }))
+
   // Handle autoScroll (Follow) on items change
   createEffect(on(() => props.items().length, (len, prevLen) => {
-    updateAutoPinHold()
     if (len > (prevLen ?? 0) && autoScroll() && !effectiveSuspendAutoPinToBottom() && !suppressAutoScrollOnce) {
       requestAnimationFrame(() => scrollToBottom(true))
     }
@@ -427,16 +434,11 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
 
   // Handle followToken change
   createEffect(on(() => props.followToken?.(), () => {
-    updateAutoPinHold()
     if (autoScroll() && !effectiveSuspendAutoPinToBottom()) {
       scrollToBottom(true)
     }
   }, { defer: true }))
 
-  createEffect(on(() => holdTargetKey(), () => {
-    updateAutoPinHold()
-  }, { defer: true }))
-
   // Reset state on resetKey change
   createEffect(on(() => props.resetKey?.(), (nextKey) => {
     if (nextKey === lastResetKey) return
@@ -459,13 +461,6 @@ export default function VirtualFollowList<T>(props: VirtualFollowListProps<T>) {
     }
   })
 
-  createEffect(() => {
-    if (typeof window === "undefined") return
-    const handleResize = () => updateAutoPinHold()
-    window.addEventListener("resize", handleResize)
-    onCleanup(() => window.removeEventListener("resize", handleResize))
-  })
-
   return (
     <div class="virtual-follow-list-shell" ref={shellElement => {
       setShellElement(shellElement)
diff --git a/packages/ui/src/stores/session-events.ts b/packages/ui/src/stores/session-events.ts
index 713c95d1..1068636c 100644
--- a/packages/ui/src/stores/session-events.ts
+++ b/packages/ui/src/stores/session-events.ts
@@ -397,7 +397,8 @@ function handleMessageUpdate(instanceId: string, event: MessageUpdateEvent | Mes
 
     const role: MessageRole = info.role === "user" ? "user" : "assistant"
     const hasError = Boolean((info as any).error)
-    const status: MessageStatus = hasError ? "error" : "complete"
+    const hasEnded = typeof timeInfo.end === "number" && timeInfo.end > 0
+    const status: MessageStatus = hasError ? "error" : hasEnded ? "complete" : "streaming"
 
     let record = store.getMessage(messageId)
     if (!record) {