ollie/toolsrv/server.go

package toolsrv

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"syscall"
	"time"

	"ollie/detach"
	"ollie/paths"
)

const (
	failureWindow = 1 * time.Minute
	maxFailures   = 5
	blockDuration = 5 * time.Minute
)

// RateLimitedError is returned by checkRateLimit when the shell is blocked
// due to too many validation failures. The agent loop uses this to detect
// the condition and inject a clear signal rather than counting it as a
// normal tool error.
type RateLimitedError struct {
	Remaining time.Duration
}

func (e *RateLimitedError) Error() string {
	return fmt.Sprintf("rate limited: too many validation failures, blocked for %v", e.Remaining)
}

// Server runs code in a sandboxed environment.
type Server struct {
	// cwd is the working directory for sandboxed commands. If empty,
	// the process working directory is used.
	wdMu sync.RWMutex
	cwd  string

	// envExtra holds per-session environment variables injected via SetEnv.
	envMu    sync.RWMutex
	envExtra map[string]string

// Hooks for lifecycle events (harnesses like 9P can inject mount logic here)
	OnClose  func()
	OnEnvSet func(key, value string)

	// Yolo skips the landrun sandbox for all execution.
	Yolo bool

	// allowTools restricts which named tool scripts can be invoked via call_tool/pipe.
	// Empty means all are allowed.
	allowTools map[string]bool

	toolRegistry   *Registry
	sessionID      string

	// OnInjection is called when a skill is loaded and its content
	// should be injected into the agent's context.
	OnInjection func(content string)

	// OnToolsChanged is called when tool_list detects the available tools
	// have changed, providing the updated preamble-format listing.
	OnToolsChanged func(listing string)

	// rate limiting state (per-Server)
	rateLimitMu        sync.Mutex
	validationFailures int
	lastFailure        time.Time
	blockedUntil       time.Time

	// Detached process management
	detachMu sync.Mutex
	detachCh chan struct{} // signal to detach the currently running process
	detached []*detach.Process
	OnDetach func(pid int, cmd string)   // hook: called when a process is detached
	OnExit   func(pid int, exitCode int) // hook: called when a detached process exits
}

// Option configures a Server.
type Option func(*Server)

// WithYolo skips the landrun sandbox.
func WithYolo() Option { return func(s *Server) { s.Yolo = true } }

// WithOnClose registers a hook called during Close.
func WithOnClose(fn func()) Option { return func(s *Server) { s.OnClose = fn } }

// SetOnExit sets the hook called when a detached process exits.
func (e *Server) SetOnExit(fn func(pid, exitCode int)) { e.OnExit = fn }

// WithOnEnvSet registers a hook called each time SetEnv is called.
func WithOnEnvSet(fn func(key, value string)) Option { return func(s *Server) { s.OnEnvSet = fn } }

// WithAllowTools restricts which tool scripts can be invoked.
func WithAllowTools(names []string) Option {
	return func(s *Server) {
		if len(names) > 0 {
			s.allowTools = make(map[string]bool, len(names))
			for _, n := range names {
				s.allowTools[n] = true
			}
		}
	}
}

// AllowTools returns the set of allowed tool names, or nil if unrestricted.
func (e *Server) AllowTools() []string {
	if len(e.allowTools) == 0 {
		return nil
	}
	out := make([]string, 0, len(e.allowTools))
	for k := range e.allowTools {
		out = append(out, k)
	}
	return out
}

// WithToolRegistry attaches a tool registry and session ID to the Server.
func WithToolRegistry(r *Registry, sessionID string) Option {
	return func(s *Server) {
		s.toolRegistry = r
		s.sessionID = sessionID
	}
}

// LoadTool loads a tool by name into this server's tool registry.
// Returns an error if no registry is configured or the tool cannot be found.
func (e *Server) LoadTool(name string) error {
	if e.toolRegistry == nil {
		return fmt.Errorf("no tool registry configured")
	}
	if e.sessionID == "" {
		return fmt.Errorf("no session ID configured")
	}
	return e.toolRegistry.Load(e.sessionID, name)
}



