client: restore host terminal on session exit

Author: seruman

client/attach.go

@@ -206,6 +206,7 @@ func (c *Client) RunAttach(name string, command []string, dk DetachKey, forwardE
 		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 +223,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/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")