Merge pull request #14822 from NousResearch/bb/tui-inline-diff-segment-anchor

fix(tui): anchor inline_diff to the segment where the edit happened

ui-tui/src/__tests__/createGatewayEventHandler.test.ts

@@ -152,91 +152,79 @@ describe('createGatewayEventHandler', () => {
     expect(appended[0]?.thinkingTokens).toBe(estimateTokensRough(fromServer))
   })
 
-  it('attaches inline_diff to the assistant completion body', () => {
+  it('anchors inline_diff as its own segment where the edit happened', () => {
     const appended: Msg[] = []
     const onEvent = createGatewayEventHandler(buildCtx(appended))
     const diff = '\u001b[31m--- a/foo.ts\u001b[0m\n\u001b[32m+++ b/foo.ts\u001b[0m\n@@\n-old\n+new'
     const cleaned = '--- a/foo.ts\n+++ b/foo.ts\n@@\n-old\n+new'
+    const block = `\`\`\`diff\n${cleaned}\n\`\`\``
 
-    onEvent({
+    // Narration → tool → tool-complete → more narration → message-complete.
-      payload: { context: 'foo.ts', name: 'patch', tool_id: 'tool-1' },
+    // The diff MUST land between the two narration segments, not tacked
-      type: 'tool.start'
+    // onto the final one.
-    } as any)
+    onEvent({ payload: { text: 'Editing the file' }, type: 'message.delta' } as any)
-    onEvent({
+    onEvent({ payload: { context: 'foo.ts', name: 'patch', tool_id: 'tool-1' }, type: 'tool.start' } as any)
-      payload: { inline_diff: diff, summary: 'patched', tool_id: 'tool-1' },
+    onEvent({ payload: { inline_diff: diff, summary: 'patched', tool_id: 'tool-1' }, type: 'tool.complete' } as any)
-      type: 'tool.complete'
-    } as any)
 
-    // Diff is buffered for message.complete and sanitized (ANSI stripped).
+    // Diff is already committed to segmentMessages as its own segment.
     expect(appended).toHaveLength(0)
-    expect(turnController.pendingInlineDiffs).toEqual([cleaned])
+    expect(turnController.segmentMessages).toEqual([
+      { role: 'assistant', text: 'Editing the file' },
+      { kind: 'diff', role: 'assistant', text: block }
+    ])
 
-    onEvent({
+    onEvent({ payload: { text: 'patch applied' }, type: 'message.complete' } as any)
-      payload: { text: 'patch applied' },
-      type: 'message.complete'
-    } as any)
 
-    // Diff is rendered in the same assistant message body as the completion.
+    // Three transcript messages: pre-tool narration → diff (kind='diff',
-    expect(appended).toHaveLength(1)
+    // so MessageLine gives it blank-line breathing room) → post-tool
-    expect(appended[0]).toMatchObject({ role: 'assistant' })
+    // narration. The final message does NOT contain a diff.
-    expect(appended[0]?.text).toContain('patch applied')
+    expect(appended).toHaveLength(3)
-    expect(appended[0]?.text).toContain('```diff')
+    expect(appended[0]?.text).toBe('Editing the file')
-    expect(appended[0]?.text).toContain(cleaned)
+    expect(appended[1]).toMatchObject({ kind: 'diff', text: block })
+    expect(appended[2]?.text).toBe('patch applied')
+    expect(appended[2]?.text).not.toContain('```diff')
   })
 
-  it('does not append inline_diff twice when assistant text already contains it', () => {
+  it('drops the diff segment when the final assistant text narrates the same diff', () => {
     const appended: Msg[] = []
     const onEvent = createGatewayEventHandler(buildCtx(appended))
     const cleaned = '--- a/foo.ts\n+++ b/foo.ts\n@@\n-old\n+new'
     const assistantText = `Done. Here's the inline diff:\n\n\`\`\`diff\n${cleaned}\n\`\`\``
 
-    onEvent({
+    onEvent({ payload: { inline_diff: cleaned, summary: 'patched', tool_id: 'tool-1' }, type: 'tool.complete' } as any)
-      payload: { inline_diff: cleaned, summary: 'patched', tool_id: 'tool-1' },
+    onEvent({ payload: { text: assistantText }, type: 'message.complete' } as any)
-      type: 'tool.complete'
-    } as any)
-    onEvent({
-      payload: { text: assistantText },
-      type: 'message.complete'
-    } as any)
 
+    // 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)
   })
 
-  it('strips the CLI "┊ review diff" header from queued inline diffs', () => {
+  it('strips the CLI "┊ review diff" header from inline diff segments', () => {
     const appended: Msg[] = []
     const onEvent = createGatewayEventHandler(buildCtx(appended))
     const raw = '  \u001b[33m┊ review diff\u001b[0m\n--- a/foo.ts\n+++ b/foo.ts\n@@\n-old\n+new'
 
-    onEvent({
+    onEvent({ payload: { inline_diff: raw, summary: 'patched', tool_id: 'tool-1' }, type: 'tool.complete' } as any)
-      payload: { inline_diff: raw, summary: 'patched', tool_id: 'tool-1' },
+    onEvent({ payload: { text: 'done' }, type: 'message.complete' } as any)
-      type: 'tool.complete'
-    } as any)
-    onEvent({
-      payload: { text: 'done' },
-      type: 'message.complete'
-    } as any)
 
-    expect(appended).toHaveLength(1)
+    // diff segment first (kind='diff'), final narration second
+    expect(appended).toHaveLength(2)
+    expect(appended[0]?.kind).toBe('diff')
     expect(appended[0]?.text).not.toContain('┊ review diff')
     expect(appended[0]?.text).toContain('--- a/foo.ts')
+    expect(appended[1]?.text).toBe('done')
   })
 
-  it('suppresses inline_diff when assistant already wrote a diff fence', () => {
+  it('drops the diff segment when assistant writes its own ```diff fence', () => {
     const appended: Msg[] = []
     const onEvent = createGatewayEventHandler(buildCtx(appended))
     const inlineDiff = '--- a/foo.ts\n+++ b/foo.ts\n@@\n-old\n+new'
     const assistantText = 'Done. Clean swap:\n\n```diff\n-old\n+new\n```'
 
-    onEvent({
+    onEvent({ payload: { inline_diff: inlineDiff, summary: 'patched', tool_id: 'tool-1' }, type: 'tool.complete' } as any)
-      payload: { inline_diff: inlineDiff, summary: 'patched', tool_id: 'tool-1' },
+    onEvent({ payload: { text: assistantText }, type: 'message.complete' } as any)
-      type: 'tool.complete'
-    } as any)
-    onEvent({
-      payload: { text: assistantText },
-      type: 'message.complete'
-    } as any)
 
     expect(appended).toHaveLength(1)
     expect(appended[0]?.text).toBe(assistantText)
@@ -252,15 +240,18 @@ describe('createGatewayEventHandler', () => {
       payload: { inline_diff: diff, name: 'review_diff', summary: diff, tool_id: 'tool-1' },
       type: 'tool.complete'
     } as any)
-    onEvent({
+    onEvent({ payload: { text: 'done' }, type: 'message.complete' } as any)
-      payload: { text: 'done' },
-      type: 'message.complete'
-    } as any)
 
-    expect(appended).toHaveLength(1)
+    // Two segments: the diff block (kind='diff', no tool row) and the final
-    expect(appended[0]?.tools?.[0]).toContain('Review Diff')
+    // narration (tool row belongs here since pendingSegmentTools carries
-    expect(appended[0]?.tools?.[0]).not.toContain('--- a/foo.ts')
+    // across the flushStreamingSegment call).
+    expect(appended).toHaveLength(2)
+    expect(appended[0]?.kind).toBe('diff')
     expect(appended[0]?.text).toContain('```diff')
+    expect(appended[0]?.tools ?? []).toEqual([])
+    expect(appended[1]?.text).toBe('done')
+    expect(appended[1]?.tools?.[0]).toContain('Review Diff')
+    expect(appended[1]?.tools?.[0]).not.toContain('--- a/foo.ts')
   })
 
   it('shows setup panel for missing provider startup error', () => {

ui-tui/src/app/createGatewayEventHandler.ts

@@ -385,10 +385,12 @@ export function createGatewayEventHandler(ctx: GatewayEventHandlerContext): (ev:
           return
         }
 
-        // Keep inline diffs attached to the assistant completion body so
+        // Anchor the diff to where the edit happened in the turn — between
-        // they render in the same message flow, not as a standalone system
+        // the narration that preceded the tool call and whatever the agent
-        // artifact that can look out-of-place around tool rows.
+        // streams afterwards. The previous end-merge put the diff at the
-        turnController.queueInlineDiff(inlineDiffText)
+        // 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
       }

ui-tui/src/app/turnController.ts

@@ -19,6 +19,20 @@ const INTERRUPT_COOLDOWN_MS = 1500
 const ACTIVITY_LIMIT = 8
 const TRAIL_LIMIT = 8
 
+// Extracts the raw patch from a diff-only segment produced by
+// pushInlineDiffSegment. Used at message.complete to dedupe against final
+// assistant text that narrates the same patch. Returns null for anything
+// else so real assistant narration never gets touched.
+const diffSegmentBody = (msg: Msg): null | string => {
+  if (msg.kind !== 'diff') {
+    return null
+  }
+
+  const m = msg.text.match(/^```diff\n([\s\S]*?)\n```$/)
+
+  return m ? m[1]! : null
+}
+
 export interface InterruptDeps {
   appendMessage: (msg: Msg) => void
   gw: { request: <T = unknown>(method: string, params?: Record<string, unknown>) => Promise<T> }
@@ -40,7 +54,6 @@ class TurnController {
   bufRef = ''
   interrupted = false
   lastStatusNote = ''
-  pendingInlineDiffs: string[] = []
   persistedToolLabels = new Set<string>()
   persistSpawnTree?: (subagents: SubagentProgress[], sessionId: null | string) => Promise<void>
   protocolWarned = false
@@ -79,7 +92,6 @@ class TurnController {
     this.activeTools = []
     this.streamTimer = clear(this.streamTimer)
     this.bufRef = ''
-    this.pendingInlineDiffs = []
     this.pendingSegmentTools = []
     this.segmentMessages = []
 
@@ -186,18 +198,35 @@ class TurnController {
     }, REASONING_PULSE_MS)
   }
 
-  queueInlineDiff(diffText: string) {
+  pushInlineDiffSegment(diffText: string) {
     // Strip CLI chrome the gateway emits before the unified diff (e.g. a
     // leading "┊ review diff" header written by `_emit_inline_diff` for the
     // terminal printer). That header only makes sense as stdout dressing,
     // not inside a markdown ```diff block.
-    const text = diffText.replace(/^\s*┊[^\n]*\n?/, '').trim()
+    const stripped = diffText.replace(/^\s*┊[^\n]*\n?/, '').trim()
 
-    if (!text || this.pendingInlineDiffs.includes(text)) {
+    if (!stripped) {
       return
     }
 
-    this.pendingInlineDiffs = [...this.pendingInlineDiffs, text]
+    // Flush any in-progress streaming text as its own segment first, so the
+    // diff lands BETWEEN the assistant narration that preceded the edit and
+    // whatever the agent streams afterwards — not glued onto the final
+    // message. This is the whole point of segment-anchored diffs: the diff
+    // renders where the edit actually happened.
+    this.flushStreamingSegment()
+
+    const block = `\`\`\`diff\n${stripped}\n\`\`\``
+
+    // Skip consecutive duplicates (same tool firing tool.complete twice, or
+    // two edits producing the same patch). Keeping this cheap — deeper
+    // dedupe against the final assistant text happens at message.complete.
+    if (this.segmentMessages.at(-1)?.text === block) {
+      return
+    }
+
+    this.segmentMessages = [...this.segmentMessages, { kind: 'diff', role: 'assistant', text: block }]
+    patchTurnState({ streamSegments: this.segmentMessages })
   }
 
   pushActivity(text: string, tone: ActivityItem['tone'] = 'info', replaceLabel?: string) {
@@ -234,7 +263,6 @@ class TurnController {
     this.idle()
     this.clearReasoning()
     this.clearStatusTimer()
-    this.pendingInlineDiffs = []
     this.pendingSegmentTools = []
     this.segmentMessages = []
     this.turnTools = []
@@ -245,31 +273,31 @@ class TurnController {
     const rawText = (payload.rendered ?? payload.text ?? this.bufRef).trimStart()
     const split = splitReasoning(rawText)
     const finalText = split.text
-    // Skip appending if the assistant already narrated the diff inside a
-    // markdown fence of its own — otherwise we render two stacked diff
-    // blocks for the same edit.
-    const assistantAlreadyHasDiff = /```(?:diff|patch)\b/i.test(finalText)
-
-    const remainingInlineDiffs = assistantAlreadyHasDiff
-      ? []
-      : this.pendingInlineDiffs.filter(diff => !finalText.includes(diff))
-
-    const inlineDiffBlock = remainingInlineDiffs.length
-      ? `\`\`\`diff\n${remainingInlineDiffs.join('\n\n')}\n\`\`\``
-      : ''
-
-    const mergedText = [finalText, inlineDiffBlock].filter(Boolean).join('\n\n')
     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
-    const finalMessages = [...this.segmentMessages]
 
-    if (mergedText) {
+    // Drop diff-only segments the agent is about to narrate in the final
+    // reply. Without this, a closing "here's the diff …" message would
+    // render two stacked copies of the same patch. Only touches segments
+    // with `kind: 'diff'` emitted by pushInlineDiffSegment — real
+    // assistant narration stays put.
+    const finalHasOwnDiffFence = /```(?:diff|patch)\b/i.test(finalText)
+
+    const segments = this.segmentMessages.filter(msg => {
+      const body = diffSegmentBody(msg)
+
+      return body === null || (!finalHasOwnDiffFence && !finalText.includes(body))
+    })
+
+    const finalMessages = [...segments]
+
+    if (finalText) {
       finalMessages.push({
         role: 'assistant',
-        text: mergedText,
+        text: finalText,
         thinking: savedReasoning || undefined,
         thinkingTokens: savedReasoning ? savedReasoningTokens : undefined,
         toolTokens: savedToolTokens || undefined,
@@ -300,7 +328,7 @@ class TurnController {
     this.bufRef = ''
     patchTurnState({ activity: [], outcome: '' })
 
-    return { finalMessages, finalText: mergedText, wasInterrupted }
+    return { finalMessages, finalText, wasInterrupted }
   }
 
   recordMessageDelta({ rendered, text }: { rendered?: string; text?: string }) {
@@ -406,7 +434,6 @@ class TurnController {
     this.bufRef = ''
     this.interrupted = false
     this.lastStatusNote = ''
-    this.pendingInlineDiffs = []
     this.pendingSegmentTools = []
     this.protocolWarned = false
     this.segmentMessages = []
@@ -452,7 +479,6 @@ class TurnController {
     this.endReasoningPhase()
     this.clearReasoning()
     this.activeTools = []
-    this.pendingInlineDiffs = []
     this.turnTools = []
     this.toolTokenAcc = 0
     this.persistedToolLabels.clear()

ui-tui/src/components/messageLine.tsx

@@ -81,11 +81,16 @@ export const MessageLine = memo(function MessageLine({
     return <Text {...(body ? { color: body } : {})}>{msg.text}</Text>
   })()
 
+  // Diff segments (emitted by pushInlineDiffSegment between narration
+  // segments) need a blank line on both sides so the patch doesn't butt up
+  // against the prose around it.
+  const isDiffSegment = msg.kind === 'diff'
+
   return (
     <Box
       flexDirection="column"
-      marginBottom={msg.role === 'user' ? 1 : 0}
+      marginBottom={msg.role === 'user' || isDiffSegment ? 1 : 0}
-      marginTop={msg.role === 'user' || msg.kind === 'slash' ? 1 : 0}
+      marginTop={msg.role === 'user' || msg.kind === 'slash' || isDiffSegment ? 1 : 0}
     >
       {showDetails && (
         <Box flexDirection="column" marginBottom={1}>

ui-tui/src/types.ts

@@ -102,7 +102,7 @@ export interface ClarifyReq {
 
 export interface Msg {
   info?: SessionInfo
-  kind?: 'intro' | 'panel' | 'slash' | 'trail'
+  kind?: 'diff' | 'intro' | 'panel' | 'slash' | 'trail'
   panelData?: PanelData
   role: Role
   text: string