ollie/cmd/toolsrv/internal/server/proc.go

// proc.go - Process management and tool operations for toolsrvclient.
package server

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

	"ollie/cmd/toolsrv/internal/exec"
	"ollie/cmd/toolsrv/internal/quirks"
	"ollie/cmd/toolsrv/internal/registry"
	"ollie/toolsrv/protocol"
	"ollie/util"
)

// State implements the filesystem state for toolsrvclient.
type State struct {
	mu sync.RWMutex

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

	// Path-based lock coordination for cross-agent serialization.
	// All foreground tool calls acquire a lock based on scope+path before
	// execution, ensuring writes to the same path serialize regardless of
	// which agent initiated the call.
	pathLocks pathLockTable

	// Process management
	procMu       sync.Mutex
	procs        map[int]*Proc
	nextID       int
	procLimit    int
	procReserved int
	procGCOnce   sync.Once

	// Callbacks
	OnToolsChanged func()
	OnProcExit     func(proc *Proc) // called when a background proc exits
}

// Proc represents a running or completed process.
type Proc struct {
	ID        int
	Tool      string
	Args      map[string]string
	Cmd       string // human-readable command
	AgentID   string // agent that owns this proc
	SessionID string // session that owns this proc
	StartTime time.Time
	EndTime   time.Time
	ExitCode  int
	Exited    bool
	LastRead  time.Time // last time output was read; GC TTL counts from here

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

// Lock locks the proc mutex.
func (p *Proc) Lock() {
	p.mu.Lock()
}

// Unlock unlocks the proc mutex.
func (p *Proc) Unlock() {
	p.mu.Unlock()
}

// OutputString returns the output buffer as a string (caller must hold lock).
func (p *Proc) OutputString() string {
	return p.output.String()
}

// NewState creates a new filesystem state.
func NewState(cwd string) *State {
	return &State{
		cwd:       cwd,
		env:       make(map[string]string),
		procs:     make(map[int]*Proc),
		nextID:    1,
		procLimit: 32,
	}
}

// SetRegistry configures the tool registry.
func (st *State) SetRegistry(r *registry.Registry) {
	st.mu.Lock()
	st.registry = r
	st.mu.Unlock()
}

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

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

// SetCWD sets the current working directory.
func (st *State) SetCWD(dir string) {
	st.mu.Lock()
	st.cwd = util.ExpandHome(dir)
	st.mu.Unlock()
}

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

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

// --- Tool Registry ---

// ListTools returns loaded tools for the given agent.
func (st *State) ListTools(agentID string) []protocol.ToolInfo {
	if agentID == "" {
		return nil
	}
	st.mu.RLock()
	reg := st.registry
	st.mu.RUnlock()

	if reg == nil {
		return nil
	}
	return reg.Loaded(agentID)
}

// Revision returns the current tool registry revision for an agent.
func (st *State) Revision(agentID string) uint64 {
	if agentID == "" {
		return 0
	}
	st.mu.RLock()
	reg := st.registry
	st.mu.RUnlock()
	if reg == nil {
		return 0
	}
	return reg.Revision(agentID)
}

// LoadTool loads a tool by name for the given agent.
func (st *State) LoadTool(agentID, name string) error {
	if agentID == "" {
		return fmt.Errorf("agent ID required")
	}
	st.mu.RLock()
	reg := st.registry
	st.mu.RUnlock()

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

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

// UnloadTool removes a tool from the given agent's registry.
func (st *State) UnloadTool(agentID, name string) error {
	if agentID == "" {
		return fmt.Errorf("agent ID required")
	}
	st.mu.RLock()
	reg := st.registry
	st.mu.RUnlock()

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

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

// ClearAgent removes all tools for an agent's registry.
func (st *State) ClearAgent(agentID string) {
	st.mu.RLock()
	reg := st.registry
	st.mu.RUnlock()

	if reg == nil {
		return
	}

	reg.ClearAgent(agentID)
	if st.OnToolsChanged != nil {
		st.OnToolsChanged()
	}
}

// --- Process Management ---

// allocID allocates a new process ID.
func (st *State) allocID() int {
	st.procMu.Lock()
	defer st.procMu.Unlock()
	pid := st.nextID
	st.nextID++
	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 (st *State) NewProc(ctx context.Context, payload string, background bool) (result string, pid int, err error) {
	// Parse payload
	args := protocol.ParsePayload(payload)
	toolName := args["tool"]
	if toolName == "" {
		return "", 0, fmt.Errorf("missing 'tool' in payload")
	}

	// Reserve a process slot before setup so concurrent callers cannot all pass
	// the limit check while their processes are still being initialized.
	st.procMu.Lock()
	if len(st.procs)+st.procReserved >= st.procLimit {
		st.procMu.Unlock()
		return "", 0, fmt.Errorf("process limit reached (%d)", st.procLimit)
	}
	st.procReserved++
	st.procMu.Unlock()
	reserved := true
	defer func() {
		if reserved {
			st.procMu.Lock()
			st.procReserved--
			st.procMu.Unlock()
		}
	}()

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

	aid := args["agent"]
	if aid == "" {
		return "", 0, fmt.Errorf("missing 'agent' in payload")
	}
	envCopy["OLLIE_UNAME"] = aid

	// Get session ID from environment
	sid := envCopy["OLLIE_SESSION_ID"]

	// Get human-readable command
	cmd := args["cmd"]
	if cmd == "" {
		cmd = toolName
	}

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

	// CAPABILITY CHECK: This is the structural enforcement point.
	// The agent can only call tools that exist in its registry (namespace).
	// This check is NOT behavioral—it cannot be bypassed by prompt injection,
	// model persuasion, or clever argument formatting. If the tool isn't
	// in the registry, execution is denied unconditionally.
	//
	// See registry/doc.go for the capability model.
	info, ok := reg.Lookup(aid, toolName)
	if !ok {
		return "", 0, fmt.Errorf("tool not loaded: %s (add to tools in agent config)", toolName)
	}

	// Reject shell calls that invoke a native tool.
	// This is a hack for stupid, broken models that ignore prompt instructions
	// and shell out to tools instead of calling them natively.
	if toolName == "shell" {
		if conflicting := quirks.ShellInvokesNativeTool(args["cmd"], reg.All()); conflicting != "" {
			return "", 0, fmt.Errorf("STOP. You called shell to run '%s', but '%s' is a native tool. Call it directly: %s(...). NEVER use shell to invoke native tools", conflicting, conflicting, conflicting)
		}
	}

	// Create proc context. For background procs, detach from request context
	// so the process continues after the 9P request returns.
	var procCtx context.Context
	var cancel context.CancelFunc
	if background {
		procCtx, cancel = context.WithCancel(context.Background())
	} else {
		procCtx, cancel = context.WithCancel(ctx)
	}
	proc := &Proc{
		ID:        st.allocID(),
		Tool:      toolName,
		Args:      args,
		Cmd:       cmd,
		AgentID:   aid,
		SessionID: sid,
		StartTime: time.Now(),
		done:      make(chan struct{}),
		cancel:    cancel,
	}

	st.procMu.Lock()
	st.procs[proc.ID] = proc
	st.procReserved--
	st.procMu.Unlock()
	reserved = false

	// For background procs, stream output directly into the proc buffer.
	// For foreground, use nil (local buffer in exec path).
	var outputWriter io.Writer
	var startedCh chan exec.StartResult
	timeout := 30 // default for foreground
	if info.Timeout >= 0 {
		timeout = info.Timeout // tool declares its own default
	}
	if background {
		outputWriter = &procWriter{proc: proc}
		startedCh = make(chan exec.StartResult, 1)
		timeout = 0 // background procs never timeout
	}

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

		// Acquire path-based lock for foreground procs to serialize
		// conflicting writes across agents. Background procs are
		// long-running and don't hold locks.
		if !background {
			release := st.pathLocks.acquire(info.Scope, args["path"])
			defer release()
		}

		out, exitCode := st.executeTool(procCtx, info, args, cwd, envCopy, yolo, timeout, outputWriter, startedCh)

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

		// Notify on background proc exit
		if background && st.OnProcExit != nil {
			st.OnProcExit(proc)
		}
	}()

	// For background procs, wait for the process to start and store its reference.
	if background {
		started := <-startedCh
		if started.Err != nil {
			proc.mu.Lock()
			proc.Exited = true
			proc.ExitCode = 1
			proc.EndTime = time.Now()
			proc.mu.Unlock()
			st.procMu.Lock()
			delete(st.procs, proc.ID)
			st.procMu.Unlock()
			return "", 0, started.Err
		}
		if started.Process != nil {
			proc.mu.Lock()
			proc.proc = started.Process
			proc.mu.Unlock()
		}
		return fmt.Sprintf("%d", proc.ID), proc.ID, nil
	}

	// Wait for completion or context cancellation
	select {
	case <-proc.done:
		// Normal completion
	case <-ctx.Done():
		// Context was canceled - cancel the proc and return error
		cancel()
		// Wait briefly for the process to clean up
		select {
		case <-proc.done:
		case <-time.After(100 * time.Millisecond):
		}
		st.procMu.Lock()
		delete(st.procs, proc.ID)
		st.procMu.Unlock()
		return "", proc.ID, ctx.Err()
	}

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

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

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

// executeTool runs a tool and returns output + exit code.
// timeout: 0 = no timeout, >0 = seconds.
func (st *State) executeTool(ctx context.Context, info protocol.ToolInfo, args map[string]string, cwd string, envExtra map[string]string, yolo bool, timeout int, output io.Writer, started chan exec.StartResult) (string, int) {
	// Convert args map to JSON for the execution path
	jsonArgs := argsToJSON(args)

	cfg := exec.Config{
		CWD:     cwd,
		Env:     envExtra,
		Yolo:    yolo,
		Timeout: timeout,
		Output:  output,
		Started: started,
	}

	result, err := exec.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 (st *State) GetProc(pid int) *Proc {
	st.procMu.Lock()
	defer st.procMu.Unlock()
	return st.procs[pid]
}

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

// ListProcsForAgent returns a formatted list of processes for the given agent.
// If agentID is empty, returns all processes.
// Format: "pid\tstate\ttool\tcmd\n"
func (st *State) ListProcsForAgent(agentID string) string {
	st.procMu.Lock()
	defer st.procMu.Unlock()
	var buf bytes.Buffer
	for pid, p := range st.procs {
		p.mu.Lock()
		if agentID != "" && p.AgentID != agentID {
			p.mu.Unlock()
			continue
		}
		var state string
		if p.Exited {
			state = fmt.Sprintf("exited (%d)", p.ExitCode)
		} else {
			state = "running"
		}
		tool := p.Tool
		cmd := p.Cmd
		p.mu.Unlock()
		fmt.Fprintf(&buf, "%-6d %-12s %-16s %s\n", pid, state, tool, cmd)
	}
	return buf.String()
}

// KillAll sends SIGKILL to all running procs and cancels their contexts.
// Called on shutdown to ensure no orphaned processes.
func (st *State) KillAll() {
	st.procMu.Lock()
	procs := make([]*Proc, 0, len(st.procs))
	for _, p := range st.procs {
		procs = append(procs, p)
	}
	st.procMu.Unlock()

	for _, p := range procs {
		p.mu.Lock()
		proc := p.proc
		exited := p.Exited
		p.mu.Unlock()
		if !exited && proc != nil {
			syscall.Kill(-proc.Pid, syscall.SIGKILL)
		}
		p.cancel()
	}
}

// DismissProc removes a completed process from the list. Running processes
// must be terminated before they can be dismissed so they remain visible to
// shutdown and process cleanup.
func (st *State) DismissProc(pid int) bool {
	st.procMu.Lock()
	p, ok := st.procs[pid]
	if !ok {
		st.procMu.Unlock()
		return false
	}
	p.mu.Lock()
	exited := p.Exited
	p.mu.Unlock()
	if !exited {
		st.procMu.Unlock()
		return false
	}
	delete(st.procs, pid)
	st.procMu.Unlock()
	return true
}

const (
	procGCTTL       = 10 * time.Minute
	procOutputLimit = 10 * 1024 * 1024
)

// StartProcGC runs a background goroutine that removes exited procs
// whose output has been read and whose LastRead is older than procGCTTL.
func (st *State) StartProcGC(ctx context.Context) {
	st.procGCOnce.Do(func() {
		go func() {
			ticker := time.NewTicker(1 * time.Minute)
			defer ticker.Stop()
			for {
				select {
				case <-ctx.Done():
					return
				case <-ticker.C:
					st.gcProcs()
				}
			}
		}()
	})
}

func (st *State) gcProcs() {
	st.procMu.Lock()
	defer st.procMu.Unlock()
	now := time.Now()
	for pid, p := range st.procs {
		p.mu.Lock()
		exited := p.Exited
		lastRead := p.LastRead
		endTime := p.EndTime
		p.mu.Unlock()
		if exited && ((!lastRead.IsZero() && now.Sub(lastRead) >= procGCTTL) || (!endTime.IsZero() && now.Sub(endTime) >= procGCTTL)) {
			delete(st.procs, pid)
		}
	}
}

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

	// Signal the process group (we set Setpgid: true)
	proc.mu.Lock()
	p := proc.proc
	proc.mu.Unlock()
	if p != nil {
		syscall.Kill(-p.Pid, sig)
	} else if sig == syscall.SIGTERM || sig == syscall.SIGKILL {
		// No OS process reference (e.g., bypass background proc waiting for approval).
		// Cancel the context to stop the operation.
		proc.cancel()
	}

	// Cancel context on SIGKILL (hard stop)
	if sig == syscall.SIGKILL {
		proc.cancel()
	}

	return nil
}

// --- Host Info ---

// HostInfo returns host information.
func (st *State) HostInfo() string {
	st.mu.RLock()
	cwd := st.cwd
	st.mu.RUnlock()

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

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

// --- Request Handlers ---

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

	switch parts[0] {
	case "load":
		if len(parts) < 3 {
			return fmt.Errorf("load requires agent ID and tool name")
		}
		return st.LoadTool(parts[1], parts[2])
	case "unload":
		if len(parts) < 3 {
			return fmt.Errorf("unload requires agent ID and tool name")
		}
		return st.UnloadTool(parts[1], parts[2])
	case "clear":
		if len(parts) < 2 {
			return fmt.Errorf("clear requires agent ID")
		}
		st.ClearAgent(parts[1])
		return nil
	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")
		}
		key := kv[:idx]
		if key == "OLLIE_UNAME" {
			return fmt.Errorf("OLLIE_UNAME is set per-request via agent ID, not via env")
		}
		st.SetEnv(key, kv[idx+1:])
		return nil
	case "cwd":
		if len(parts) < 2 {
			return fmt.Errorf("cwd requires path")
		}
		st.SetCWD(parts[1])
		return nil
	default:
		return fmt.Errorf("unknown command: %s", parts[0])
	}
}

// HandleToolsRevision returns the registry revision for an agent.
func (st *State) HandleToolsRevision(agentID string) string {
	return strconv.FormatUint(st.Revision(strings.TrimSpace(agentID)), 10)
}

// HandleToolsRequest handles the tools rdwr: write agent ID, read tool list.
func (st *State) HandleToolsRequest(agentID string) string {
	tools := st.ListTools(agentID)
	data, _ := json.Marshal(tools)
	return string(data)
}

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

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

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

	switch parts[0] {
	case "term":
		return st.SignalProc(pid, syscall.SIGTERM)
	case "kill":
		return st.SignalProc(pid, syscall.SIGKILL)
	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 st.SignalProc(pid, syscall.Signal(sig))
	case "dismiss":
		st.DismissProc(pid)
		return nil
	default:
		return fmt.Errorf("unknown command: %s", parts[0])
	}
}

// --- procWriter ---

// procWriter is a thread-safe writer that streams into a Proc's output buffer.
type procWriter struct {
	proc *Proc
}

func (pw *procWriter) Write(p []byte) (int, error) {
	pw.proc.mu.Lock()
	defer pw.proc.mu.Unlock()
	remaining := procOutputLimit - pw.proc.output.Len()
	if remaining > 0 {
		if len(p) > remaining {
			_, _ = pw.proc.output.Write(p[:remaining])
		} else {
			_, _ = pw.proc.output.Write(p)
		}
	}
	return len(p), nil
}

// --- Proc Methods ---

// Output returns the current output buffer and updates LastRead.
func (p *Proc) Output() string {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.LastRead = time.Now()
	return p.output.String()
}

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

// Stat returns a status string in key=value format.
func (p *Proc) Stat() string {
	p.mu.Lock()
	defer p.mu.Unlock()

	rt := time.Since(p.StartTime)
	if p.Exited {
		rt = p.EndTime.Sub(p.StartTime)
		return fmt.Sprintf("id=%d\ncommand=%s\nexited=true\nexit_code=%d\nruntime=%v\ntool=%s\n", p.ID, p.Cmd, p.ExitCode, rt, p.Tool)
	}
	return fmt.Sprintf("id=%d\ncommand=%s\nexited=false\nruntime=%v\ntool=%s\n", p.ID, p.Cmd, rt, p.Tool)
}

// --- Helpers ---

// argsToJSON converts the args map to JSON for the existing tool path.
func argsToJSON(args map[string]string) []byte {
	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
}