fix(attach): own the terminal state across attach and detach

internal/session/attach.go

@@ -27,6 +27,32 @@
 // binding C-z to a warning.
 const ctrlZ = 0x1a
 
+// The attach client is a verbatim byte bridge: the host replay opens with a
+// clear-screen and every escape sequence the agent emits while a client is
+// attached lands on the user's real terminal. `tmux attach` confined all of
+// that by running the session inside its own alternate screen and resetting
+// modes on detach; without the same ownership the session scribbles over the
+// user's primary screen (still visible after uam exits) and leaked modes
+// corrupt the TUI that resumes after detach.
+
+// screenEnter opens the attach client's own alternate screen, saving the
+// primary screen and cursor underneath.
+const screenEnter = "\x1b[?1049h"
+
+// screenExit resets every mode the agent could have toggled mid-attach, then
+// leaves the alternate screen. Terminals ignore sequences they don't
+// implement, so the suffix is safe to emit unconditionally.
+const screenExit = "\x1b[<u" + // pop the kitty keyboard flags agents push
+	"\x1b[=0;1u" + // and zero them in case the agent pushed more than once
+	"\x1b[?1000;1002;1003;1004;1005;1006;1015l" + // mouse tracking + focus reporting off
+	"\x1b[?2004l" + // bracketed paste off
+	"\x1b[?2026l" + // synchronized output off
+	"\x1b[!p" + // DECSTR: cursor keys, origin, margins, SGR, insert mode
+	"\x1b>" + // numeric keypad (DECKPNM; DECSTR leaves keypad mode alone)
+	"\x1b(B" + // G0 charset back to ASCII
+	"\x1b[?25h" + // cursor visible
+	"\x1b[?1049l" // leave the alt screen: primary buffer and cursor restored
+
 // RunAttach is the entry point of `uam __attach`: it puts the terminal in raw
 // mode and bridges it to a session host — the native replacement for
 // `tmux attach`. It returns when the user detaches (Ctrl+B d, or a bare left
@@ -43,7 +69,7 @@
 	return runAttach(*dir, fs.Arg(0), os.Stdin, os.Stdout)
 }
 
 func runAttach(dir, name string, stdin *os.File, stdout *os.File) error {
 	if err := ValidateName(name); err != nil {
 		return err
 	}
@@ -69,20 +95,42 @@
 		return fmt.Errorf("attach %s: %s", name, resp.Err)
 	}
 
-	restore := func() {}
+	var ttyState *term.State
 	if term.IsTerminal(stdin.Fd()) {
 		state, err := term.MakeRaw(stdin.Fd())
 		if err != nil {
 			return fmt.Errorf("set raw mode: %w", err)
 		}
-		var once sync.Once
-		restore = func() { once.Do(func() { _ = term.Restore(stdin.Fd(), state) }) }
-		defer restore()
+		ttyState = state
+	}
+	ownScreen := term.IsTerminal(stdout.Fd())
+	if ownScreen {
+		_, _ = stdout.WriteString(screenEnter)
+	}
+	var once sync.Once
+	restore := func() {
+		once.Do(func() {
+			if ownScreen {
+				_, _ = stdout.WriteString(screenExit)
+			}
+			if ttyState != nil {
+				_ = term.Restore(stdin.Fd(), ttyState)
+			}
+		})
 	}
+	defer restore()
 
 	winch := make(chan os.Signal, 1)
 	signal.Notify(winch, syscall.SIGWINCH)
 	defer signal.Stop(winch)
