Fix styled switch-back repaint

internal/app/app.go

@@ -215,7 +215,9 @@ type uiState struct {
 	// renderer confines focused-child live output to the main viewport.
 	// A fresh renderer is allocated per focused child so partial-escape
 	// state cannot bleed between panes.
-	renderer *viewportRenderer
+	renderer             *viewportRenderer
+	repaintNextPTY       string
+	repaintNextPTYBudget int
 
 	// attention is the latest request_human_attention surfaced via MCP;
 	// rendered in the status line until cleared.
@@ -350,15 +352,33 @@ func (st *uiState) OnChildExited(c *Child) {
 // disabled only around the replay so long styled runs cannot wrap into
 // the right rail.
 func (st *uiState) OnPTYOut(childID string, chunk []byte) {
+	layout := st.layoutSnapshot()
 	st.mu.Lock()
 	focus := st.focusedID
 	palOpen := st.palette != nil
 	renderer := st.renderer
+	forceRepaint := focus == childID && st.repaintNextPTY == childID && st.repaintNextPTYBudget > 0
+	if forceRepaint {
+		renderer = newViewportRenderer(layout)
+		st.renderer = renderer
+		st.repaintNextPTYBudget--
+		if st.repaintNextPTYBudget == 0 {
+			st.repaintNextPTY = ""
+		}
+	}
 	st.mu.Unlock()
 	if palOpen || focus != childID || renderer == nil {
 		return
 	}
-	out := renderer.Render(chunk)
+	var out []byte
+	if forceRepaint {
+		out = st.renderFocusedSnapshot(childID, renderer, layout)
+		if len(out) == 0 {
+			return
+		}
+	} else {
+		out = renderer.Render(chunk)
+	}
 	st.outMu.Lock()
 	_, _ = os.Stdout.Write([]byte("\x1b[?7l"))
 	_, _ = os.Stdout.Write(out)
@@ -798,11 +818,11 @@ func (st *uiState) flashTransient(msg string) {
 // Callers must NOT hold st.mu — repaintFocused takes it
 // briefly itself.
 //
-// We feed the emulator's VT serialization through the viewport
-// renderer so SGR styling, alt-screen state, and the cursor position
-// survive a focus switch. The plain-text path (renderScreenSnapshot)
-// is kept as a fallback for environments where SerializeVT is
-// unavailable (e.g. the nocgo stub).
+// We replay the emulator's padded grid snapshot rather than its VT
+// serialization. SerializeVT can preserve style, but for diff-based TUIs
+// we've seen it replay stale prompt layout that no longer matches the
+// emulator grid; the padded snapshot is the source of truth for visible
+// cells.
 func (st *uiState) repaintFocused() {
 	st.mu.Lock()
 	id := st.focusedID
@@ -824,16 +844,28 @@ func (st *uiState) repaintFocused() {
 		defer c.NudgeRedraw(cols, rows)
 	}
 
+	out := st.renderFocusedSnapshot(id, renderer, layout)
+	if len(out) == 0 {
+		return
+	}
+	st.mu.Lock()
+	if st.focusedID == id {
+		st.repaintNextPTY = id
+		st.repaintNextPTYBudget = 8
+	}
+	st.mu.Unlock()
+	st.outMu.Lock()
+	defer st.outMu.Unlock()
+	_, _ = os.Stdout.Write(out)
+}
+
+func (st *uiState) renderFocusedSnapshot(id string, renderer *viewportRenderer, layout terminalLayout) []byte {
+	text, cursor, err := st.sess.SnapshotChild(id)
+	if err != nil {
+		return nil
+	}
 	if renderer != nil {
-		if serialized, err := st.sess.SerializeChild(id); err == nil && len(serialized) > 0 {
-			// Reset host terminal state before replaying so leftover
-			// modes from the previously-focused child (DECSTBM,
-			// DECOM, SGR) don't distort the snapshot. The DECSTBM is
-			// pinned to the viewport region in host coordinates; the
-			// cursor parks at the viewport's top-left. The replayed
-			// SerializeVT may re-set these modes if the child
-			// configured them, which is fine — we're just guaranteeing
-			// a known starting baseline.
+		if styled, err := st.sess.StyledSnapshotChild(id); err == nil && len(styled) > 0 {
 			mainBottom := int(layout.statusRow) - statusRows
 			prelude := fmt.Sprintf(
 				"\x1b[0m\x1b[?6l\x1b[%d;%dr\x1b[?25h\x1b[%d;%dH",
@@ -841,40 +873,19 @@ func (st *uiState) repaintFocused() {
 				int(layout.mainTop), int(layout.mainLeft),
 			)
 			out := []byte(prelude)
-			out = append(out, renderer.Render(serialized)...)
-			// Ghostty's VT serialization emits the cursor CUP, then
-			// DECSTBM (which moves the cursor to region home as a
-			// documented side effect), then tab-stop setup using CHA
-			// (\x1b[NG) — which leaves the renderer's internal vr.col
-			// tracking pointing at the last tab-stop column, not
-			// where the cursor actually ended up. Re-emit the saved
-			// cursor as a child-space CUP through the renderer so
-			// (a) the host cursor lands at the right place and (b)
-			// the renderer's internal row/col tracking is brought
-			// back in sync with the host. Without this, subsequent
-			// relative moves (CSI C/D) and erase-line widths (CSI K
-			// uses vr.col) operate from a stale column and the input
-			// box gets drawn at the wrong width / row.
-			if _, cursor, err := st.sess.SnapshotChild(id); err == nil {
-				cup := fmt.Sprintf("\x1b[%d;%dH",
-					int(cursor.Row)+1, int(cursor.Col)+1)
-				out = append(out, renderer.Render([]byte(cup))...)
-			}
-			st.outMu.Lock()
-			defer st.outMu.Unlock()
-			_, _ = os.Stdout.Write(out)
-			return
+			out = append(out, renderer.ClearViewport()...)
+			out = append(out, renderer.Render(styled)...)
+			cup := fmt.Sprintf("\x1b[%d;%dH", int(cursor.Row)+1, int(cursor.Col)+1)
+			out = append(out, renderer.Render([]byte(cup))...)
+			return out
 		}
 	}
-
-	text, cursor, err := st.sess.SnapshotChild(id)
-	if err != nil {
-		return
-	}
 	out := renderScreenSnapshot(text, cursor, layout)
-	st.outMu.Lock()
-	defer st.outMu.Unlock()
-	_, _ = os.Stdout.Write(out)
+	if renderer != nil {
+		cup := fmt.Sprintf("\x1b[%d;%dH", int(cursor.Row)+1, int(cursor.Col)+1)
+		out = append(out, renderer.Render([]byte(cup))...)
+	}
+	return out
 }
 
 func (st *uiState) requestExit() {