ollie/cmd/toolsrv/fs9p.go

// Package toolsrv provides a 9P filesystem interface to the tool execution server.
//
// Namespace:
//
//	/ctl           write: load <tool>, unload <tool>
//	/tools         read: list loaded tools (name\tdesc per line)
//	               write: tool name to load
//	/proc/         dir: running/completed processes
//	/proc/new      rdwr: write tool+args, blocks, read result (auto-cleanup)
//	/proc/new.bg   write: write tool+args, returns pid immediately (detached)
//	/proc/<pid>/   dir: process directory
//	/proc/<pid>/out   read: output stream
//	/proc/<pid>/wait  read: blocks until exit, returns exit code
//	/proc/<pid>/stat  read: status (running, exited N, runtime)
//	/proc/<pid>/ctl   write: signal <N>, dismiss
//	/info          read: host info (JSON)
//
// Tool call payload format: key=value lines, newline delimited.
// Example:
//
//	tool=shell
//	cmd=ls -la
//	timeout=30
package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"ollie/toolsrv"
)

// FS9P implements a 9P-style tool server.
// This is the core logic; the actual 9P serving is handled separately.
type FS9P struct {
	mu sync.RWMutex

	cwd      string
	env      map[string]string
	registry *Registry
	sessID   string
	yolo     bool

	// Process management
	procMu    sync.Mutex
	procs     map[int]*Proc9P
	nextPID   int
	procLimit int

	// Callbacks
	OnToolsChanged func()
}

// Proc9P represents a running or completed process.
type Proc9P struct {
	PID       int
	Tool      string
	Args      map[string]string
	StartTime time.Time
	EndTime   time.Time
	ExitCode  int
	Exited    bool

	mu     sync.Mutex
	output bytes.Buffer
	done   chan struct{}
	cancel context.CancelFunc
	proc   *os.Process // for signaling
}

// NewFS9P creates a new 9P tool server.
func NewFS9P(cwd string) *FS9P {
	return &FS9P{
		cwd:       cwd,
		env:       make(map[string]string),
		procs:     make(map[int]*Proc9P),
		nextPID:   1,
		procLimit: 32,
	}
}

// SetRegistry configures the tool registry and session ID.
func (fs *FS9P) SetRegistry(r *Registry, sessID string) {
	fs.mu.Lock()
	fs.registry = r
	fs.sessID = sessID
	fs.mu.Unlock()
}

// SetYolo enables/disables sandbox bypass.
func (fs *FS9P) SetYolo(yolo bool) {
	fs.mu.Lock()
	fs.yolo = yolo
	fs.mu.Unlock()
}

// CWD returns the current working directory.
func (fs *FS9P) CWD() string {
	fs.mu.RLock()
	defer fs.mu.RUnlock()
	return fs.cwd
}

// SetCWD sets the current working directory.
func (fs *FS9P) SetCWD(dir string) {
	fs.mu.Lock()
	fs.cwd = dir
	fs.mu.Unlock()
}

// SetEnv sets an environment variable.
func (fs *FS9P) SetEnv(key, value string) {
	fs.mu.Lock()
	if fs.env == nil {
		fs.env = make(map[string]string)
	}
	fs.env[key] = value
	fs.mu.Unlock()
}

// GetEnv returns an environment variable value.
func (fs *FS9P) GetEnv(key string) string {
	fs.mu.RLock()
	defer fs.mu.RUnlock()
	return fs.env[key]
}

// --- Tool Registry ---

// ListTools returns loaded tools for the current session.
func (fs *FS9P) ListTools() []toolsrv.ToolInfo {
	fs.mu.RLock()
	reg := fs.registry
	sid := fs.sessID
	fs.mu.RUnlock()

	if reg == nil || sid == "" {
		return nil
	}
	return reg.Loaded(sid)
}

