claude-web-chat/pty_session.go

package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"sync"
	"time"

	"github.com/creack/pty"
)


const (
	outputBufferSize = 128 * 1024 // 128KB ring buffer for replay
	ptyIdleTimeout   = 60 * time.Minute
	ptyCleanupTick   = 5 * time.Minute
)

// RingBuffer is a fixed-size circular buffer for terminal output.
type RingBuffer struct {
	data []byte
	size int
	pos  int
	full bool
	mu   sync.Mutex
}

func NewRingBuffer(size int) *RingBuffer {
	return &RingBuffer{data: make([]byte, size), size: size}
}

func (rb *RingBuffer) Write(p []byte) {
	rb.mu.Lock()
	defer rb.mu.Unlock()
	for _, b := range p {
		rb.data[rb.pos] = b
		rb.pos++
		if rb.pos >= rb.size {
			rb.pos = 0
			rb.full = true
		}
	}
}

func (rb *RingBuffer) Bytes() []byte {
	rb.mu.Lock()
	defer rb.mu.Unlock()
	if !rb.full {
		result := make([]byte, rb.pos)
		copy(result, rb.data[:rb.pos])
		return result
	}
	result := make([]byte, rb.size)
	n := copy(result, rb.data[rb.pos:])
	copy(result[n:], rb.data[:rb.pos])
	return result
}

// PTYSession manages a single Claude CLI running in a pseudo-terminal.
type PTYSession struct {
	ID         string
	ProjectDir string
	Ptmx       *os.File
	Cmd        *exec.Cmd

	buffer      *RingBuffer
	subscribers map[string]chan []byte
	mu          sync.Mutex
	done        chan struct{}
	lastActive  time.Time
}

// SpawnPTY starts a new Claude CLI in a pseudo-terminal.
func SpawnPTY(projectDir string) (*PTYSession, error) {
	cmd := exec.Command("claude")
	cmd.Dir = projectDir

	// Filter env vars to prevent nested session detection
	env := filterEnvMulti(os.Environ(), []string{
		"CLAUDECODE",
		"CLAUDE_CODE_ENTRYPOINT",
		"TERM",
		"COLORTERM",
	})
	// Set proper terminal type for full color and cursor support
	env = append(env, "TERM=xterm-256color", "COLORTERM=truecolor")
	cmd.Env = env

	ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Rows: 24, Cols: 80})
	if err != nil {
		return nil, fmt.Errorf("start pty: %w", err)
	}

	sess := &PTYSession{
		ID:          fmt.Sprintf("%d", time.Now().UnixNano()),
		ProjectDir:  projectDir,
		Ptmx:        ptmx,
		Cmd:         cmd,
		buffer:      NewRingBuffer(outputBufferSize),
		subscribers: make(map[string]chan []byte),
		done:        make(chan struct{}),
		lastActive:  time.Now(),
	}

	go sess.readLoop()
	go sess.waitExit()

	return sess, nil
}

// readLoop reads PTY output, writes to ring buffer, and forwards to subscribers.
func (s *PTYSession) readLoop() {
	buf := make([]byte, 4096)
	for {
		n, err := s.Ptmx.Read(buf)
		if err != nil {
			return
		}
		if n == 0 {
			continue
		}
		data := make([]byte, n)
		copy(data, buf[:n])

		s.buffer.Write(data)

		s.mu.Lock()
		s.lastActive = time.Now()
		for id, ch := range s.subscribers {
			select {
			case ch <- data:
			default:
				log.Printf("PTY subscriber %s slow, dropping data", id)
			}
		}
		s.mu.Unlock()
	}
}

// waitExit waits for the CLI process to exit and signals done.
func (s *PTYSession) waitExit() {
	if s.Cmd.Process != nil {
		s.Cmd.Wait()
	}
	close(s.done)
}

// Subscribe adds a subscriber for PTY output.
func (s *PTYSession) Subscribe(id string) chan []byte {
	s.mu.Lock()
	defer s.mu.Unlock()
	ch := make(chan []byte, 256)
	s.subscribers[id] = ch
	return ch
}