+
+	// An external SIGINT/SIGTERM/SIGHUP must restore the screen and termios
+	// like a detach would, or the terminal is left raw on the agent's output.
+	// Ctrl+C inside the session never lands here: raw mode clears ISIG, so it
+	// reaches the agent as a plain byte.
+	quit := make(chan os.Signal, 1)
+	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
+	defer signal.Stop(quit)
 	go func() {
 		for range winch {
 			if w, h, err := term.GetSize(stdout.Fd()); err == nil {
@@ -100,11 +148,13 @@
 	}()
 
 	// host → terminal (the main loop): ends when the host closes the
-	// connection (agent exited) or the user detached.
+	// connection (agent exited) or the user detached. done is closed once the
+	// pump has fully drained, so a second receive never blocks.
 	done := make(chan error, 1)
 	go func() {
 		_, err := io.Copy(stdout, br)
 		done <- err
+		close(done)
 	}()
 	var note string
 	select {
@@ -113,7 +163,17 @@
 		note = "detached"
 	case <-done:
 		note = "session ended"
+	case <-quit:
+		_ = writeFrame(conn, frameDetach, nil)
+		note = "detached"
 	}
+	// Stop the host→terminal pump and drain it before restoring the screen:
+	// bytes still buffered from the socket must land inside the alternate
+	// screen, not on the primary screen revealed after screenExit. On the
+	// session-ended path the pump has already finished and done is closed, so
+	// this returns immediately.
+	_ = conn.Close()
+	<-done
 	restore()
 	_, _ = fmt.Fprintf(stdout, "\r\n[uam: %s]\r\n", note)
 	return nil

internal/session/session_test.go

@@ -5,9 +5,12 @@
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
+	"github.com/creack/pty"
+
 	"github.com/RandomCodeSpace/unified-agent-manager/internal/store"
 )
 
@@ -296,6 +299,182 @@
 	if !strings.Contains(out, "[uam: detached]") {
 		t.Fatalf("missing detach notice: %q", out)
 	}
+	// Terminal-state ownership is tty-only: piped output must stay free of
+	// alternate-screen switches.
+	if strings.Contains(out, "\x1b[?1049") {
+		t.Fatalf("non-tty attach must not emit alt-screen sequences: %q", out)
+	}
+}
+
+// On a real terminal the attach client must own the screen the way
+// `tmux attach` did: bridge the session inside its own alternate screen (so
+// the replay's clear and live agent output never land on the user's primary
+// screen — the buffer uam's TUI reveals again when it exits), and on detach
+// reset the modes an agent could have leaked (mouse tracking, bracketed
+// paste, kitty keyboard, hidden cursor) before handing the primary screen
+// back. Regression test for the post-tmux rendering bugs: corrupted TUI after
+// attach/detach and session residue left on the terminal after quitting uam.
+func TestAttachOwnsTerminalStateOnTTY(t *testing.T) {
+	c := newTestClient(t)
+	ctx := context.Background()
+	name := "uam-fake-aaaa9999"
+	if err := c.CreateSession(ctx, name, t.TempDir(), nil, []string{"/bin/sh", "-c", "echo banner; sleep 60"}); err != nil {
+		t.Fatalf("CreateSession: %v", err)
+	}
+	waitFor(t, "banner", func() bool {
+		out, _ := c.Capture(ctx, name, 10)
+		return strings.Contains(out, "banner")
+	})
+
+	ptmx, tty, err := pty.Open()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ptmx.Close(); _ = tty.Close() }()
+
+	done := make(chan error, 1)
+	go func() { done <- runAttach(c.Dir, name, tty, tty) }()
+
+	var mu sync.Mutex
+	var got []byte
+	go func() {
+		buf := make([]byte, 4096)
+		for {
+			n, rerr := ptmx.Read(buf)
+			if n > 0 {
+				mu.Lock()
+				got = append(got, buf[:n]...)
+				mu.Unlock()
+			}
+			if rerr != nil {
+				return
+			}
+		}
+	}()
+	snapshot := func() string {
+		mu.Lock()
+		defer mu.Unlock()
+		return string(got)
+	}
+
+	waitFor(t, "replay banner", func() bool { return strings.Contains(snapshot(), "banner") })
+	pre := snapshot()
+	enter := strings.Index(pre, "\x1b[?1049h")
+	if enter < 0 {
+		t.Fatalf("tty attach must enter its own alternate screen: %q", pre)
+	}
+	if clear := strings.Index(pre, "\x1b[2J"); clear >= 0 && clear < enter {
+		t.Fatalf("replay clear must land inside the alt screen, not on the primary: %q", pre)
+	}
+
+	if _, err := ptmx.Write([]byte{0x02, 'd'}); err != nil { // Ctrl+B d
+		t.Fatal(err)
+	}
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("runAttach: %v", err)
+		}
+	case <-time.After(10 * time.Second):
+		t.Fatal("detach chord did not detach")
+	}
+	waitFor(t, "detach notice", func() bool { return strings.Contains(snapshot(), "[uam: detached]") })
+	full := snapshot()
+	exit := strings.Index(full, "\x1b[?1049l")
+	if exit < 0 {
+		t.Fatalf("detach must leave the alternate screen: %q", full)
+	}
+	for _, reset := range []string{
+		"\x1b[?1000;1002;1003;1004;1005;1006;1015l", // mouse tracking + focus reporting off
+		"\x1b[?2004l", // bracketed paste off
+		"\x1b[?25h",   // cursor visible
+	} {
+		idx := strings.Index(full, reset)
+		if idx < 0 {
+			t.Fatalf("detach must reset leakable terminal modes (missing %q): %q", reset, full)
+		}
+		if idx > exit {
+			t.Fatalf("mode reset %q must precede the alt-screen exit: %q", reset, full)
+		}
+	}
+	if note := strings.Index(full, "[uam: detached]"); note < exit {
+		t.Fatalf("detach notice must print on the restored primary screen: %q", full)
+	}
+	if !c.HasSession(ctx, name) {
+		t.Fatal("session must keep running after detach")
+	}
+}
+
+// Detaching while the agent is mid-burst must not scribble the primary
+// screen: bytes still buffered in the host→terminal pump at detach time have
+// to be drained inside the alternate screen, so nothing follows the
+// alt-screen exit except the detach notice.
+func TestAttachDetachDrainsPumpBeforeScreenRestore(t *testing.T) {
+	c := newTestClient(t)
+	ctx := context.Background()
+	name := "uam-fake-bbbb0000"
+	err := c.CreateSession(ctx, name, t.TempDir(), nil, []string{"/bin/sh", "-c", "i=0; while [ $i -lt 2000 ]; do echo spam$i; i=$((i+1)); done; sleep 60"})
+	if err != nil {
+		t.Fatalf("CreateSession: %v", err)
+	}
+
+	ptmx, tty, err := pty.Open()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ptmx.Close(); _ = tty.Close() }()
+
+	done := make(chan error, 1)
+	go func() { done <- runAttach(c.Dir, name, tty, tty) }()
+
+	var mu sync.Mutex
+	var got []byte
+	go func() {
+		buf := make([]byte, 4096)
+		for {
+			n, rerr := ptmx.Read(buf)
+			if n > 0 {
+				mu.Lock()
+				got = append(got, buf[:n]...)
+				mu.Unlock()
+			}
+			if rerr != nil {
+				return
+			}
+		}
+	}()
+	snapshot := func() string {
+		mu.Lock()
+		defer mu.Unlock()
+		return string(got)
+	}
+
+	// Detach as soon as the burst starts flowing, while plenty of it is
+	// still in flight between the host and the terminal.
+	waitFor(t, "burst start", func() bool { return strings.Contains(snapshot(), "spam") })
+	if _, err := ptmx.Write([]byte{0x02, 'd'}); err != nil { // Ctrl+B d
+		t.Fatal(err)
+	}
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("runAttach: %v", err)
+		}
+	case <-time.After(10 * time.Second):
+		t.Fatal("detach chord did not detach")
+	}
+	waitFor(t, "detach notice", func() bool { return strings.Contains(snapshot(), "[uam: detached]") })
+	full := snapshot()
+	exit := strings.LastIndex(full, "\x1b[?1049l")
+	if exit < 0 {
+		t.Fatalf("detach must leave the alternate screen: %q", full)
+	}
+	// The note is written after termios are restored, so the pty line
+	// discipline may ONLCR-translate its newlines — compare CR-insensitively.
+	tail := strings.ReplaceAll(full[exit+len("\x1b[?1049l"):], "\r", "")
+	if tail != "\n[uam: detached]\n" {
+		t.Fatalf("only the detach notice may follow the alt-screen exit, got %q", tail)
+	}
 }
 
 // The left-arrow quick detach works end to end through the real attach