// LoadTool loads a tool by name into the session registry.
func (fs *FS9P) LoadTool(name string) error {
	fs.mu.RLock()
	reg := fs.registry
	sid := fs.sessID
	fs.mu.RUnlock()

	if reg == nil {
		return fmt.Errorf("no tool registry configured")
	}
	if sid == "" {
		return fmt.Errorf("no session ID configured")
	}

	if err := reg.Load(sid, name); err != nil {
		return err
	}
	if fs.OnToolsChanged != nil {
		fs.OnToolsChanged()
	}
	return nil
}

// UnloadTool removes a tool from the session registry.
func (fs *FS9P) UnloadTool(name string) error {
	fs.mu.RLock()
	reg := fs.registry
	sid := fs.sessID
	fs.mu.RUnlock()

	if reg == nil {
		return fmt.Errorf("no tool registry configured")
	}
	if sid == "" {
		return fmt.Errorf("no session ID configured")
	}

	if err := reg.Unload(sid, name); err != nil {
		return err
	}
	if fs.OnToolsChanged != nil {
		fs.OnToolsChanged()
	}
	return nil
}

// --- Process Management ---

// allocPID allocates a new process ID.
func (fs *FS9P) allocPID() int {
	fs.procMu.Lock()
	defer fs.procMu.Unlock()
	pid := fs.nextPID
	fs.nextPID++
	return pid
}

// NewProc creates a new process and starts execution.
// If background is false, blocks until completion and returns result.
// If background is true, returns immediately with pid.
func (fs *FS9P) NewProc(ctx context.Context, payload string, background bool) (result string, pid int, err error) {
	// Parse payload
	args := parsePayload(payload)
	toolName := args["tool"]
	if toolName == "" {
		return "", 0, fmt.Errorf("missing 'tool' in payload")
	}

	// Check proc limit
	fs.procMu.Lock()
	if len(fs.procs) >= fs.procLimit {
		fs.procMu.Unlock()
		return "", 0, fmt.Errorf("process limit reached (%d)", fs.procLimit)
	}
	fs.procMu.Unlock()

	// Look up tool
	fs.mu.RLock()
	reg := fs.registry
	sid := fs.sessID
	cwd := fs.cwd
	yolo := fs.yolo
	envCopy := make(map[string]string)
	for k, v := range fs.env {
		envCopy[k] = v
	}
	fs.mu.RUnlock()

	if reg == nil || sid == "" {
		return "", 0, fmt.Errorf("tool registry not configured")
	}

	info, ok := reg.Lookup(sid, toolName)
	if !ok {
		return "", 0, fmt.Errorf("tool not found: %s", toolName)
	}

	// Create proc
	procCtx, cancel := context.WithCancel(ctx)
	proc := &Proc9P{
		PID:       fs.allocPID(),
		Tool:      toolName,
		Args:      args,
		StartTime: time.Now(),
		done:      make(chan struct{}),
		cancel:    cancel,
	}

	fs.procMu.Lock()
	fs.procs[proc.PID] = proc
	fs.procMu.Unlock()

	// Execute in goroutine
	go func() {
		defer close(proc.done)
		defer cancel()

		out, exitCode := fs.executeTool(procCtx, info, args, cwd, envCopy, yolo)

		proc.mu.Lock()
		proc.output.WriteString(out)
		proc.ExitCode = exitCode
		proc.Exited = true
		proc.EndTime = time.Now()
		proc.mu.Unlock()
	}()

	if background {
		return fmt.Sprintf("%d", proc.PID), proc.PID, nil
	}

	// Wait for completion
	<-proc.done

	proc.mu.Lock()
	result = proc.output.String()
	exitCode := proc.ExitCode
	proc.mu.Unlock()

	// Auto-cleanup for foreground procs
	fs.procMu.Lock()
	delete(fs.procs, proc.PID)
	fs.procMu.Unlock()

	if exitCode != 0 {
		return result, proc.PID, fmt.Errorf("exit %d", exitCode)
	}
	return result, proc.PID, nil
}

// executeTool runs a tool and returns output + exit code.
func (fs *FS9P) executeTool(ctx context.Context, info toolsrv.ToolInfo, args map[string]string, cwd string, envExtra map[string]string, yolo bool) (string, int) {
	// Convert args map to JSON for the execution path
	jsonArgs := argsToJSON(args)

	cfg := ExecConfig{
		CWD:  cwd,
		Env:  envExtra,
		Yolo: yolo,
	}

	result, err := ExecuteTool(ctx, info, jsonArgs, cfg)
	if err != nil {
		return fmt.Sprintf("error: %v", err), 1
	}
	return string(result), 0
}

