Fix exit not closing session

nativessh/pty.go

@@ -1,9 +1,11 @@
 package nativessh
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"os/exec"
+	"sync"
 
 	"github.com/creack/pty"
 )
@@ -11,8 +13,10 @@ import (
 // PTYSession is a running shell process attached to a PTY.
 // It implements io.ReadWriteCloser so it can be bridged to any transport.
 type PTYSession struct {
-	ptmx *os.File
-	cmd  *exec.Cmd
+	ptmx     *os.File
+	cmd      *exec.Cmd
+	waitOnce sync.Once
+	exitCode int
 }
 
 // findShell returns the path to the best available interactive shell by
@@ -54,9 +58,32 @@ func (p *PTYSession) Resize(cols, rows uint16) error {
 	return pty.Setsize(p.ptmx, &pty.Winsize{Cols: cols, Rows: rows})
 }
 
+// wait waits for the child process to exit exactly once and records its exit
+// code. Safe to call concurrently or multiple times.
+func (p *PTYSession) wait() {
+	p.waitOnce.Do(func() {
+		err := p.cmd.Wait()
+		if err != nil {
+			var exitErr *exec.ExitError
+			if errors.As(err, &exitErr) {
+				p.exitCode = exitErr.ExitCode()
+				return
+			}
+			p.exitCode = 1
+		}
+	})
+}
+
+// ExitCode waits for the shell process to exit and returns its exit code.
+// It is safe to call before or after Close.
+func (p *PTYSession) ExitCode() int {
+	p.wait()
+	return p.exitCode
+}
+
 // Close closes the PTY and waits for the child process to exit.
 func (p *PTYSession) Close() error {
 	err := p.ptmx.Close()
-	_ = p.cmd.Wait()
+	p.wait()
 	return err
 }

nativessh/server.go

@@ -177,8 +177,18 @@ func (s *Server) handleSession(ch ssh.Channel, requests <-chan *ssh.Request) {
 			// PTY output → SSH channel.
 			go func() {
 				_, _ = io.Copy(ch, sess)
+				// Notify the client of the shell's exit status so it can
+				// disconnect cleanly instead of requiring a manual disconnect.
+				exitCode := sess.ExitCode()
+				exitStatusPayload := ssh.Marshal(struct{ Status uint32 }{uint32(exitCode)})
+				_, _ = ch.SendRequest("exit-status", false, exitStatusPayload)
 				_ = ch.CloseWrite()
 				sess.Close() //nolint:errcheck
+				// Close the channel so the ssh library closes the requests
+				// channel, which unblocks the for-range loop in handleSession
+				// and allows the deferred ch.Close() to run. Without this,
+				// handleSession blocks forever waiting for requests to drain.
+				_ = ch.Close()
 			}()
 			// SSH channel input → PTY stdin.
 			go func() {