vibetunnel/linux/pkg/session/session.go
2025-06-20 05:30:13 +02:00

524 lines
12 KiB
Go

package session
import (
"encoding/json"
"fmt"
"log"
"os"
"path/filepath"
"strings"
"sync"
"syscall"
"time"
"github.com/google/uuid"
)
// GenerateID generates a new unique session ID
func GenerateID() string {
return uuid.New().String()
}
type Status string
const (
StatusStarting Status = "starting"
StatusRunning Status = "running"
StatusExited Status = "exited"
)
type Config struct {
Name string
Cmdline []string
Cwd string
Env []string
Width int
Height int
}
type Info struct {
ID string `json:"id"`
Name string `json:"name"`
Cmdline string `json:"cmdline"`
Cwd string `json:"cwd"`
Pid int `json:"pid,omitempty"`
Status string `json:"status"`
ExitCode *int `json:"exit_code,omitempty"`
StartedAt time.Time `json:"started_at"`
Term string `json:"term"`
Width int `json:"width"`
Height int `json:"height"`
Env map[string]string `json:"env,omitempty"`
Args []string `json:"-"` // Internal use only
}
type Session struct {
ID string
controlPath string
info *Info
pty *PTY
stdinPipe *os.File
stdinMutex sync.Mutex
mu sync.RWMutex
}
func newSession(controlPath string, config Config) (*Session, error) {
id := uuid.New().String()
return newSessionWithID(controlPath, id, config)
}
func newSessionWithID(controlPath string, id string, config Config) (*Session, error) {
sessionPath := filepath.Join(controlPath, id)
// Only log in debug mode
if os.Getenv("VIBETUNNEL_DEBUG") != "" {
log.Printf("[DEBUG] Creating new session %s with config: Name=%s, Cmdline=%v, Cwd=%s",
id[:8], config.Name, config.Cmdline, config.Cwd)
}
if err := os.MkdirAll(sessionPath, 0755); err != nil {
return nil, fmt.Errorf("failed to create session directory: %w", err)
}
if config.Name == "" {
config.Name = id[:8]
}
// Set default command if empty
if len(config.Cmdline) == 0 {
shell := os.Getenv("SHELL")
if shell == "" {
shell = "/bin/bash"
}
config.Cmdline = []string{shell}
if os.Getenv("VIBETUNNEL_DEBUG") != "" {
log.Printf("[DEBUG] Session %s: Set default command to %v", id[:8], config.Cmdline)
}
}
// Set default working directory if empty
if config.Cwd == "" {
cwd, err := os.Getwd()
if err != nil {
config.Cwd = os.Getenv("HOME")
if config.Cwd == "" {
config.Cwd = "/"
}
} else {
config.Cwd = cwd
}
if os.Getenv("VIBETUNNEL_DEBUG") != "" {
log.Printf("[DEBUG] Session %s: Set default working directory to %s", id[:8], config.Cwd)
}
}
term := os.Getenv("TERM")
if term == "" {
term = "xterm-256color"
}
// Set default terminal dimensions if not provided
width := config.Width
if width <= 0 {
width = 120 // Better default for modern terminals
}
height := config.Height
if height <= 0 {
height = 30 // Better default for modern terminals
}
info := &Info{
ID: id,
Name: config.Name,
Cmdline: strings.Join(config.Cmdline, " "),
Cwd: config.Cwd,
Status: string(StatusStarting),
StartedAt: time.Now(),
Term: term,
Width: width,
Height: height,
Args: config.Cmdline,
}
if err := info.Save(sessionPath); err != nil {
os.RemoveAll(sessionPath)
return nil, fmt.Errorf("failed to save session info: %w", err)
}
return &Session{
ID: id,
controlPath: controlPath,
info: info,
}, nil
}
func loadSession(controlPath, id string) (*Session, error) {
sessionPath := filepath.Join(controlPath, id)
info, err := LoadInfo(sessionPath)
if err != nil {
return nil, err
}
session := &Session{
ID: id,
controlPath: controlPath,
info: info,
}
// If session is running, we need to reconnect to the PTY for operations like resize
// For now, we'll handle this by checking if we need PTY access in individual methods
return session, nil
}
func (s *Session) Path() string {
return filepath.Join(s.controlPath, s.ID)
}
func (s *Session) StreamOutPath() string {
return filepath.Join(s.Path(), "stream-out")
}
func (s *Session) StdinPath() string {
return filepath.Join(s.Path(), "stdin")
}
func (s *Session) NotificationPath() string {
return filepath.Join(s.Path(), "notification-stream")
}
func (s *Session) Start() error {
pty, err := NewPTY(s)
if err != nil {
return fmt.Errorf("failed to create PTY: %w", err)
}
s.pty = pty
s.info.Status = string(StatusRunning)
s.info.Pid = pty.Pid()
if err := s.info.Save(s.Path()); err != nil {
pty.Close()
return fmt.Errorf("failed to update session info: %w", err)
}
go func() {
if err := s.pty.Run(); err != nil {
if os.Getenv("VIBETUNNEL_DEBUG") != "" {
log.Printf("[DEBUG] Session %s: PTY.Run() exited with error: %v", s.ID[:8], err)
}
} else {
if os.Getenv("VIBETUNNEL_DEBUG") != "" {
log.Printf("[DEBUG] Session %s: PTY.Run() exited normally", s.ID[:8])
}
}
}()
// Start control listener
s.startControlListener()
// Process status will be checked on first access - no artificial delay needed
if os.Getenv("VIBETUNNEL_DEBUG") != "" {
log.Printf("[DEBUG] Session %s: Started successfully", s.ID[:8])
}
return nil
}
func (s *Session) Attach() error {
if s.pty == nil {
return fmt.Errorf("session not started")
}
return s.pty.Attach()
}
func (s *Session) SendKey(key string) error {
return s.sendInput([]byte(key))
}
func (s *Session) SendText(text string) error {
return s.sendInput([]byte(text))
}
func (s *Session) sendInput(data []byte) error {
s.stdinMutex.Lock()
defer s.stdinMutex.Unlock()
// Open pipe if not already open
if s.stdinPipe == nil {
stdinPath := s.StdinPath()
pipe, err := os.OpenFile(stdinPath, os.O_WRONLY, 0)
if err != nil {
return fmt.Errorf("failed to open stdin pipe: %w", err)
}
s.stdinPipe = pipe
}
_, err := s.stdinPipe.Write(data)
if err != nil {
// If write fails, close and reset the pipe for next attempt
s.stdinPipe.Close()
s.stdinPipe = nil
return fmt.Errorf("failed to write to stdin pipe: %w", err)
}
return nil
}
func (s *Session) Signal(sig string) error {
if s.info.Pid == 0 {
return fmt.Errorf("no process to signal")
}
// Check if process is still alive before signaling
if !s.IsAlive() {
// Process is already dead, update status and return success
s.info.Status = string(StatusExited)
exitCode := 0
s.info.ExitCode = &exitCode
s.info.Save(s.Path())
return nil
}
proc, err := os.FindProcess(s.info.Pid)
if err != nil {
return err
}
switch sig {
case "SIGTERM":
return proc.Signal(os.Interrupt)
case "SIGKILL":
err = proc.Kill()
// If kill fails with "process already finished", that's okay
if err != nil && strings.Contains(err.Error(), "process already finished") {
return nil
}
return err
default:
return fmt.Errorf("unsupported signal: %s", sig)
}
}
func (s *Session) Stop() error {
return s.Signal("SIGTERM")
}
func (s *Session) Kill() error {
// First check if the session is already dead
if s.info.Status == string(StatusExited) {
// Already exited, just cleanup and return success
s.cleanup()
return nil
}
// Try to kill the process
err := s.Signal("SIGKILL")
s.cleanup()
// If the error is because the process doesn't exist, that's fine
if err != nil && (strings.Contains(err.Error(), "no such process") ||
strings.Contains(err.Error(), "process already finished")) {
return nil
}
return err
}
func (s *Session) cleanup() {
s.stdinMutex.Lock()
defer s.stdinMutex.Unlock()
if s.stdinPipe != nil {
s.stdinPipe.Close()
s.stdinPipe = nil
}
}
func (s *Session) Resize(width, height int) error {
if s.pty == nil {
return fmt.Errorf("session not started")
}
// Check if session is still alive
if s.info.Status == string(StatusExited) {
return fmt.Errorf("cannot resize exited session")
}
// Validate dimensions
if width <= 0 || height <= 0 {
return fmt.Errorf("invalid dimensions: width=%d, height=%d", width, height)
}
// Update session info
s.info.Width = width
s.info.Height = height
// Save updated session info
if err := s.info.Save(s.Path()); err != nil {
log.Printf("[ERROR] Failed to save session info after resize: %v", err)
}
// Resize the PTY
return s.pty.Resize(width, height)
}
func (s *Session) IsAlive() bool {
if s.info.Pid == 0 {
return false
}
// Check if process exists and is not a zombie
if isZombie(s.info.Pid) {
return false
}
proc, err := os.FindProcess(s.info.Pid)
if err != nil {
return false
}
err = proc.Signal(syscall.Signal(0))
return err == nil
}
// isZombie checks if a process is in zombie state
func isZombie(pid int) bool {
// Read process status from /proc/[pid]/stat
statPath := fmt.Sprintf("/proc/%d/stat", pid)
data, err := os.ReadFile(statPath)
if err != nil {
// If we can't read the stat file, process doesn't exist
return true
}
// The process state is the third field after the command name in parentheses
// Find the last ')' to handle processes with ')' in their names
statStr := string(data)
lastParen := strings.LastIndex(statStr, ")")
if lastParen == -1 {
return true
}
// Parse fields after the command name
fields := strings.Fields(statStr[lastParen+1:])
if len(fields) < 1 {
return true
}
// State is the first field after the command
// Z = zombie
state := fields[0]
return state == "Z"
}
func (s *Session) UpdateStatus() error {
if s.info.Status == string(StatusExited) {
return nil
}
if !s.IsAlive() {
s.info.Status = string(StatusExited)
exitCode := 0
s.info.ExitCode = &exitCode
return s.info.Save(s.Path())
}
return nil
}
// GetInfo returns the session info
func (s *Session) GetInfo() *Info {
s.mu.RLock()
defer s.mu.RUnlock()
return s.info
}
func (i *Info) Save(sessionPath string) error {
data, err := json.MarshalIndent(i, "", " ")
if err != nil {
return err
}
return os.WriteFile(filepath.Join(sessionPath, "session.json"), data, 0644)
}
// RustSessionInfo represents the session format used by the Rust server
type RustSessionInfo struct {
ID string `json:"id,omitempty"`
Name string `json:"name"`
Cmdline []string `json:"cmdline"`
Cwd string `json:"cwd"`
Pid *int `json:"pid,omitempty"`
Status string `json:"status"`
ExitCode *int `json:"exit_code,omitempty"`
StartedAt *time.Time `json:"started_at,omitempty"`
Term string `json:"term"`
SpawnType string `json:"spawn_type,omitempty"`
Cols *int `json:"cols,omitempty"`
Rows *int `json:"rows,omitempty"`
Env map[string]string `json:"env,omitempty"`
}
func LoadInfo(sessionPath string) (*Info, error) {
data, err := os.ReadFile(filepath.Join(sessionPath, "session.json"))
if err != nil {
return nil, err
}
// First try to unmarshal as Go format
var info Info
if err := json.Unmarshal(data, &info); err == nil {
// Successfully parsed as Go format
return &info, nil
}
// If that fails, try Rust format
var rustInfo RustSessionInfo
if err := json.Unmarshal(data, &rustInfo); err != nil {
return nil, fmt.Errorf("failed to parse session.json: %w", err)
}
// Convert Rust format to Go format
info = Info{
ID: rustInfo.ID,
Name: rustInfo.Name,
Cmdline: strings.Join(rustInfo.Cmdline, " "),
Cwd: rustInfo.Cwd,
Status: rustInfo.Status,
ExitCode: rustInfo.ExitCode,
Term: rustInfo.Term,
Args: rustInfo.Cmdline,
Env: rustInfo.Env,
}
// Handle PID conversion
if rustInfo.Pid != nil {
info.Pid = *rustInfo.Pid
}
// Handle dimensions: use cols/rows if available, otherwise defaults
if rustInfo.Cols != nil {
info.Width = *rustInfo.Cols
} else {
info.Width = 120
}
if rustInfo.Rows != nil {
info.Height = *rustInfo.Rows
} else {
info.Height = 30
}
// Handle timestamp
if rustInfo.StartedAt != nil {
info.StartedAt = *rustInfo.StartedAt
} else {
info.StartedAt = time.Now()
}
// If ID is empty (Rust doesn't store it in JSON), derive it from directory name
if info.ID == "" {
info.ID = filepath.Base(sessionPath)
}
return &info, nil
}