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

This commit is contained in:
Ken
2026-05-28 06:09:39 +00:00
parent ecc6c6979c
commit 3f0e31007f
5 changed files with 363 additions and 5 deletions
+1 -4
View File
@@ -2,7 +2,4 @@ module github.com/muxplex/muxterm
go 1.22 go 1.22
require ( require github.com/creack/pty v1.1.24
github.com/creack/pty v1.1.24
github.com/gorilla/websocket v1.5.3
)
+1 -1
View File
@@ -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/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE=
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+129
View File
@@ -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)
}
}
+224
View File
@@ -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")
}
}
+8
View File
@@ -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)
}