diff --git a/ui-tui/src/__tests__/createGatewayEventHandler.test.ts b/ui-tui/src/__tests__/createGatewayEventHandler.test.ts index 658ca571b..c3cb5095d 100644 --- a/ui-tui/src/__tests__/createGatewayEventHandler.test.ts +++ b/ui-tui/src/__tests__/createGatewayEventHandler.test.ts @@ -82,15 +82,13 @@ describe('createGatewayEventHandler', () => { type: 'message.complete' } as any) - expect(appended).toHaveLength(1) - expect(appended[0]).toMatchObject({ - role: 'assistant', - text: 'final answer', - thinking: 'mapped the page' - }) - expect(appended[0]?.tools).toHaveLength(1) - expect(appended[0]?.tools?.[0]).toContain('hero cards') - expect(appended[0]?.toolTokens).toBeGreaterThan(0) + expect(appended).toHaveLength(3) + expect(appended[0]).toMatchObject({ kind: 'trail', role: 'system', text: '', thinking: 'mapped the page' }) + expect(appended[1]).toMatchObject({ kind: 'trail', role: 'system', text: '' }) + expect(appended[1]?.tools).toHaveLength(1) + expect(appended[1]?.tools?.[0]).toContain('hero cards') + expect(appended[1]?.toolTokens).toBeGreaterThan(0) + expect(appended[2]).toMatchObject({ role: 'assistant', text: 'final answer' }) }) it('keeps tool tokens across handler recreation mid-turn', () => { @@ -118,9 +116,10 @@ describe('createGatewayEventHandler', () => { type: 'message.complete' } as any) - expect(appended).toHaveLength(1) - expect(appended[0]?.tools).toHaveLength(1) - expect(appended[0]?.toolTokens).toBeGreaterThan(0) + expect(appended).toHaveLength(3) + expect(appended[1]?.tools).toHaveLength(1) + expect(appended[1]?.toolTokens).toBeGreaterThan(0) + expect(appended[2]).toMatchObject({ role: 'assistant', text: 'final answer' }) }) it('streams legacy thinking.delta into visible reasoning state', () => { @@ -148,9 +147,10 @@ describe('createGatewayEventHandler', () => { onEvent({ payload: { text: fallback }, type: 'reasoning.available' } as any) onEvent({ payload: { text: 'final answer' }, type: 'message.complete' } as any) - expect(appended).toHaveLength(1) + expect(appended).toHaveLength(2) expect(appended[0]?.thinking).toBe(streamed) expect(appended[0]?.thinkingTokens).toBe(estimateTokensRough(streamed)) + expect(appended[1]).toMatchObject({ role: 'assistant', text: 'final answer' }) }) it('uses message.complete reasoning when no streamed reasoning ref', () => { @@ -161,9 +161,10 @@ describe('createGatewayEventHandler', () => { onEvent({ payload: { reasoning: fromServer, text: 'final answer' }, type: 'message.complete' } as any) - expect(appended).toHaveLength(1) + expect(appended).toHaveLength(2) expect(appended[0]?.thinking).toBe(fromServer) expect(appended[0]?.thinkingTokens).toBe(estimateTokensRough(fromServer)) + expect(appended[1]).toMatchObject({ role: 'assistant', text: 'final answer' }) }) it('anchors inline_diff as its own segment where the edit happened', () => { @@ -184,21 +185,19 @@ describe('createGatewayEventHandler', () => { expect(appended).toHaveLength(0) expect(turnController.segmentMessages).toEqual([ { role: 'assistant', text: 'Editing the file' }, + { kind: 'trail', role: 'system', text: '', tools: ['Patch("foo.ts") ✓'] }, { kind: 'diff', role: 'assistant', text: block } ]) onEvent({ payload: { text: 'patch applied' }, type: 'message.complete' } as any) - // Four transcript messages: pre-tool narration → tool trail → diff - // (kind='diff', so MessageLine gives it blank-line breathing room) → - // post-tool narration. The final message does NOT contain a diff. - expect(appended).toHaveLength(4) + expect(appended).toHaveLength(5) expect(appended[0]?.text).toBe('Editing the file') expect(appended[1]).toMatchObject({ kind: 'trail' }) expect(appended[1]?.tools?.[0]).toContain('Patch') expect(appended[2]).toMatchObject({ kind: 'diff', text: block }) - expect(appended[3]?.text).toBe('patch applied') - expect(appended[3]?.text).not.toContain('```diff') + expect(appended[4]?.text).toBe('patch applied') + expect(appended[4]?.text).not.toContain('```diff') }) it('drops the diff segment when the final assistant text narrates the same diff', () => { @@ -212,9 +211,10 @@ describe('createGatewayEventHandler', () => { // Only the final message — diff-only segment dropped so we don't // render two stacked copies of the same patch. - expect(appended).toHaveLength(1) - expect(appended[0]?.text).toBe(assistantText) - expect((appended[0]?.text.match(/```diff/g) ?? []).length).toBe(1) + expect(appended).toHaveLength(2) + expect(appended[0]).toMatchObject({ kind: 'trail' }) + expect(appended[1]?.text).toBe(assistantText) + expect((appended[1]?.text.match(/```diff/g) ?? []).length).toBe(1) }) it('strips the CLI "┊ review diff" header from inline diff segments', () => { @@ -246,9 +246,10 @@ describe('createGatewayEventHandler', () => { } as any) onEvent({ payload: { text: assistantText }, type: 'message.complete' } as any) - expect(appended).toHaveLength(1) - expect(appended[0]?.text).toBe(assistantText) - expect((appended[0]?.text.match(/```diff/g) ?? []).length).toBe(1) + expect(appended).toHaveLength(2) + expect(appended[0]).toMatchObject({ kind: 'trail' }) + expect(appended[1]?.text).toBe(assistantText) + expect((appended[1]?.text.match(/```diff/g) ?? []).length).toBe(1) }) it('keeps tool trail terse when inline_diff is present', () => { diff --git a/ui-tui/src/app/createGatewayEventHandler.ts b/ui-tui/src/app/createGatewayEventHandler.ts index 94e82c56c..502f5387c 100644 --- a/ui-tui/src/app/createGatewayEventHandler.ts +++ b/ui-tui/src/app/createGatewayEventHandler.ts @@ -379,6 +379,10 @@ export function createGatewayEventHandler(ctx: GatewayEventHandlerContext): (ev: const inlineDiffText = ev.payload.inline_diff && getUiState().inlineDiffs ? stripAnsi(String(ev.payload.inline_diff)).trim() : '' + if (inlineDiffText) { + turnController.flushStreamingSegment() + } + turnController.recordToolComplete( ev.payload.tool_id, ev.payload.name, @@ -386,17 +390,10 @@ export function createGatewayEventHandler(ctx: GatewayEventHandlerContext): (ev: inlineDiffText ? '' : ev.payload.summary ) - if (!inlineDiffText) { - return + if (inlineDiffText) { + turnController.pushInlineDiffSegment(inlineDiffText) } - // Anchor the diff to where the edit happened in the turn — between - // the narration that preceded the tool call and whatever the agent - // streams afterwards. The previous end-merge put the diff at the - // bottom of the final message even when the edit fired mid-turn, - // which read as "the agent wrote this after saying that". - turnController.pushInlineDiffSegment(inlineDiffText) - return } diff --git a/ui-tui/src/app/turnController.ts b/ui-tui/src/app/turnController.ts index 3240c4e89..1269409dd 100644 --- a/ui-tui/src/app/turnController.ts +++ b/ui-tui/src/app/turnController.ts @@ -38,11 +38,7 @@ const diffSegmentBody = (msg: Msg): null | string => { return m ? m[1]! : null } -const insertBeforeFirstDiff = (segments: Msg[], msg: Msg): Msg[] => { - const index = segments.findIndex(segment => segment.kind === 'diff') - - return index < 0 ? [...segments, msg] : [...segments.slice(0, index), msg, ...segments.slice(index)] -} +const hasDetails = (msg: Msg): boolean => Boolean(msg.thinking || msg.tools?.length || msg.toolTokens) export interface InterruptDeps { appendMessage: (msg: Msg) => void @@ -69,6 +65,7 @@ class TurnController { persistSpawnTree?: (subagents: SubagentProgress[], sessionId: null | string) => Promise protocolWarned = false reasoningText = '' + reasoningSegmentOffset = 0 segmentMessages: Msg[] = [] pendingSegmentTools: string[] = [] statusTimer: Timer = null @@ -94,6 +91,7 @@ class TurnController { clearReasoning() { this.reasoningTimer = clear(this.reasoningTimer) this.reasoningText = '' + this.reasoningSegmentOffset = 0 this.toolTokenAcc = 0 patchTurnState({ reasoning: '', reasoningTokens: 0, toolTokens: 0 }) } @@ -181,29 +179,33 @@ class TurnController { flushStreamingSegment() { const raw = this.bufRef.trimStart() - - if (!raw) { - return - } - - const split = hasReasoningTag(raw) ? splitReasoning(raw) : { reasoning: '', text: raw } + const split = raw ? (hasReasoningTag(raw) ? splitReasoning(raw) : { reasoning: '', text: raw }) : { reasoning: '', text: '' } if (split.reasoning && !this.reasoningText.trim()) { this.reasoningText = split.reasoning patchTurnState({ reasoning: this.reasoningText, reasoningTokens: estimateTokensRough(this.reasoningText) }) } - const text = split.text + const thinking = this.reasoningText.slice(this.reasoningSegmentOffset).trim() + const msg: Msg = { + role: split.text ? 'assistant' : 'system', + text: split.text, + ...(!split.text && { kind: 'trail' as const }), + ...(thinking && { + thinking, + thinkingTokens: estimateTokensRough(thinking) + }), + ...(this.pendingSegmentTools.length && { tools: this.pendingSegmentTools }) + } this.streamTimer = clear(this.streamTimer) - if (text) { - const tools = this.pendingSegmentTools - - this.segmentMessages = [...this.segmentMessages, { role: 'assistant', text, ...(tools.length && { tools }) }] - this.pendingSegmentTools = [] + if (split.text || hasDetails(msg)) { + this.segmentMessages = [...this.segmentMessages, msg] } + this.reasoningSegmentOffset = this.reasoningText.length + this.pendingSegmentTools = [] this.bufRef = '' patchTurnState({ streamPendingTools: [], streamSegments: this.segmentMessages, streaming: '' }) } @@ -295,7 +297,6 @@ class TurnController { const finalText = split.text const existingReasoning = this.reasoningText.trim() || String(payload.reasoning ?? '').trim() const savedReasoning = [existingReasoning, existingReasoning ? '' : split.reasoning].filter(Boolean).join('\n\n') - const savedReasoningTokens = savedReasoning ? estimateTokensRough(savedReasoning) : 0 const savedToolTokens = this.toolTokenAcc const tools = this.pendingSegmentTools @@ -312,32 +313,20 @@ class TurnController { return body === null || (!finalHasOwnDiffFence && !finalText.includes(body)) }) - const hasDiffSegment = segments.some(msg => msg.kind === 'diff') - const detailsBelongBeforeDiff = hasDiffSegment && (tools.length > 0 || Boolean(savedReasoning)) - - const finalMessages = detailsBelongBeforeDiff - ? insertBeforeFirstDiff(segments, { - kind: 'trail', - role: 'system', - text: '', - thinking: savedReasoning || undefined, - thinkingTokens: savedReasoning ? savedReasoningTokens : undefined, - toolTokens: savedToolTokens || undefined, - ...(tools.length && { tools }) - }) - : [...segments] + const finalThinking = savedReasoning.slice(this.reasoningSegmentOffset).trim() + const finalDetails: Msg = { + kind: 'trail', + role: 'system', + text: '', + thinking: finalThinking || undefined, + thinkingTokens: finalThinking ? estimateTokensRough(finalThinking) : undefined, + toolTokens: savedToolTokens || undefined, + ...(tools.length && { tools }) + } + const finalMessages = hasDetails(finalDetails) ? [...segments, finalDetails] : [...segments] if (finalText) { - finalMessages.push({ - role: 'assistant', - text: finalText, - ...(!detailsBelongBeforeDiff && { - thinking: savedReasoning || undefined, - thinkingTokens: savedReasoning ? savedReasoningTokens : undefined, - toolTokens: savedToolTokens || undefined, - ...(tools.length && { tools }) - }) - }) + finalMessages.push({ role: 'assistant', text: finalText }) } const wasInterrupted = this.interrupted