From 3688cd7301703e687cc3d8562cebe2558e6f2d06 Mon Sep 17 00:00:00 2001 From: Amit Kumar Date: Thu, 11 Jun 2026 07:57:36 +0000 Subject: [PATCH] fix(attach): own the terminal state across attach and detach The native attach client was a verbatim byte bridge with no terminal ownership: the host replay opens with \x1b[2J on whatever screen the terminal shows, and detach restored only termios. Attaching from the TUI therefore painted the session over the user's PRIMARY screen (tea.ExecProcess leaves the alt screen before the child runs), so quitting uam revealed session output instead of the shell; and any modes the agent enabled mid-attach (mouse tracking, kitty keyboard, bracketed paste, hidden cursor, charsets) stayed active on the real terminal, corrupting the TUI that resumes after detach. Own the window the way `tmux attach` did: - on a tty, enter our own alternate screen (\x1b[?1049h) before the replay, so the session never touches the primary screen - on every exit path (detach, session end, and SIGINT/SIGTERM/SIGHUP, now handled like a detach) emit a mode-reset suffix - kitty pop + zero, mouse modes and focus reporting off, bracketed paste off, synchronized output off, DECSTR, keypad and charset reset, cursor visible - then leave the alt screen, then restore termios, all under one sync.Once - stop and drain the host-to-terminal pump before the restore, so bytes still buffered from the socket land inside the alt screen, never on the restored primary Piped (non-tty) attaches stay byte-identical, asserted by test. Verified end to end on a real PTY pair through the session host. --- internal/session/attach.go | 70 +++++++++++- internal/session/session_test.go | 179 +++++++++++++++++++++++++++++++ 2 files changed, 244 insertions(+), 5 deletions(-) diff --git a/internal/session/attach.go b/internal/session/attach.go index f972fb8..e1054eb 100644 --- a/internal/session/attach.go +++ b/internal/session/attach.go @@ -27,6 +27,32 @@ const detachPrefix = 0x02 // 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[" + // 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 @@ -69,20 +95,42 @@ func runAttach(dir, name string, stdin *os.File, stdout *os.File) error { 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 @@ func runAttach(dir, name string, stdin *os.File, stdout *os.File) error { }() // 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 @@ func runAttach(dir, name string, stdin *os.File, stdout *os.File) error { 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 diff --git a/internal/session/session_test.go b/internal/session/session_test.go index ea24acc..96eb5ff 100644 --- a/internal/session/session_test.go +++ b/internal/session/session_test.go @@ -5,9 +5,12 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" "time" + "github.com/creack/pty" + "github.com/RandomCodeSpace/unified-agent-manager/internal/store" ) @@ -296,6 +299,182 @@ func TestAttachStreamsAndDetaches(t *testing.T) { 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