// Unsubscribe removes a subscriber.
func (s *PTYSession) Unsubscribe(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if ch, ok := s.subscribers[id]; ok {
		close(ch)
		delete(s.subscribers, id)
	}
}

// Resize changes the PTY terminal size.
func (s *PTYSession) Resize(rows, cols uint16) error {
	return pty.Setsize(s.Ptmx, &pty.Winsize{Rows: rows, Cols: cols})
}

// WriteInput sends keyboard input to the PTY.
func (s *PTYSession) WriteInput(data []byte) (int, error) {
	s.mu.Lock()
	s.lastActive = time.Now()
	s.mu.Unlock()
	return s.Ptmx.Write(data)
}

// GetBuffer returns the ring buffer contents for replay.
func (s *PTYSession) GetBuffer() []byte {
	return s.buffer.Bytes()
}

// Done returns a channel that closes when the process exits.
func (s *PTYSession) Done() <-chan struct{} {
	return s.done
}

// Close terminates the PTY session.
func (s *PTYSession) Close() {
	s.mu.Lock()
	for id, ch := range s.subscribers {
		close(ch)
		delete(s.subscribers, id)
	}
	s.mu.Unlock()

	s.Ptmx.Close()
	if s.Cmd.Process != nil {
		s.Cmd.Process.Kill()
	}
}

// PTYSessionManager manages PTY sessions keyed by project name.
type PTYSessionManager struct {
	sessions map[string]*PTYSession
	mu       sync.RWMutex
	stopCh   chan struct{}
}

func NewPTYSessionManager() *PTYSessionManager {
	m := &PTYSessionManager{
		sessions: make(map[string]*PTYSession),
		stopCh:   make(chan struct{}),
	}
	go m.cleanup()
	return m
}

// GetOrCreate returns an existing session or creates a new one.
func (m *PTYSessionManager) GetOrCreate(project, projectDir string) (*PTYSession, bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if sess, ok := m.sessions[project]; ok {
		// Check if process is still alive
		select {
		case <-sess.Done():
			// Process exited, remove and create new
			sess.Close()
			delete(m.sessions, project)
		default:
			sess.mu.Lock()
			sess.lastActive = time.Now()
			sess.mu.Unlock()
			return sess, false, nil
		}
	}

	sess, err := SpawnPTY(projectDir)
	if err != nil {
		return nil, false, err
	}

	m.sessions[project] = sess
	return sess, true, nil
}

// Remove terminates and removes a session.
func (m *PTYSessionManager) Remove(project string) {
	m.mu.Lock()
	sess, ok := m.sessions[project]
	if ok {
		delete(m.sessions, project)
	}
	m.mu.Unlock()

	if ok {
		sess.Close()
	}
}

// Stop shuts down the manager and all sessions.
func (m *PTYSessionManager) Stop() {
	close(m.stopCh)

	m.mu.Lock()
	for id, sess := range m.sessions {
		sess.Close()
		delete(m.sessions, id)
	}
	m.mu.Unlock()
}

func (m *PTYSessionManager) cleanup() {
	ticker := time.NewTicker(ptyCleanupTick)
	defer ticker.Stop()

	for {
		select {
		case <-m.stopCh:
			return
		case <-ticker.C:
			m.cleanupIdle()
		}
	}
}

func (m *PTYSessionManager) cleanupIdle() {
	m.mu.Lock()
	var toRemove []string
	for id, sess := range m.sessions {
		sess.mu.Lock()
		idle := time.Since(sess.lastActive) > ptyIdleTimeout
		sess.mu.Unlock()

		// Also check if process has exited
		exited := false
		select {
		case <-sess.Done():
			exited = true
		default:
		}

		if idle || exited {
			toRemove = append(toRemove, id)
		}
	}
	m.mu.Unlock()

	for _, id := range toRemove {
		log.Printf("Cleaning up PTY session: %s", id)
		m.Remove(id)
	}
}