// GetProc returns a process by PID.
func (fs *FS9P) GetProc(pid int) *Proc9P {
	fs.procMu.Lock()
	defer fs.procMu.Unlock()
	return fs.procs[pid]
}

// ListProcs returns all process PIDs.
func (fs *FS9P) ListProcs() []int {
	fs.procMu.Lock()
	defer fs.procMu.Unlock()
	pids := make([]int, 0, len(fs.procs))
	for pid := range fs.procs {
		pids = append(pids, pid)
	}
	return pids
}

// DismissProc removes a process from the list.
func (fs *FS9P) DismissProc(pid int) bool {
	fs.procMu.Lock()
	defer fs.procMu.Unlock()
	if _, ok := fs.procs[pid]; ok {
		delete(fs.procs, pid)
		return true
	}
	return false
}

// SignalProc sends a signal to a process.
func (fs *FS9P) SignalProc(pid int, sig syscall.Signal) error {
	proc := fs.GetProc(pid)
	if proc == nil {
		return fmt.Errorf("process not found: %d", pid)
	}

	// Cancel the context (for graceful stop)
	if sig == syscall.SIGTERM || sig == syscall.SIGINT {
		proc.cancel()
	}

	// Also signal the underlying process if available
	proc.mu.Lock()
	p := proc.proc
	proc.mu.Unlock()
	if p != nil {
		p.Signal(sig)
	}

	return nil
}

// --- Host Info ---

// HostInfo9P returns host information.
func (fs *FS9P) HostInfo9P() string {
	fs.mu.RLock()
	cwd := fs.cwd
	fs.mu.RUnlock()

	return fmt.Sprintf("platform=%s\narch=%s\ncwd=%s\ngit=%v\n",
		runtime.GOOS, runtime.GOARCH, cwd, isGitRepo9P(cwd))
}

func isGitRepo9P(dir string) bool {
	info, err := os.Stat(dir + "/.git")
	return err == nil && info.IsDir()
}

// --- Payload Parsing ---

// parsePayload parses key=value lines into a map.
func parsePayload(payload string) map[string]string {
	args := make(map[string]string)
	for _, line := range strings.Split(payload, "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		idx := strings.Index(line, "=")
		if idx < 0 {
			continue
		}
		key := line[:idx]
		value := line[idx+1:]
		// Unescape \n and \\ in one pass to handle \\n correctly
		value = unescapeValue(value)
		args[key] = value
	}
	return args
}

// unescapeValue reverses escapeValue: \\n -> newline, \\\\ -> backslash
func unescapeValue(s string) string {
	result := make([]byte, 0, len(s))
	for i := 0; i < len(s); i++ {
		if s[i] == '\\' && i+1 < len(s) {
			switch s[i+1] {
			case 'n':
				result = append(result, '\n')
				i++
				continue
			case '\\':
				result = append(result, '\\')
				i++
				continue
			}
		}
		result = append(result, s[i])
	}
	return string(result)
}

// argsToJSON converts the args map to JSON for the existing tool path.
func argsToJSON(args map[string]string) []byte {
	// Build JSON manually to avoid import cycle
	var buf bytes.Buffer
	buf.WriteByte('{')
	first := true
	for k, v := range args {
		if k == "tool" {
			continue // tool name is not part of the args
		}
		if !first {
			buf.WriteByte(',')
		}
		first = false
		buf.WriteByte('"')
		buf.WriteString(escapeJSON(k))
		buf.WriteString("\":")
		buf.WriteByte('"')
		buf.WriteString(escapeJSON(v))
		buf.WriteByte('"')
	}
	buf.WriteByte('}')
	return buf.Bytes()
}

