scobleclaw/pkg/floxenv/floxenv.go

// Package floxenv manages Flox environment lifecycles via flox activate --detach.
// When the flox binary is not available, it falls back to a dev-mode stub suitable
// for local development and testing.
package floxenv

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

var (
	// ErrNotFound is returned when flox is not installed and dev mode is disabled.
	ErrNotFound = errors.New("flox command not found")
	// ErrAlreadyRunning is returned if an env with the same manifest is already active.
	ErrAlreadyRunning = errors.New("environment already active")
	// ErrActivationFailed is returned when flox activate exits non-zero.
	ErrActivationFailed = errors.New("flox activation failed")
	// ErrTimeout is returned when graceful shutdown exceeds the deadline.
	ErrTimeout = errors.New("environment shutdown timed out")
)

const (
	defaultFloxPath = "flox"
	devModeEnvVar   = "FLOXENV_DEV_MODE"
	gracePeriod     = 10 * time.Second
)

// Registry tracks active environment metadata by envID.
type Registry struct {
	mu   sync.RWMutex
	m    map[string]Info
	stub bool // true when flox binary is absent and dev mode is enabled
}

// Info holds runtime details for an activated environment.
type Info struct {
	EnvID       string
	PID         int
	ManifestPath string
	StartedAt   time.Time
}

// NewRegistry creates a new environment registry.
// It detects whether flox is available and enables dev-mode stub behaviour if not.
func NewRegistry() *Registry {
	r := &Registry{m: make(map[string]Info)}
	_, err := exec.LookPath(defaultFloxPath)
	if err != nil || os.Getenv(devModeEnvVar) == "1" {
		r.stub = true
	}
	return r
}

// IsDevMode returns true when the registry is running in stub/dev mode.
func (r *Registry) IsDevMode() bool {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return r.stub
}

// Activate starts a Flox environment in detached mode.
// manifestPath should point to a valid manifest.toml or flox.nix.
// On success it returns an envID (directory name + timestamp hash) and the PID of the
// detached activation process.
func (r *Registry) Activate(manifestPath string) (string, int, error) {
	absPath, err := filepath.Abs(manifestPath)
	if err != nil {
		return "", 0, fmt.Errorf("resolve manifest path: %w", err)
	}

	envID := generateEnvID(absPath)

	r.mu.Lock()
	if existing, ok := r.m[envID]; ok && IsAlive(existing.PID) {
		r.mu.Unlock()
		return "", 0, fmt.Errorf("%w: envID=%s pid=%d", ErrAlreadyRunning, envID, existing.PID)
	}
	r.mu.Unlock()

	if r.IsDevMode() {
		return r.stubActivate(envID, absPath)
	}

	// Run:  flox activate --detach -d <manifestDir>
	manifestDir := filepath.Dir(absPath)
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	cmd := exec.CommandContext(ctx, defaultFloxPath, "activate", "--detach", "-d", manifestDir)
	var out bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		return "", 0, fmt.Errorf("%w: %w (stderr: %s)", ErrActivationFailed, err, strings.TrimSpace(stderr.String()))
	}

	pid, err := parsePID(out.String())
	if err != nil {
		return "", 0, fmt.Errorf("parse activation pid: %w (output: %s)", err, strings.TrimSpace(out.String()))
	}

	info := Info{
		EnvID:        envID,
		PID:          pid,
		ManifestPath: absPath,
		StartedAt:    time.Now(),
	}

	r.mu.Lock()
	r.m[envID] = info
	r.mu.Unlock()

	return envID, pid, nil
}

// stubActivate is the dev-mode fallback when flox is unavailable.
// It forks a lightweight long-running shell command and records the PID.
func (r *Registry) stubActivate(envID, manifestPath string) (string, int, error) {
	cmd := exec.Command("sh", "-c", "sleep 3600")
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true, // avoid dying with the parent
	}
	if err := cmd.Start(); err != nil {
		return "", 0, fmt.Errorf("stub activate failed: %w", err)
	}

	info := Info{
		EnvID:        envID,
		PID:          cmd.Process.Pid,
		ManifestPath: manifestPath,
		StartedAt:    time.Now(),
	}

	r.mu.Lock()
	r.m[envID] = info
	r.mu.Unlock()

	return envID, cmd.Process.Pid, nil
}

// Deactivate gracefully shuts down an active Flox environment.
// It sends SIGTERM, waits up to 10 seconds, then SIGKILL if still alive.
func (r *Registry) Deactivate(envID string, pid int) error {
	r.mu.RLock()
	info, ok := r.m[envID]
	r.mu.RUnlock()

	if !ok {
		// Allow deactivation of PIDs we didn't start (defensive).
		info = Info{EnvID: envID, PID: pid}
	}

	if info.PID != pid {
		return fmt.Errorf("pid mismatch: registry has %d, caller provided %d", info.PID, pid)
	}

	if !IsAlive(pid) {
		// Already dead, clean-up registry anyway.
		r.mu.Lock()
		delete(r.m, envID)
		r.mu.Unlock()
		return nil
	}

	// Phase 1: SIGTERM (syscall.Kill avoids os.FindProcess Linux quirk)
	if err := syscall.Kill(pid, syscall.SIGTERM); err != nil {
		return fmt.Errorf("send SIGTERM to %d: %w", pid, err)
	}

	deadline := time.After(gracePeriod)
	ticker := time.NewTicker(200 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-deadline:
			// Phase 2: SIGKILL
			_ = syscall.Kill(pid, syscall.SIGKILL)
			// Wait a hair for the kernel to reap.
			time.Sleep(100 * time.Millisecond)
			if IsAlive(pid) {
				return fmt.Errorf("%w: envID=%s pid=%d", ErrTimeout, envID, pid)
			}
			r.mu.Lock()
			delete(r.m, envID)
			r.mu.Unlock()
			return nil
		case <-ticker.C:
			if !IsAlive(pid) {
				r.mu.Lock()
				delete(r.m, envID)
				r.mu.Unlock()
				return nil
			}
		}
	}
}

// List returns a snapshot of currently tracked environments.
func (r *Registry) List() []Info {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]Info, 0, len(r.m))
	for _, v := range r.m {
		out = append(out, v)
	}
	return out
}

// generateEnvID creates a deterministic identifier for a manifest path.
func generateEnvID(path string) string {
	dir := filepath.Dir(path)
	base := filepath.Base(dir)
	t := strconv.FormatInt(time.Now().Unix(), 36)
	return fmt.Sprintf("%s-%s", base, t)
}

// parsePID extracts an integer PID from flox activate --detach output.
// Flox itself does not always print the PID, so fall back to the shell
// process spawned by this package (handled separately).  This helper
// attempts to find a PID in the output and returns an error if none found.
func parsePID(s string) (int, error) {
	f := strings.Fields(s)
	for _, tok := range f {
		if n, err := strconv.Atoi(tok); err == nil && n > 0 {
			return n, nil
		}
	}
	return 0, errors.New("no PID found in output")
}