go-llm/v2/sandbox/sandbox.go

package sandbox

import (
	"context"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
)

// Config holds all configuration for creating sandboxes.
type Config struct {
	Proxmox  ProxmoxConfig
	SSH      SSHConfig
	Defaults ContainerConfig
}

// Option configures a Sandbox before creation.
type Option func(*createOpts)

type createOpts struct {
	hostname string
	cpus     int
	memoryMB int
	diskGB   int
	internet bool
}

// WithHostname sets the container hostname.
func WithHostname(name string) Option {
	return func(o *createOpts) { o.hostname = name }
}

// WithCPUs sets the number of CPU cores for the container.
func WithCPUs(n int) Option {
	return func(o *createOpts) { o.cpus = n }
}

// WithMemoryMB sets the memory limit in megabytes.
func WithMemoryMB(mb int) Option {
	return func(o *createOpts) { o.memoryMB = mb }
}

// WithDiskGB sets the root filesystem size in gigabytes.
func WithDiskGB(gb int) Option {
	return func(o *createOpts) { o.diskGB = gb }
}

// WithInternet enables outbound HTTP/HTTPS access on creation.
func WithInternet(enabled bool) Option {
	return func(o *createOpts) { o.internet = enabled }
}

// Sandbox represents an isolated Linux container environment with SSH access.
// It wraps a Proxmox LXC container and provides command execution and file operations.
type Sandbox struct {
	// ID is the Proxmox VMID of this container.
	ID int

	// IP is the container's IP address on the isolated bridge.
	IP string

	// Internet indicates whether outbound HTTP/HTTPS is enabled.
	Internet bool

	proxmox *ProxmoxClient
	ssh     *SSHExecutor
}

// Manager creates and manages sandbox instances.
type Manager struct {
	proxmox  *ProxmoxClient
	sshKey   ssh.Signer
	defaults ContainerConfig
	sshCfg   SSHConfig
}

// NewManager creates a new sandbox manager from the given configuration.
func NewManager(cfg Config) (*Manager, error) {
	if cfg.SSH.Signer == nil {
		return nil, fmt.Errorf("SSH signer is required")
	}

	return &Manager{
		proxmox:  NewProxmoxClient(cfg.Proxmox),
		sshKey:   cfg.SSH.Signer,
		defaults: cfg.Defaults,
		sshCfg:   cfg.SSH,
	}, nil
}

// Create provisions a new sandbox container: clones the template, starts it,
// waits for SSH, and optionally enables internet access.
// The returned Sandbox must be destroyed with Destroy when no longer needed.
func (m *Manager) Create(ctx context.Context, opts ...Option) (*Sandbox, error) {
	o := &createOpts{
		hostname: m.defaults.Hostname,
		cpus:     m.defaults.CPUs,
		memoryMB: m.defaults.MemoryMB,
		diskGB:   m.defaults.DiskGB,
	}
	for _, opt := range opts {
		opt(o)
	}

	// Apply defaults for zero values.
	if o.cpus <= 0 {
		o.cpus = 1
	}
	if o.memoryMB <= 0 {
		o.memoryMB = 1024
	}
	if o.diskGB <= 0 {
		o.diskGB = 8
	}

	// Get next VMID.
	vmid, err := m.proxmox.NextAvailableID(ctx)
	if err != nil {
		return nil, fmt.Errorf("get next VMID: %w", err)
	}

	containerCfg := ContainerConfig{
		Hostname: o.hostname,
		CPUs:     o.cpus,
		MemoryMB: o.memoryMB,
		DiskGB:   o.diskGB,
	}

	// Clone template.
	if err := m.proxmox.CloneTemplate(ctx, vmid, containerCfg); err != nil {
		return nil, fmt.Errorf("clone template: %w", err)
	}

	// Configure container resources.
	if err := m.proxmox.ConfigureContainer(ctx, vmid, containerCfg); err != nil {
		// Clean up the cloned container on failure.
		_ = m.proxmox.DestroyContainer(ctx, vmid)
		return nil, fmt.Errorf("configure container: %w", err)
	}

	// Start container.
	if err := m.proxmox.StartContainer(ctx, vmid); err != nil {
		_ = m.proxmox.DestroyContainer(ctx, vmid)
		return nil, fmt.Errorf("start container: %w", err)
	}

	// Discover IP address (with timeout).
	ipCtx, ipCancel := context.WithTimeout(ctx, 30*time.Second)
	defer ipCancel()

	ip, err := m.proxmox.GetContainerIP(ipCtx, vmid)
	if err != nil {
		_ = m.proxmox.DestroyContainer(ctx, vmid)
		return nil, fmt.Errorf("discover IP: %w", err)
	}

	// Connect SSH (with timeout).
	sshExec := NewSSHExecutor(ip, m.sshCfg)

	sshCtx, sshCancel := context.WithTimeout(ctx, 30*time.Second)
	defer sshCancel()

	if err := sshExec.Connect(sshCtx); err != nil {
		_ = m.proxmox.DestroyContainer(ctx, vmid)
		return nil, fmt.Errorf("ssh connect: %w", err)
	}

	sb := &Sandbox{
		ID:      vmid,
		IP:      ip,
		proxmox: m.proxmox,
		ssh:     sshExec,
	}

	// Enable internet if requested.
	if o.internet {
		if err := sb.SetInternet(ctx, true); err != nil {
			sb.Destroy(ctx)
			return nil, fmt.Errorf("enable internet: %w", err)
		}
	}

	return sb, nil
}

