From 3f0e31007f39156e113783255e118e3f4d6efc95 Mon Sep 17 00:00:00 2001 From: Ken Date: Thu, 28 May 2026 06:09:39 +0000 Subject: [PATCH] feat(muxterm): PTY pool with attach, resize, cleanup --- muxterm/go.mod | 5 +- muxterm/go.sum | 2 +- muxterm/pool.go | 129 +++++++++++++++++++++++++ muxterm/pool_test.go | 224 +++++++++++++++++++++++++++++++++++++++++++ muxterm/ws_stub.go | 8 ++ 5 files changed, 363 insertions(+), 5 deletions(-) create mode 100644 muxterm/pool.go create mode 100644 muxterm/pool_test.go create mode 100644 muxterm/ws_stub.go diff --git a/muxterm/go.mod b/muxterm/go.mod index 63a2b05..09cf5f4 100644 --- a/muxterm/go.mod +++ b/muxterm/go.mod @@ -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 diff --git a/muxterm/go.sum b/muxterm/go.sum index 14d693a..5310a43 100644 --- a/muxterm/go.sum +++ b/muxterm/go.sum @@ -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= diff --git a/muxterm/pool.go b/muxterm/pool.go new file mode 100644 index 0000000..fc8c983 --- /dev/null +++ b/muxterm/pool.go @@ -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) + } +} diff --git a/muxterm/pool_test.go b/muxterm/pool_test.go new file mode 100644 index 0000000..6c5531f --- /dev/null +++ b/muxterm/pool_test.go @@ -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") + } +} diff --git a/muxterm/ws_stub.go b/muxterm/ws_stub.go new file mode 100644 index 0000000..7df95e3 --- /dev/null +++ b/muxterm/ws_stub.go @@ -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) +}