// ListTools implements Server, returning tools loaded in the session's
// tool registry. With zero built-in tools, only dynamically loaded tools
// appear here.
func (e *Server) ListTools() ([]ToolInfo, error) {
	var all []ToolInfo

	if e.toolRegistry != nil && e.sessionID != "" {
		all = append(all, e.toolRegistry.Loaded(e.sessionID)...)
	}
	return all, nil
}

// CallTool implements Runner.
// With zero built-in tools, all tools must be loaded in the registry first.
func (e *Server) CallTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
	if e.toolRegistry != nil && e.sessionID != "" {
		if _, promoted := e.toolRegistry.Lookup(e.sessionID, tool); promoted {
			return e.callPromotedTool(ctx, tool, args)
		}
	}
	return nil, fmt.Errorf("unknown tool: %s (not loaded)", tool)
}

// New creates a new Server with the given working directory.
func New(cwd string) *Server { return &Server{cwd: paths.ExpandHome(cwd)} }

// SetCWD updates the working directory used for subsequent command executions.
func (e *Server) SetCWD(dir string) {
	e.wdMu.Lock()
	e.cwd = paths.ExpandHome(dir)
	e.wdMu.Unlock()
}

// SetAllowTools restricts which tool scripts can be invoked.
func (e *Server) SetAllowTools(names []string) {
	if len(names) > 0 {
		e.allowTools = make(map[string]bool, len(names))
		for _, n := range names {
			e.allowTools[n] = true
		}
	} else {
		e.allowTools = nil
	}
}

// SetEnv adds a session-scoped environment variable injected into all
// subsequent subprocess invocations for this session.
func (e *Server) SetEnv(key, value string) {
	e.envMu.Lock()
	if e.envExtra == nil {
		e.envExtra = make(map[string]string)
	}
	e.envExtra[key] = value
	e.envMu.Unlock()
	if e.OnEnvSet != nil {
		e.OnEnvSet(key, value)
	}
}

// SetToolRegistry attaches a session-local tool registry.
func (e *Server) SetToolRegistry(r *Registry, sessionID string) {
	e.toolRegistry = r
	e.sessionID = sessionID
}

// callPromotedTool executes a tool promoted via the registry by running the
// script file inside the sandbox, piping the JSON args to stdin.
func (e *Server) callPromotedTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
	// Resolve script path.
	if strings.Contains(tool, "/") || strings.Contains(tool, "..") {
		return nil, fmt.Errorf("invalid tool name")
	}
	path, err := ResolveTool(tool)
	if err != nil {
		return nil, err
	}

	// Extract elevated flag (dispatch-level concern, not passed to tool).
	elevated := false
	var argMap map[string]interface{}
	if err := json.Unmarshal(args, &argMap); err == nil {
		if e, ok := argMap["elevated"]; ok {
			switch v := e.(type) {
			case bool:
				elevated = v
			case string:
				elevated = v == "true" || v == "1"
			}
			// Remove elevated from the args passed to the tool.
			delete(argMap, "elevated")
			args, _ = json.Marshal(argMap)
		}
	}

	// Check if tool requires sudo (from .meta).
	needsSudo := false
	if m, err := LoadMetaFile(tool); err == nil && m != nil {
		resolved := m.Resolve()
		if resolved != nil && resolved.Sudo {
			needsSudo = true
			elevated = true // sudo implies elevation
		}
	}

	// The tool script receives JSON args on stdin.
	code := path
	stdinData := string(args)

	var result string
	if needsSudo {
		e.wdMu.RLock()
		workDir := e.cwd
		e.wdMu.RUnlock()
		sudoCode := fmt.Sprintf("cat <<'OLLIE_EOF' | %s\n%s\nOLLIE_EOF", code, stdinData)
		result, err = e.executeElevatedSudo(ctx, sudoCode, workDir, 30)
	} else if elevated {
		e.wdMu.RLock()
		workDir := e.cwd
		e.wdMu.RUnlock()
		// Broker protocol has no stdin support; pipe JSON via heredoc.
		elevatedCode := fmt.Sprintf("cat <<'OLLIE_EOF' | %s\n%s\nOLLIE_EOF", code, stdinData)
		result, err = e.executeElevated(ctx, elevatedCode, workDir, 30)
	} else {
		result, err = e.executeWithStdin(ctx, code, "bash", 30, "default", false, stdinData)
	}
	if err != nil {
		return json.Marshal(map[string]interface{}{
			"isError": true,
			"content": []map[string]string{{"type": "text", "text": result + ": " + err.Error()}},
		})
	}
	return json.Marshal(map[string]interface{}{
		"content": []map[string]string{{"type": "text", "text": result}},
	})
}