// Attach reconnects to an existing sandbox container by VMID.
// This is useful for recovering sessions after a restart.
func (m *Manager) Attach(ctx context.Context, vmid int) (*Sandbox, error) {
	status, err := m.proxmox.GetContainerStatus(ctx, vmid)
	if err != nil {
		return nil, fmt.Errorf("get container status: %w", err)
	}
	if status.Status != "running" {
		return nil, fmt.Errorf("container %d is not running (status: %s)", vmid, status.Status)
	}

	ip, err := m.proxmox.GetContainerIP(ctx, vmid)
	if err != nil {
		return nil, fmt.Errorf("get container IP: %w", err)
	}

	sshExec := NewSSHExecutor(ip, m.sshCfg)
	if err := sshExec.Connect(ctx); err != nil {
		return nil, fmt.Errorf("ssh connect: %w", err)
	}

	return &Sandbox{
		ID:      vmid,
		IP:      ip,
		proxmox: m.proxmox,
		ssh:     sshExec,
	}, nil
}

// Exec runs a shell command in the sandbox and returns the result.
func (s *Sandbox) Exec(ctx context.Context, command string) (ExecResult, error) {
	return s.ssh.Exec(ctx, command)
}

// WriteFile creates or overwrites a file in the sandbox.
func (s *Sandbox) WriteFile(ctx context.Context, path, content string) error {
	return s.ssh.Upload(ctx, strings.NewReader(content), path, 0644)
}

// ReadFile reads a file from the sandbox and returns its contents.
func (s *Sandbox) ReadFile(ctx context.Context, path string) (string, error) {
	rc, err := s.ssh.Download(ctx, path)
	if err != nil {
		return "", err
	}
	defer rc.Close()

	data, err := io.ReadAll(rc)
	if err != nil {
		return "", fmt.Errorf("read file %s: %w", path, err)
	}
	return string(data), nil
}

// Upload copies data from an io.Reader to a file in the sandbox.
func (s *Sandbox) Upload(ctx context.Context, reader io.Reader, remotePath string, mode os.FileMode) error {
	return s.ssh.Upload(ctx, reader, remotePath, mode)
}

// Download returns an io.ReadCloser for a file in the sandbox.
// The caller must close the returned reader.
func (s *Sandbox) Download(ctx context.Context, remotePath string) (io.ReadCloser, error) {
	return s.ssh.Download(ctx, remotePath)
}

// SetInternet enables or disables outbound HTTP/HTTPS access for the sandbox.
func (s *Sandbox) SetInternet(ctx context.Context, enabled bool) error {
	if enabled {
		if err := s.proxmox.EnableInternet(ctx, s.IP); err != nil {
			return err
		}
	} else {
		if err := s.proxmox.DisableInternet(ctx, s.IP); err != nil {
			return err
		}
	}
	s.Internet = enabled
	return nil
}

// Status returns the current resource usage of the sandbox container.
func (s *Sandbox) Status(ctx context.Context) (ContainerStatus, error) {
	return s.proxmox.GetContainerStatus(ctx, s.ID)
}

// IsConnected returns true if the SSH connection to the sandbox is active.
func (s *Sandbox) IsConnected() bool {
	return s.ssh.IsConnected()
}

// Destroy stops the container, removes internet access, closes SSH connections,
// and permanently deletes the container from Proxmox.
func (s *Sandbox) Destroy(ctx context.Context) error {
	var errs []error

	// Remove internet access first (ignore errors — container is being destroyed).
	if s.Internet {
		_ = s.proxmox.DisableInternet(ctx, s.IP)
	}

	// Close SSH connections.
	if err := s.ssh.Close(); err != nil {
		errs = append(errs, fmt.Errorf("close ssh: %w", err))
	}

	// Destroy the container.
	if err := s.proxmox.DestroyContainer(ctx, s.ID); err != nil {
		errs = append(errs, fmt.Errorf("destroy container: %w", err))
	}

	if len(errs) > 0 {
		return fmt.Errorf("destroy sandbox %d: %v", s.ID, errs)
	}
	return nil
}

// DestroyByID destroys a container by VMID without requiring an active SSH connection.
// This is useful for cleaning up orphaned containers after a restart.
func (m *Manager) DestroyByID(ctx context.Context, vmid int) error {
	return m.proxmox.DestroyContainer(ctx, vmid)
}