fix: capture TERM from PTY request to restore terminal colors on SSH sessions

session/interaction/interaction.go

@@ -53,6 +53,8 @@ type interaction struct {
 	cancel          context.CancelFunc
 	mode            types.InteractiveMode
 	programMu       sync.Mutex
+	width           int
+	height          int
 }
 
 func (i *interaction) SetMode(m types.InteractiveMode) {
@@ -71,8 +73,14 @@ func (i *interaction) Send(message string) error {
 	return nil
 }
 func (i *interaction) SetWH(w, h int) {
-	if i.program != nil {
-		i.program.Send(tea.WindowSizeMsg{
+	i.programMu.Lock()
+	i.width = w
+	i.height = h
+	prog := i.program
+	i.programMu.Unlock()
+
+	if prog != nil {
+		prog.Send(tea.WindowSizeMsg{
 			Width:  w,
 			Height: h,
 		})
@@ -158,8 +166,13 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 }
 
 func (i *interaction) Redraw() {
-	if i.program != nil {
-		i.program.Send(tea.WindowSizeMsg{})
+	i.programMu.Lock()
+	prog := i.program
+	w, h := i.width, i.height
+	i.programMu.Unlock()
+
+	if prog != nil {
+		prog.Send(tea.WindowSizeMsg{Width: w, Height: h})
 	}
 }
 
@@ -246,18 +259,29 @@ func (i *interaction) Start() {
 	}
 
 	i.programMu.Lock()
-	i.program = tea.NewProgram(
-		m,
+	w, h := i.width, i.height
+	i.programMu.Unlock()
+
+	opts := []tea.ProgramOption{
 		tea.WithInput(i.channel),
 		tea.WithOutput(i.channel),
 		tea.WithColorProfile(colorprofile.TrueColor),
 		tea.WithoutSignals(),
 		tea.WithoutSignalHandler(),
 		tea.WithFPS(30),
-	)
+		tea.WithEnvironment([]string{"TERM=xterm-256color", "COLORTERM=truecolor"}),
+	}
+	if w > 0 && h > 0 {
+		opts = append(opts, tea.WithWindowSize(w, h))
+	}
+
+	prog := tea.NewProgram(m, opts...)
+
+	i.programMu.Lock()
+	i.program = prog
 	i.programMu.Unlock()
 
-	_, err := i.program.Run()
+	_, err := prog.Run()
 	if err != nil {
 		log.Printf("Cannot close tea: %s \n", err)
 	}

session/session.go

@@ -232,13 +232,34 @@ func (s *session) handleWindowChange(req *ssh.Request) error {
 	return req.Reply(true, nil)
 }
 
+func (s *session) handlePtyReq(req *ssh.Request) error {
+	var ptyPayload struct {
+		Term   string
+		Width  uint32
+		Height uint32
+		PxW    uint32
+		PxH    uint32
+		Modes  string
+	}
+	if err := ssh.Unmarshal(req.Payload, &ptyPayload); err == nil {
+		if ptyPayload.Width > 0 && ptyPayload.Height > 0 {
+			s.interaction.SetWH(int(ptyPayload.Width), int(ptyPayload.Height))
+		}
+	}
+	return req.Reply(true, nil)
+}
+
 func (s *session) HandleGlobalRequest(GlobalRequest <-chan *ssh.Request) error {
 	for req := range GlobalRequest {
 		switch req.Type {
-		case "shell", "pty-req":
+		case "shell":
 			if err := req.Reply(true, nil); err != nil {
 				return err
 			}
+		case "pty-req":
+			if err := s.handlePtyReq(req); err != nil {
+				return err
+			}
 		case "window-change":
 			if err := s.handleWindowChange(req); err != nil {
 				return err