client: restore host terminal on session exit

Author: seruman

client/attach.go

@@ -117,8 +117,6 @@ func (c *Client) RunAttach(name string, command []string, dk DetachKey, forwardE
 		}
 	}
 
-	// Enter raw mode after writing the state dump so OPOST translates
-	// \n in the VT dump to \r\n correctly while in cooked mode.
 	oldState, err := term.MakeRaw(fd)
 	if err != nil {
 		return fmt.Errorf("set raw mode: %w", err)
@@ -201,11 +199,18 @@ func (c *Client) RunAttach(name string, command []string, dk DetachKey, forwardE
 		}
 	}()
 
+	if attached.Created && len(command) > 0 && len(attached.ScreenDump) > 0 {
+		if _, err := os.Stdout.Write(attached.ScreenDump); err != nil {
+			return fmt.Errorf("write state dump: %w", err)
+		}
+	}
+
 	handleMsg := func(msg protocol.Message) error {
 		switch m := msg.(type) {
 		case *protocol.Output:
 			os.Stdout.Write(m.Data)
 		case *protocol.Exited:
+			restoreHostTerminal(fd, oldState, "[hauntty] session exited\n")
 			return &ExitError{Code: int(m.ExitCode)}
 		case *protocol.Error:
 			term.Restore(fd, oldState)
@@ -222,36 +227,40 @@ func (c *Client) RunAttach(name string, command []string, dk DetachKey, forwardE
 		if err != nil {
 			close(done)
 			if err == io.EOF || isConnClosed(err) {
-				// Use 1047 (not 1049) to exit alt screen: 1047
-				// just switches the buffer without restoring the
-				// saved cursor, so session content on the primary
-				// screen stays intact. No-op when already on primary.
-				//
-				// Reset modes, show cursor, pop kitty keyboard,
-				// reset SGR, erase from cursor to end of screen.
-				// Session content above the cursor is preserved.
-				os.Stdout.Write([]byte(
-					"\x1b[?1047;1;1000;1002;1003;1006;1004;2004;2048;2026l" +
-						"\x1b[?25h" +
-						"\x1b[<u" +
-						"\x1b[0m" +
-						"\x1b[J"))
-				drainStdin(fd, 20*time.Millisecond)
-				term.Restore(fd, oldState)
-				fmt.Fprintf(os.Stderr, "[hauntty] detached\n")
+				restoreHostTerminal(fd, oldState, "[hauntty] detached\n")
 				return nil
 			}
 			term.Restore(fd, oldState)
 			return fmt.Errorf("read message: %w", err)
 		}
 		if err := handleMsg(msg); err != nil {
 			close(done)
-			term.Restore(fd, oldState)
 			return err
 		}
 	}
 }
 
+func restoreHostTerminal(fd int, oldState *term.State, message string) {
+	// Use 1047 (not 1049) to exit alt screen: 1047 just switches the
+	// buffer without restoring the saved cursor, so session content on
+	// the primary screen stays intact. No-op when already on primary.
+	//
+	// Reset modes, show cursor, pop kitty keyboard, reset SGR, erase
+	// from cursor to end of screen. Session content above the cursor is
+	// preserved.
+	os.Stdout.Write([]byte(
+		"\x1b[?1047;1;1000;1002;1003;1006;1004;2004;2048;2026l" +
+			"\x1b[?25h" +
+			"\x1b[<u" +
+			"\x1b[0m" +
+			"\x1b[J"))
+	drainStdin(fd, 20*time.Millisecond)
+	term.Restore(fd, oldState)
+	if message != "" {
+		fmt.Fprint(os.Stderr, message)
+	}
+}
+
 // readCursorRow reads the DSR response (\x1b[{row};{col}R) from fd
 // and returns the cursor row. Must be called in raw mode after sending
 // \x1b[6n. Returns fallback if the response cannot be parsed.

cmd/ht/e2e_test/automation_test.go

@@ -160,8 +160,8 @@ func TestDumpDeadSessionPreservesFormats(t *testing.T) {
 	sh.Type("$HT_BIN attach dead-dump-formats -- /bin/sh -c \"printf '\\033[31mred\\033[0m\\nplain\\n'; sleep 30\"\n")
 	sh.WaitFor("red")
 	sh.WaitFor("plain")
+	sh.WaitStable(250*time.Millisecond, termtest.WaitTimeout(2*time.Second))
 	sh.Key(libghostty.KeyCode(']'), libghostty.ModCtrl)
-	sh.WaitFor("detached")
 	e.waitHostPrompt(sh)
 
 	kill := e.run("kill", "dead-dump-formats")
@@ -193,6 +193,7 @@ func TestDumpDeadSessionPreservesJoinFlag(t *testing.T) {
 	e.waitHostPrompt(sh)
 	sh.Type("$HT_BIN attach dead-dump-join -- /bin/sh -c \"printf 'aaaaaaaaaaaaaaaaaaaabbbbbbbbbb\\n'; sleep 30\"\n")
 	sh.WaitFor("bbbb")
+	sh.WaitStable(250*time.Millisecond, termtest.WaitTimeout(2*time.Second))
 	sh.Key(libghostty.KeyCode(']'), libghostty.ModCtrl)
 	e.waitHostPrompt(sh)
 

cmd/ht/e2e_test/lifecycle_test.go

@@ -370,6 +370,53 @@ func TestRestoreDeadSession(t *testing.T) {
 	e.waitHostPrompt(restoreSh)
 }
 
+func TestRestoreDeadSessionAfterHostOutput(t *testing.T) {
+	cfg := config.Default()
+	cfg.Client.DetachKeybind = "ctrl+]"
+	cfg.Daemon.StatePersistence = true
+	e := setup(t, cfg)
+
+	daemon := e.term([]string{htBin, "daemon"})
+	daemon.WaitFor("daemon listening")
+
+	sh := e.term([]string{"/bin/sh"}, termtest.WithEnv("PS1=$ ", "SHELL=/bin/sh"))
+	e.waitHostPrompt(sh)
+	sh.Type("$HT_BIN attach restore-host-output\n")
+	sh.WaitFor("created session")
+	e.waitAttachedPrompt(sh)
+
+	sh.Type("echo restore-marker\n")
+	sh.WaitFor("restore-marker")
+	sh.WaitStable(250*time.Millisecond, termtest.WaitTimeout(2*time.Second))
+
+	sh.Key(libghostty.KeyCode(']'), libghostty.ModCtrl)
+	sh.WaitFor("detached")
+	e.waitHostPrompt(sh)
+
+	kill := e.run("kill", "restore-host-output")
+	kill.Assert(t, icmd.Expected{ExitCode: 0})
+	time.Sleep(500 * time.Millisecond)
+
+	dump := e.run("dump", "restore-host-output")
+	dump.Assert(t, icmd.Expected{ExitCode: 0})
+
+	restoreSh := e.term([]string{"/bin/sh"}, termtest.WithEnv("PS1=$ ", "SHELL=/bin/sh"))
+	e.waitHostPrompt(restoreSh)
+	restoreSh.Type("echo host-before\n")
+	restoreSh.WaitFor("host-before")
+	restoreSh.Type("$HT_BIN restore restore-host-output\n")
+	restoreSh.WaitFor("attached to session")
+	restoreSh.WaitFor("restore-marker")
+	e.waitAttachedPrompt(restoreSh)
+
+	restoreSh.Type("echo restored-ok\n")
+	restoreSh.WaitFor("restored-ok")
+
+	restoreSh.Key(libghostty.KeyCode(']'), libghostty.ModCtrl)
+	restoreSh.WaitFor("detached")
+	e.waitHostPrompt(restoreSh)
+}
+
 func TestRestoreRunningSessionFails(t *testing.T) {
 	cfg := config.Default()
 	cfg.Daemon.AutoExit = true

cmd/ht/e2e_test/main_test.go

@@ -18,7 +18,7 @@ import (
 
 var (
 	htBin        string
-	hostPromptRE = regexp.MustCompile(`\$ ?$`)
+	hostPromptRE = regexp.MustCompile(`^\$ ?$`)
 )
 
 type testEnv struct {

cmd/ht/e2e_test/terminal_safety_test.go

@@ -295,6 +295,7 @@ func TestDetachFailurePathHostCleanup(t *testing.T) {
 	kill := e.run("kill", "fail-cleanup")
 	kill.Assert(t, icmd.Expected{ExitCode: 0, Out: "killed session \"fail-cleanup\"\n"})
 
+	sh.WaitFor("session exited")
 	e.waitHostPrompt(sh)
 	sh.Type("echo failure-cleanup-ok\n")
 	sh.WaitFor("failure-cleanup-ok")

internal/daemon/session.go

@@ -24,6 +24,11 @@ var feedPool = sync.Pool{
 	},
 }
 
+type feedItem struct {
+	data *[]byte
+	seq  uint64
+}
+
 // All fields are owned exclusively by the session's run loop.
 type sessionClient struct {
 	id        string
@@ -95,13 +100,14 @@ type Session struct {
 	ptmx    *os.File
 	cmd     *exec.Cmd
 	term    *libghostty.Terminal
-	feedCh  chan *[]byte
+	feedCh  chan feedItem
 	tempDir string
 
-	actions  chan any
-	ptyOut   chan []byte
-	done     chan struct{}
-	exitCode int32
+	actions     chan any
+	ptyOut      chan []byte
+	done        chan struct{}
+	exitCode    int32
+	feedApplied atomic.Uint64
 
 	// sizeVal packs cols|rows as (cols<<16)|rows for lock-free reads.
 	sizeVal atomic.Uint32
@@ -191,7 +197,7 @@ func newSession(ctx context.Context, name string, command []string, env []string
 		ptmx:         ptmx,
 		cmd:          cmd,
 		term:         term,
-		feedCh:       make(chan *[]byte, 64),
+		feedCh:       make(chan feedItem, 64),
 		tempDir:      tempDir,
 		actions:      make(chan any, 16),
 		ptyOut:       make(chan []byte, 64),
@@ -270,7 +276,7 @@ func restoreSession(ctx context.Context, name string, command []string, env []st
 		ptmx:         ptmx,
 		cmd:          cmd,
 		term:         term,
-		feedCh:       make(chan *[]byte, 64),
+		feedCh:       make(chan feedItem, 64),
 		tempDir:      tempDir,
 		actions:      make(chan any, 16),
 		ptyOut:       make(chan []byte, 64),
@@ -287,12 +293,20 @@ func restoreSession(ctx context.Context, name string, command []string, env []st
 }
 
 func (s *Session) feedLoop(ctx context.Context) {
-	for bp := range s.feedCh {
-		if err := s.term.Feed(ctx, *bp); err != nil {
+	for item := range s.feedCh {
+		if err := s.term.Feed(ctx, *item.data); err != nil {
 			slog.Debug("wasm feed error", "session", s.Name, "err", err)
 		}
-		*bp = (*bp)[:cap(*bp)]
-		feedPool.Put(bp)
+		s.feedApplied.Store(item.seq)
+		*item.data = (*item.data)[:cap(*item.data)]
+		feedPool.Put(item.data)
+	}
+}
+
+// waitFeedApplied blocks until feedLoop has applied every PTY chunk up to target.
+func (s *Session) waitFeedApplied(target uint64) {
+	for s.feedApplied.Load() < target {
+		time.Sleep(100 * time.Microsecond)
 	}
 }
 
@@ -331,16 +345,19 @@ func (s *Session) run() {
 	// pendingFeed holds data waiting to be sent to feedCh. While
 	// non-nil, we stop reading ptyOut (backpressure) but keep
 	// processing actions so detach/kick/list don't stall.
-	var pendingFeed *[]byte
+	var pendingFeed *feedItem
+	var nextFeedSeq uint64
 
 	for {
 		// Nil-channel trick: only one of ptyCh/feedSend is active
 		// at a time. When pendingFeed is nil, read ptyOut. When
 		// non-nil, send to feedCh. Actions are always processed.
 		var ptyCh <-chan []byte
-		var feedSend chan<- *[]byte
+		var feedSend chan<- feedItem
+		var feedItemToSend feedItem
 		if pendingFeed != nil {
 			feedSend = s.feedCh
+			feedItemToSend = *pendingFeed
 		} else {
 			ptyCh = s.ptyOut
 		}
@@ -373,9 +390,10 @@ func (s *Session) run() {
 			d := (*bp)[:len(data)]
 			copy(d, data)
 			*bp = d
-			pendingFeed = bp
+			nextFeedSeq++
+			pendingFeed = &feedItem{data: bp, seq: nextFeedSeq}
 
-		case feedSend <- pendingFeed:
+		case feedSend <- feedItemToSend:
 			pendingFeed = nil
 
 		case action := <-s.actions:
@@ -384,6 +402,12 @@ func (s *Session) run() {
 				if !a.readOnly {
 					s.resizeForPending(clients, a.cols, a.rows, a.xpixel, a.ypixel)
 				}
+				if pendingFeed != nil {
+					s.feedCh <- *pendingFeed
+					pendingFeed = nil
+				}
+				// Attach dumps must reflect every PTY chunk we've already accepted.
+				s.waitFeedApplied(nextFeedSeq)
 
 				dump, err := s.term.DumpScreen(s.ctx, libghostty.DumpVTFull)
 				if err != nil {
@@ -485,7 +509,7 @@ func (s *Session) run() {
 				// Clients see connection close (EOF), not Exited — this is
 				// the kill/shutdown path.
 				if pendingFeed != nil {
-					feedPool.Put(pendingFeed)
+					feedPool.Put(pendingFeed.data)
 					pendingFeed = nil
 				}
 				close(s.feedCh)