feat(muxterm): PTY pool with attach, resize, cleanup

2026-05-28 06:09:39 +00:00
parent ecc6c6979c
commit 3f0e31007f
5 changed files with 363 additions and 5 deletions
@@ -2,7 +2,4 @@ module github.com/muxplex/muxterm

 go 1.22

-require (
-	github.com/creack/pty v1.1.24
-	github.com/gorilla/websocket v1.5.3
-)
+require github.com/creack/pty v1.1.24
@@ -1,2 +1,2 @@
+github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s=
 github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE=
-github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
@@ -0,0 +1,129 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"sync"
+	"syscall"
+
+	"github.com/creack/pty"
+)
+
+// Session represents a PTY-backed tmux attach process.
+type Session struct {
+	PTY  *os.File
+	Cmd  *exec.Cmd
+	Cols uint16
+	Rows uint16
+}
+
+// Pool manages named tmux PTY sessions.
+type Pool struct {
+	mu       sync.Mutex
+	sessions map[string]*Session
+}
+
+// NewPool returns an empty Pool.
+func NewPool() *Pool {
+	return &Pool{
+		sessions: make(map[string]*Session),
+	}
+}
+
+// Attach reuses an existing live session or spawns a new tmux attach process.
+// Dead sessions are cleaned up before spawning a replacement.
+func (p *Pool) Attach(name string, cols, rows uint16) (*Session, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if sess, ok := p.sessions[name]; ok {
+		// Check if process is still alive via Signal(0).
+		if sess.Cmd.Process != nil && sess.Cmd.Process.Signal(syscall.Signal(0)) == nil {
+			return sess, nil
+		}
+		// Dead session — clean up.
+		sess.PTY.Close()
+		delete(p.sessions, name)
+	}
+
+	cmd := exec.Command("tmux", "attach", "-t", name)
+	winSize := &pty.Winsize{Cols: cols, Rows: rows}
+
+	ptmx, err := pty.StartWithSize(cmd, winSize)
+	if err != nil {
+		return nil, fmt.Errorf("pty start: %w", err)
+	}
+
+	sess := &Session{
+		PTY:  ptmx,
+		Cmd:  cmd,
+		Cols: cols,
+		Rows: rows,
+	}
+	p.sessions[name] = sess
+
+	// Background goroutine: wait for process exit and clean up.
+	go func() {
+		cmd.Wait()
+		p.mu.Lock()
+		defer p.mu.Unlock()
+		// Only delete if the map still holds this exact session
+		// (Attach may have replaced it already).
+		if cur, ok := p.sessions[name]; ok && cur == sess {
+			sess.PTY.Close()
+			delete(p.sessions, name)
+		}
+	}()
+
+	return sess, nil
+}
+
+// Resize changes the PTY window size for a named session.
+func (p *Pool) Resize(name string, cols, rows uint16) error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	sess, ok := p.sessions[name]
+	if !ok {
+		return fmt.Errorf("session %q not found", name)
+	}
+
+	if err := pty.Setsize(sess.PTY, &pty.Winsize{Cols: cols, Rows: rows}); err != nil {
+		return fmt.Errorf("setsize: %w", err)
+	}
+	sess.Cols = cols
+	sess.Rows = rows
+	return nil
+}
+
+// Get returns the session for name, or nil if not present.
+func (p *Pool) Get(name string) *Session {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.sessions[name]
+}
+
+// IsAlive returns true if the named session exists and its process is alive.
+func (p *Pool) IsAlive(name string) bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	sess, ok := p.sessions[name]
+	if !ok {
+		return false
+	}
+	return sess.Cmd.Process != nil && sess.Cmd.Process.Signal(syscall.Signal(0)) == nil
+}
+
+// CloseAll sends SIGTERM to all sessions and closes their PTYs.
+func (p *Pool) CloseAll() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	for name, sess := range p.sessions {
+		if sess.Cmd.Process != nil {
+			sess.Cmd.Process.Signal(syscall.SIGTERM)
+		}
+		sess.PTY.Close()
+		delete(p.sessions, name)
+	}
+}
@@ -0,0 +1,224 @@
+package main
+
+import (
+	"os/exec"
+	"testing"
+	"time"
+)
+
+func TestNewPool(t *testing.T) {
+	p := NewPool()
+	if p == nil {
+		t.Fatal("NewPool returned nil")
+	}
+}
+
+func TestGetEmpty(t *testing.T) {
+	p := NewPool()
+	s := p.Get("nonexistent")
+	if s != nil {
+		t.Fatal("Get on empty pool should return nil")
+	}
+}
+
+func TestIsAliveEmpty(t *testing.T) {
+	p := NewPool()
+	if p.IsAlive("nonexistent") {
+		t.Fatal("IsAlive on empty pool should return false")
+	}
+}
+
+// setupTmuxSession creates a tmux session for testing. Returns a cleanup function.
+func setupTmuxSession(t *testing.T, name string) func() {
+	t.Helper()
+	// Kill any leftover session with this name
+	exec.Command("tmux", "kill-session", "-t", name).Run()
+	cmd := exec.Command("tmux", "new-session", "-d", "-s", name, "-x", "80", "-y", "24")
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("failed to create tmux session %q: %v", name, err)
+	}
+	return func() {
+		exec.Command("tmux", "kill-session", "-t", name).Run()
+	}
+}
+
+func TestAttachCreatesSession(t *testing.T) {
+	name := "test-pool-attach"
+	cleanup := setupTmuxSession(t, name)
+	defer cleanup()
+
+	p := NewPool()
+	defer p.CloseAll()
+
+	sess, err := p.Attach(name, 80, 24)
+	if err != nil {
+		t.Fatalf("Attach failed: %v", err)
+	}
+	if sess == nil {
+		t.Fatal("Attach returned nil session")
+	}
+	if sess.PTY == nil {
+		t.Fatal("Session PTY is nil")
+	}
+	if sess.Cmd == nil {
+		t.Fatal("Session Cmd is nil")
+	}
+	if sess.Cols != 80 {
+		t.Errorf("Cols = %d, want 80", sess.Cols)
+	}
+	if sess.Rows != 24 {
+		t.Errorf("Rows = %d, want 24", sess.Rows)
+	}
+}
+
+func TestGetReturnsSession(t *testing.T) {
+	name := "test-pool-get"
+	cleanup := setupTmuxSession(t, name)
+	defer cleanup()
+
+	p := NewPool()
+	defer p.CloseAll()
+
+	_, err := p.Attach(name, 80, 24)
+	if err != nil {
+		t.Fatalf("Attach failed: %v", err)
+	}
+
+	sess := p.Get(name)
+	if sess == nil {
+		t.Fatal("Get returned nil after Attach")
+	}
+}
+
+func TestIsAliveAfterAttach(t *testing.T) {
+	name := "test-pool-alive"
+	cleanup := setupTmuxSession(t, name)
+	defer cleanup()
+
+	p := NewPool()
+	defer p.CloseAll()
+
+	_, err := p.Attach(name, 80, 24)
+	if err != nil {
+		t.Fatalf("Attach failed: %v", err)
+	}
+
+	if !p.IsAlive(name) {
+		t.Fatal("IsAlive should return true after Attach")
+	}
+}
+
+func TestAttachReusesLiveSession(t *testing.T) {
+	name := "test-pool-reuse"
+	cleanup := setupTmuxSession(t, name)
+	defer cleanup()
+
+	p := NewPool()
+	defer p.CloseAll()
+
+	sess1, err := p.Attach(name, 80, 24)
+	if err != nil {
+		t.Fatalf("First Attach failed: %v", err)
+	}
+
+	sess2, err := p.Attach(name, 80, 24)
+	if err != nil {
+		t.Fatalf("Second Attach failed: %v", err)
+	}
+
+	if sess1 != sess2 {
+		t.Fatal("Attach should reuse existing live session")
+	}
+}
+
+func TestResize(t *testing.T) {
+	name := "test-pool-resize"
+	cleanup := setupTmuxSession(t, name)
+	defer cleanup()
+
+	p := NewPool()
+	defer p.CloseAll()
+
+	_, err := p.Attach(name, 80, 24)
+	if err != nil {
+		t.Fatalf("Attach failed: %v", err)
+	}
+
+	err = p.Resize(name, 120, 40)
+	if err != nil {
+		t.Fatalf("Resize failed: %v", err)
+	}
+
+	sess := p.Get(name)
+	if sess.Cols != 120 {
+		t.Errorf("After resize Cols = %d, want 120", sess.Cols)
+	}
+	if sess.Rows != 40 {
+		t.Errorf("After resize Rows = %d, want 40", sess.Rows)
+	}
+}
+
+func TestResizeNonexistent(t *testing.T) {
+	p := NewPool()
+	err := p.Resize("nonexistent", 80, 24)
+	if err == nil {
+		t.Fatal("Resize on nonexistent session should return error")
+	}
+}
+
+func TestCloseAll(t *testing.T) {
+	name := "test-pool-closeall"
+	cleanup := setupTmuxSession(t, name)
+	defer cleanup()
+
+	p := NewPool()
+
+	_, err := p.Attach(name, 80, 24)
+	if err != nil {
+		t.Fatalf("Attach failed: %v", err)
+	}
+
+	p.CloseAll()
+
+	// Give the background goroutine time to clean up
+	time.Sleep(100 * time.Millisecond)
+
+	if p.IsAlive(name) {
+		t.Fatal("IsAlive should return false after CloseAll")
+	}
+}
+
+func TestAttachCleansUpDeadSession(t *testing.T) {
+	name := "test-pool-dead-cleanup"
+	cleanup := setupTmuxSession(t, name)
+	defer cleanup()
+
+	p := NewPool()
+	defer p.CloseAll()
+
+	sess1, err := p.Attach(name, 80, 24)
+	if err != nil {
+		t.Fatalf("First Attach failed: %v", err)
+	}
+
+	// Kill the process to simulate a dead session
+	sess1.Cmd.Process.Kill()
+	sess1.Cmd.Process.Wait()
+	time.Sleep(100 * time.Millisecond)
+
+	// Re-create the tmux session since killing the attach killed our handle
+	exec.Command("tmux", "kill-session", "-t", name).Run()
+	setupCmd := exec.Command("tmux", "new-session", "-d", "-s", name, "-x", "80", "-y", "24")
+	if err := setupCmd.Run(); err != nil {
+		t.Fatalf("failed to recreate tmux session: %v", err)
+	}
+
+	sess2, err := p.Attach(name, 80, 24)
+	if err != nil {
+		t.Fatalf("Second Attach after dead session failed: %v", err)
+	}
+
+	if sess1 == sess2 {
+		t.Fatal("Attach should create new session when old one is dead")
+	}
+}
@@ -0,0 +1,8 @@
+package main
+
+import "net/http"
+
+// HandleWebSocket is a placeholder that will be replaced in a later task.
+func HandleWebSocket(pool *Pool, secret string, w http.ResponseWriter, r *http.Request) {
+	http.Error(w, "not implemented", http.StatusNotImplemented)
+}