// Close is called when the session ends. Calls OnClose hook if registered.
func (e *Server) Close() {
	e.cleanupDetached()
	if e.OnClose != nil {
		e.OnClose()
	}
}

// Detach signals the currently running process to be detached from the agent.
// The process continues running; its output is captured in a ring buffer.
// Returns false if no process is currently running.
func (e *Server) Detach() bool {
	e.detachMu.Lock()
	ch := e.detachCh
	e.detachMu.Unlock()
	if ch == nil {
		return false
	}
	select {
	case ch <- struct{}{}:
		return true
	default:
		return false
	}
}

// ListDetached returns all detached processes (running and exited).
func (e *Server) ListDetached() []*detach.Process {
	e.detachMu.Lock()
	defer e.detachMu.Unlock()
	out := make([]*detach.Process, len(e.detached))
	copy(out, e.detached)
	return out
}

// ListDetachedInfo returns plain-data snapshots of all detached processes.
func (e *Server) ListDetachedInfo() []detach.InfoData {
	e.detachMu.Lock()
	defer e.detachMu.Unlock()
	out := make([]detach.InfoData, len(e.detached))
	for i, p := range e.detached {
		out[i] = p.Info()
	}
	return out
}

// ListDetachedRaw returns detached process info as []any (each element is map[string]any)
// for consumption by packages that can't import this package directly.
func (e *Server) ListDetachedRaw() []any {
	e.detachMu.Lock()
	defer e.detachMu.Unlock()
	out := make([]any, len(e.detached))
	for i, p := range e.detached {
		info := p.Info()
		out[i] = map[string]any{
			"pid":       info.PID,
			"command":   info.Command,
			"started":   info.Started,
			"exited":    info.Exited,
			"exit_code": info.ExitCode,
		}
	}
	return out
}

// SignalDetached sends a signal to a detached process by PID.
func (e *Server) SignalDetached(pid int, sig syscall.Signal) error {
	e.detachMu.Lock()
	defer e.detachMu.Unlock()
	for _, p := range e.detached {
		if p.PID == pid {
			return p.Signal(sig)
		}
	}
	return fmt.Errorf("no detached process with pid %d", pid)
}

// GetDetachedOutput returns the ring buffer contents for a detached process.
func (e *Server) GetDetachedOutput(pid int) (string, error) {
	e.detachMu.Lock()
	defer e.detachMu.Unlock()
	for _, p := range e.detached {
		if p.PID == pid {
			return p.Output(), nil
		}
	}
	return "", fmt.Errorf("no detached process with pid %d", pid)
}

// DismissDetached removes an exited process from the list.
func (e *Server) DismissDetached(pid int) bool {
	e.detachMu.Lock()
	defer e.detachMu.Unlock()
	for i, p := range e.detached {
		if p.PID == pid && p.Exited {
			e.detached = append(e.detached[:i], e.detached[i+1:]...)
			return true
		}
	}
	return false
}

// cleanupDetached sends SIGTERM to all running detached processes.
func (e *Server) cleanupDetached() {
	e.detachMu.Lock()
	defer e.detachMu.Unlock()
	for _, p := range e.detached {
		p.Mu.Lock()
		if !p.Exited && p.Cmd != nil {
			syscall.Kill(-p.PID, syscall.SIGTERM)
		}
		p.Mu.Unlock()
	}
}