From cafed85c2e85e8f8d8213e75fe44fefdae662378 Mon Sep 17 00:00:00 2001
From: bagas <bagas@fossy.my.id>
Date: Tue, 10 Mar 2026 20:47:52 +0700
Subject: [PATCH] fix: capture TERM from PTY request to restore terminal colors
 on SSH sessions

---
 session/interaction/interaction.go | 45 ++++++++++++++++++++++++------
 session/session.go                 | 23 ++++++++++++++-
 2 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/session/interaction/interaction.go b/session/interaction/interaction.go
index fa39d3c..70c9319 100644
--- a/session/interaction/interaction.go
+++ b/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,34 @@ 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
+	latestW, latestH := i.width, i.height
 	i.programMu.Unlock()
 
-	_, err := i.program.Run()
+	if latestW > 0 && latestH > 0 {
+		prog.Send(tea.WindowSizeMsg{Width: latestW, Height: latestH})
+	}
+
+	_, err := prog.Run()
 	if err != nil {
 		log.Printf("Cannot close tea: %s \n", err)
 	}
diff --git a/session/session.go b/session/session.go
index cc27c4c..2062d6a 100644
--- a/session/session.go
+++ b/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