fix(web): preserve tail segments

src/apps/shared/src/assistantTurn.ts

@@ -37,15 +37,22 @@ export type WorkGroup = {
   segments: AssistantTurnSegment[]
 }
 
+export type WorkGroupSplit = {
+  workGroup: WorkGroup | null
+  finalText: string | null
+  finalTextIndex: number
+  tailSegments: AssistantTurnSegment[]
+}
+
 /**
  * Split segments into work group (pre-final) and final text.
- * The last text segment is the final answer; everything before it goes into the work group.
+ * The last text segment is the final answer; trailing segments stay after it.
  * Returns null workGroup when there's nothing meaningful to collapse.
  */
 export function splitWorkGroup(
   segments: AssistantTurnSegment[],
   durationMs: number,
-): { workGroup: WorkGroup | null; finalText: string | null } {
+): WorkGroupSplit {
   // Find the index of the last text segment
   let lastTextIndex = -1
   for (let i = segments.length - 1; i >= 0; i--) {
@@ -57,22 +64,25 @@ export function splitWorkGroup(
 
   // No text segment at all
   if (lastTextIndex === -1) {
-    return { workGroup: null, finalText: null }
+    return { workGroup: null, finalText: null, finalTextIndex: -1, tailSegments: [] }
   }
 
   const finalSegment = segments[lastTextIndex]!
   const finalText = finalSegment.type === 'text' ? finalSegment.content : null
+  const preSegments = segments.slice(0, lastTextIndex)
+  const tailSegments = segments.slice(lastTextIndex + 1)
 
   // Need at least 2 segments before final to justify a work group.
   // A single pre-final segment (text or cop) stays inline.
-  if (lastTextIndex < 2) {
-    return { workGroup: null, finalText }
+  if (preSegments.length < 2) {
+    return { workGroup: null, finalText, finalTextIndex: lastTextIndex, tailSegments }
   }
 
-  const workGroupSegments = segments.slice(0, lastTextIndex)
   return {
-    workGroup: { durationMs, segments: workGroupSegments },
+    workGroup: { durationMs, segments: preSegments },
     finalText,
+    finalTextIndex: lastTextIndex,
+    tailSegments,
   }
 }
 

src/apps/web/src/__tests__/assistantTurnSegments.test.ts

@@ -6,7 +6,9 @@ import {
   createEmptyAssistantTurnFoldState,
   finalizeAssistantTurnFoldState,
   foldAssistantTurnEvent,
+  isAssistantTurnSegmentLive,
   requestAssistantTurnThinkingBreak,
+  splitWorkGroup,
 } from '../assistantTurnSegments'
 import {
   normalizeAgentEventData,
@@ -46,6 +48,74 @@ function th(content: string, seq: number, endedByEventSeq?: number) {
   return { kind: 'thinking' as const, content, seq, startedAtMs, endedAtMs }
 }
 
+describe('splitWorkGroup', () => {
+  it('keeps a trailing tool segment after final text', () => {
+    const tailToolSegment = {
+      type: 'cop' as const,
+      title: null,
+      items: [{
+        kind: 'call' as const,
+        call: { toolCallId: 'terminal_1', toolName: 'terminal_run', arguments: {} },
+        seq: 2,
+      }],
+    }
+
+    const split = splitWorkGroup([
+      { type: 'text', content: 'done' },
+      tailToolSegment,
+    ], 1200)
+
+    expect(split.finalText).toBe('done')
+    expect(split.finalTextIndex).toBe(0)
+    expect(split.workGroup).toBeNull()
+    expect(split.tailSegments).toEqual([tailToolSegment])
+  })
+
+  it('collapses pre-final work while preserving trailing segment order', () => {
+    const firstToolSegment = {
+      type: 'cop' as const,
+      title: null,
+      items: [{
+        kind: 'call' as const,
+        call: { toolCallId: 'terminal_1', toolName: 'terminal_run', arguments: {} },
+        seq: 2,
+      }],
+    }
+    const tailToolSegment = {
+      type: 'cop' as const,
+      title: null,
+      items: [{
+        kind: 'call' as const,
+        call: { toolCallId: 'terminal_2', toolName: 'terminal_run', arguments: {} },
+        seq: 4,
+      }],
+    }
+
+    const split = splitWorkGroup([
+      { type: 'text', content: 'before' },
+      firstToolSegment,
+      { type: 'text', content: 'done' },
+      tailToolSegment,
+    ], 1200)
+
+    expect(split.finalText).toBe('done')
+    expect(split.finalTextIndex).toBe(2)
+    expect(split.workGroup?.segments).toEqual([
+      { type: 'text', content: 'before' },
+      firstToolSegment,
+    ])
+    expect(split.tailSegments).toEqual([tailToolSegment])
+  })
+})
+
+describe('isAssistantTurnSegmentLive', () => {
+  it('marks only the original last segment live for the current run', () => {
+    expect(isAssistantTurnSegmentLive(true, 2, 3)).toBe(true)
+    expect(isAssistantTurnSegmentLive(true, 1, 3)).toBe(false)
+    expect(isAssistantTurnSegmentLive(false, 2, 3)).toBe(false)
+  })
+})
+
 describe('buildAssistantTurnFromAgentEvents', () => {
   beforeEach(() => {
     vi.useFakeTimers()

src/apps/web/src/assistantTurnSegments.ts

@@ -16,6 +16,7 @@ import {
   type CopBlockItem,
   type TurnToolCallRef,
   type WorkGroup,
+  type WorkGroupSplit,
 } from '../../shared/src/assistantTurn'
 import {
   agentEventDataRecord,
@@ -40,6 +41,15 @@ export {
   type CopBlockItem,
   type TurnToolCallRef,
   type WorkGroup,
+  type WorkGroupSplit,
+}
+
+export function isAssistantTurnSegmentLive(
+  currentRunMessageLive: boolean,
+  originalIndex: number,
+  segmentCount: number,
+): boolean {
+  return currentRunMessageLive && originalIndex === segmentCount - 1
 }
 
 function toAssistantTurnEventType(type: string): string {

src/apps/web/src/components/MessageList.tsx

@@ -20,7 +20,7 @@ import { apiBaseUrl } from '@arkloop/shared/api'
 import type { AgentMessage } from '../agent-ui'
 import { copTimelinePayloadForSegment, type CopTimelinePayload, type TodoWriteRef } from '../copSegmentTimeline'
 import { buildResolvedPool, EMPTY_POOL, buildFallbackSegments } from '../copSubSegment'
-import { assistantTurnPlainText, splitWorkGroup, type AssistantTurnSegment, type WorkGroup as WorkGroupType } from '../assistantTurnSegments'
+import { assistantTurnPlainText, isAssistantTurnSegmentLive, splitWorkGroup, type AssistantTurnSegment, type WorkGroup as WorkGroupType, type WorkGroupSplit } from '../assistantTurnSegments'
 import { WorkGroup } from './WorkGroup'
 import { resolveMessageSourcesForRender } from './chatSourceResolver'
 import { createThreadShare } from '../api'
@@ -346,7 +346,9 @@ export const MessageList = memo(forwardRef<MessageListHandle, MessageListProps>(
     const messageWebFetches = msg.role === 'assistant' ? msgMeta?.webFetches : undefined
     const msgThinking = msg.role === 'assistant' ? msgMeta?.thinking : undefined
     const durationMs = historicalTurn?.durationMs ?? 0
-    const workGroupSplit = hasAssistantTurn ? splitWorkGroup(historicalSegments, durationMs) : { workGroup: null as WorkGroupType | null, finalText: null as string | null }
+    const workGroupSplit: WorkGroupSplit = hasAssistantTurn
+      ? splitWorkGroup(historicalSegments, durationMs)
+      : { workGroup: null as WorkGroupType | null, finalText: null, finalTextIndex: -1, tailSegments: [] }
     const bubbleCallbacks = bubbleCallbacksByMessageId.get(msg.id)
     return (
       <div
@@ -383,11 +385,17 @@ export const MessageList = memo(forwardRef<MessageListHandle, MessageListProps>(
                 sources: resolvedSources ?? [],
               }
 
-              const renderSegment = (seg: AssistantTurnSegment, si: number, segments: AssistantTurnSegment[], isLive: boolean) => {
+              const renderSegment = (
+                seg: AssistantTurnSegment,
+                si: number,
+                segments: AssistantTurnSegment[],
+                isLive: boolean,
+                originalIndex = si,
+              ) => {
                 if (seg.type === 'text') {
                   return (
                     <MarkdownRenderer
-                      key={`${msg.id}-at-${si}`}
+                      key={`${msg.id}-at-${originalIndex}`}
                       content={seg.content}
                       webSources={resolvedSources}
                       artifacts={msgMeta?.artifacts}
@@ -409,8 +417,8 @@ export const MessageList = memo(forwardRef<MessageListHandle, MessageListProps>(
                   .flatMap((entry) => entry.type === 'cop'
                     ? copTimelinePayloadForSegment(entry, timelinePools).todoWrites ?? []
                     : [])
-                const payload = precomputed?.payloads.get(String(si)) ?? copTimelinePayloadForSegment(seg, timelinePools)
-                const histWidgets = precomputed?.histWidgetsMap.get(String(si)) ?? historicWidgetsForCop(seg, msgWidgetsRaw)
+                const payload = precomputed?.payloads.get(String(originalIndex)) ?? copTimelinePayloadForSegment(seg, timelinePools)
+                const histWidgets = precomputed?.histWidgetsMap.get(String(originalIndex)) ?? historicWidgetsForCop(seg, msgWidgetsRaw)
 
                 const timelineTitleOverride = displayTerminalStatus != null
                   ? currentRunCopHeaderOverride({
@@ -429,9 +437,9 @@ export const MessageList = memo(forwardRef<MessageListHandle, MessageListProps>(
                 const entryComplete = !isLive
                 const promotedNodes = [(
                   <CopSegmentBlocks
-                    key={`${msg.id}-timeline-${si}`}
+                    key={`${msg.id}-timeline-${originalIndex}`}
                     segment={seg}
-                    keyPrefix={`${msg.id}-timeline-${si}`}
+                    keyPrefix={`${msg.id}-timeline-${originalIndex}`}
                     {...timelinePools}
                     isComplete={entryComplete}
                     live={isLive}
@@ -448,7 +456,7 @@ export const MessageList = memo(forwardRef<MessageListHandle, MessageListProps>(
                 )]
 
                 return (
-                  <Fragment key={`${msg.id}-acw-${si}`}>
+                  <Fragment key={`${msg.id}-acw-${originalIndex}`}>
                     {promotedNodes}
                     {histWidgets.map((w) => (
                       <WidgetBlock
@@ -464,29 +472,47 @@ export const MessageList = memo(forwardRef<MessageListHandle, MessageListProps>(
                 )
               }
 
-              if (workGroupSplit.workGroup != null) {
+              const renderFinalText = () => (
+                <MarkdownRenderer
+                  key={`${msg.id}-final`}
+                  content={workGroupSplit.finalText ?? ''}
+                  webSources={resolvedSources}
+                  artifacts={msgMeta?.artifacts}
+                  accessToken={accessToken}
+                  runId={msg.streamId ?? undefined}
+                  workFolder={workFolder}
+                  onOpenDocument={openDocumentPanel}
+                  onOpenResource={openResourcePanel}
+                  typography={isWorkMode ? 'work' : 'default'}
+                  trimTrailingMargin={true}
+                />
+              )
+              const isLiveSegment = (originalIndex: number) =>
+                isAssistantTurnSegmentLive(currentRunMessageLive, originalIndex, historicalSegments.length)
+
+              if (workGroupSplit.workGroup != null || workGroupSplit.tailSegments.length > 0) {
+                const preSegments = workGroupSplit.finalTextIndex >= 0
+                  ? historicalSegments.slice(0, workGroupSplit.finalTextIndex)
+                  : []
+                const tailStartIndex = workGroupSplit.finalTextIndex + 1
                 return (
                   <>
-                    <WorkGroup durationMs={durationMs}>
-                      {workGroupSplit.workGroup.segments.map((seg, si) =>
-                        renderSegment(seg, si, workGroupSplit.workGroup!.segments, false)
-                      )}
-                    </WorkGroup>
-                    {workGroupSplit.finalText != null && (
-                      <MarkdownRenderer
-                        key={`${msg.id}-final`}
-                        content={workGroupSplit.finalText}
-                        webSources={resolvedSources}
-                        artifacts={msgMeta?.artifacts}
-                        accessToken={accessToken}
-                        runId={msg.streamId ?? undefined}
-                        workFolder={workFolder}
-                        onOpenDocument={openDocumentPanel}
-                        onOpenResource={openResourcePanel}
-                        typography={isWorkMode ? 'work' : 'default'}
-                        trimTrailingMargin={true}
-                      />
+                    {workGroupSplit.workGroup != null ? (
+                      <WorkGroup durationMs={durationMs}>
+                        {workGroupSplit.workGroup.segments.map((seg, si) =>
+                          renderSegment(seg, si, workGroupSplit.workGroup!.segments, false, si)
+                        )}
+                      </WorkGroup>
+                    ) : (
+                      preSegments.map((seg, si) =>
+                        renderSegment(seg, si, historicalSegments, isLiveSegment(si), si)
+                      )
                     )}
+                    {workGroupSplit.finalText != null && renderFinalText()}
+                    {workGroupSplit.tailSegments.map((seg, offset) => {
+                      const originalIndex = tailStartIndex + offset
+                      return renderSegment(seg, originalIndex, historicalSegments, isLiveSegment(originalIndex), originalIndex)
+                    })}
                   </>
                 )
               }