diff --git a/src/apps/shared/src/assistantTurn.ts b/src/apps/shared/src/assistantTurn.ts index 8a3dc54e0..d1da75edd 100644 --- a/src/apps/shared/src/assistantTurn.ts +++ b/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, } } diff --git a/src/apps/web/src/__tests__/assistantTurnSegments.test.ts b/src/apps/web/src/__tests__/assistantTurnSegments.test.ts index e53017985..7813075d5 100644 --- a/src/apps/web/src/__tests__/assistantTurnSegments.test.ts +++ b/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() diff --git a/src/apps/web/src/assistantTurnSegments.ts b/src/apps/web/src/assistantTurnSegments.ts index d1fd0267b..9da86cdcb 100644 --- a/src/apps/web/src/assistantTurnSegments.ts +++ b/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 { diff --git a/src/apps/web/src/components/MessageList.tsx b/src/apps/web/src/components/MessageList.tsx index 5cad783d9..2c3528175 100644 --- a/src/apps/web/src/components/MessageList.tsx +++ b/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( 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 (
( 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 ( ( .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( const entryComplete = !isLive const promotedNodes = [( ( )] return ( - + {promotedNodes} {histWidgets.map((w) => ( ( ) } - if (workGroupSplit.workGroup != null) { + const renderFinalText = () => ( + + ) + 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 ( <> - - {workGroupSplit.workGroup.segments.map((seg, si) => - renderSegment(seg, si, workGroupSplit.workGroup!.segments, false) - )} - - {workGroupSplit.finalText != null && ( - + {workGroupSplit.workGroup != null ? ( + + {workGroupSplit.workGroup.segments.map((seg, si) => + renderSegment(seg, si, workGroupSplit.workGroup!.segments, false, si) + )} + + ) : ( + 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) + })} ) }