func escapeJSON(s string) string {
	s = strings.ReplaceAll(s, "\\", "\\\\")
	s = strings.ReplaceAll(s, "\"", "\\\"")
	s = strings.ReplaceAll(s, "\n", "\\n")
	s = strings.ReplaceAll(s, "\r", "\\r")
	s = strings.ReplaceAll(s, "\t", "\\t")
	return s
}

// --- Proc9P Methods ---

// Output returns the current output buffer.
func (p *Proc9P) Output() string {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.output.String()
}

// Wait blocks until the process exits, returns exit code.
func (p *Proc9P) Wait() int {
	<-p.done
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.ExitCode
}

// Stat returns a status string.
func (p *Proc9P) Stat() string {
	p.mu.Lock()
	defer p.mu.Unlock()

	runtime := time.Since(p.StartTime)
	if p.Exited {
		runtime = p.EndTime.Sub(p.StartTime)
		return fmt.Sprintf("exited %d\nruntime=%v\ntool=%s\n", p.ExitCode, runtime, p.Tool)
	}
	return fmt.Sprintf("running\nruntime=%v\ntool=%s\n", runtime, p.Tool)
}

// --- Request Handlers (for integration with 9P server) ---

// HandleCtl processes ctl commands.
func (fs *FS9P) HandleCtl(input string) error {
	parts := strings.Fields(input)
	if len(parts) == 0 {
		return nil
	}

	switch parts[0] {
	case "load":
		if len(parts) < 2 {
			return fmt.Errorf("load requires tool name")
		}
		return fs.LoadTool(parts[1])
	case "unload":
		if len(parts) < 2 {
			return fmt.Errorf("unload requires tool name")
		}
		return fs.UnloadTool(parts[1])
	case "env":
		if len(parts) < 2 {
			return fmt.Errorf("env requires KEY=VALUE")
		}
		// Rejoin in case value contains spaces
		kv := strings.Join(parts[1:], " ")
		idx := strings.Index(kv, "=")
		if idx < 0 {
			return fmt.Errorf("env requires KEY=VALUE format")
		}
		fs.SetEnv(kv[:idx], kv[idx+1:])
		return nil
	case "cwd":
		if len(parts) < 2 {
			return fmt.Errorf("cwd requires path")
		}
		fs.SetCWD(parts[1])
		return nil
	default:
		return fmt.Errorf("unknown command: %s", parts[0])
	}
}

// HandleTools processes reads/writes to /tools.
func (fs *FS9P) HandleToolsRead() string {
	tools := fs.ListTools()
	data, _ := json.Marshal(tools)
	return string(data)
}

// HandleToolsWrite loads a tool by name.
func (fs *FS9P) HandleToolsWrite(name string) error {
	return fs.LoadTool(strings.TrimSpace(name))
}

// HandleProcNew is the rdwr handler for /proc/new (blocking).
func (fs *FS9P) HandleProcNew(ctx context.Context, payload string) (string, error) {
	result, _, err := fs.NewProc(ctx, payload, false)
	return result, err
}

// HandleProcNewBg is the write handler for /proc/new.bg (background).
func (fs *FS9P) HandleProcNewBg(ctx context.Context, payload string) (string, error) {
	result, _, err := fs.NewProc(ctx, payload, true)
	return result, err
}

// HandleProcCtl processes writes to /proc/<pid>/ctl.
func (fs *FS9P) HandleProcCtl(pid int, input string) error {
	parts := strings.Fields(input)
	if len(parts) == 0 {
		return nil
	}

	switch parts[0] {
	case "signal":
		if len(parts) < 2 {
			return fmt.Errorf("signal requires signal number")
		}
		sig, err := strconv.Atoi(parts[1])
		if err != nil {
			return fmt.Errorf("invalid signal: %s", parts[1])
		}
		return fs.SignalProc(pid, syscall.Signal(sig))
	case "dismiss":
		fs.DismissProc(pid)
		return nil
	default:
		return fmt.Errorf("unknown command: %s", parts[0])
	}
}

// --- Atomic counters for unique IDs ---

var globalProcCounter atomic.Int64

func init() {
	globalProcCounter.Store(time.Now().UnixNano() % 10000)
}