ollie-core/toolsrv/server.go

package toolsrv

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

	"ollie/detach"
	"ollie/paths"
	"ollie/skills"
)

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

// 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)
	OnPreDispatch func()
	OnClose       func()
	OnEnvSet      func(key, value string)

	// Strict rejects inline {code} steps; only {tool} steps are allowed.
	Strict bool

	// 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

	// builtins maps tool names to their handler functions.
	builtins map[string]Handler

	toolRegistry   *Registry
	skillsRegistry *skills.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)

// WithStrict rejects inline {code} steps; only {tool} steps are allowed.
func WithStrict() Option { return func(s *Server) { s.Strict = true } }

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

// WithOnPreDispatch registers a hook called before each Dispatch.
func WithOnPreDispatch(fn func()) Option { return func(s *Server) { s.OnPreDispatch = fn } }

// 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
}

// WithBuiltins registers the built-in tool handlers.
func WithBuiltins(m map[string]Handler) Option {
	return func(s *Server) { s.builtins = m }
}

// 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
	}
}

// WithSkillsRegistry attaches a skills registry to the Server.
func WithSkillsRegistry(r *skills.Registry) Option {
	return func(s *Server) { s.skillsRegistry = r }
}


// ListTools implements Server, returning shell plus any
// tools promoted in the session's tool registry.
func (e *Server) ListTools() ([]ToolInfo, error) {
	all := []ToolInfo{
		{
			Name: "shell",
			Description: `Execute a single bash command in a sandboxed environment.

Usage: {"cmd": "your command here"}

Sandbox prevents dangerous operations. Use for computation, builds, scripting.
timeout applies to each call (default: 30s). A non-zero exit is an error.`,
			InputSchema: json.RawMessage(`{
				"type": "object",
				"required": ["cmd"],
				"properties": {
					"cmd":      {"type": "string", "description": "Bash command to execute."},
					"timeout":  {"type": "integer", "description": "Timeout in seconds (default: 30). Use 0 for no timeout."},
					"sandbox":  {"type": "string",  "description": "Sandbox profile name (default: default)."},
					"elevated": {"type": "boolean", "description": "Run outside the sandbox via elevation broker."}
				}
			}`),
		},
		{
			Name: "tool_list",
			Description: `List all available tools that can be loaded.

Usage: {"name": "toolname"} — if name is provided, loads that tool.
Otherwise lists all tools with descriptions.`,
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {}
			}`),
		},
		{
			Name: "tool_load",
			Description: `Load a tool by name into the current session.

Usage: {"name": "toolname"}

After loading, the tool becomes a native callable function.`,
			InputSchema: json.RawMessage(`{
				"type": "object",
				"required": ["name"],
				"properties": {
					"name": {"type": "string", "description": "Tool name to load."}
				}
			}`),
		},
		{
			Name: "tool_active",
			Description: `List tools currently loaded (promoted) in this session.

Usage: (no arguments)

Returns tools with descriptions, one per line.`,
			InputSchema: json.RawMessage(`{
				"type": "object",
				"properties": {}
			}`),
		},
		{
			Name: "reasoning_think",
			Description: `Scratchpad for externalizing reasoning. Recorded in history but not shown to the user.`,
			InputSchema: json.RawMessage(`{
				"type": "object",
				"required": ["thought"],
				"properties": {
					"thought": {"type": "string", "description": "Your reasoning."}
				}
			}`),
		},
	}

	if e.toolRegistry != nil && e.sessionID != "" {
		all = append(all, e.toolRegistry.Loaded(e.sessionID)...)
	}
	// Append skill tools if skills registry is available
	skillTools := ListSkillsTools(e.skillsRegistry, e.sessionID)
	all = append(all, skillTools...)
	return all, nil
}

// CallTool implements Server.
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)
		}
	}
	result, err := e.Dispatch(ctx, tool, args)
	if err != nil {
		return json.Marshal(map[string]any{
			"isError": true,
			"content": []map[string]string{{"type": "text", "text": err.Error()}},
		})
	}
	// If the tool output is already structured content-block JSON, pass it through.
	var structured map[string]json.RawMessage
	if json.Unmarshal([]byte(result), &structured) == nil {
		if _, ok := structured["content"]; ok {
			return []byte(result), nil
		}
	}
	return json.Marshal(map[string]any{
		"content": []map[string]string{{"type": "text", "text": result}},
	})
}

// 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 := filepath.Join(ToolsPath(), tool)

	// 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)
		}
	}

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

	var result string
	var err error
	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()
	}
}

// Dispatch routes tool calls to the appropriate handler.
// Dispatch routes tool calls to the appropriate built-in handler.
func (e *Server) Dispatch(ctx context.Context, name string, args json.RawMessage) (string, error) {
	if e.OnPreDispatch != nil {
		e.OnPreDispatch()
	}
	if h, ok := e.builtins[name]; ok {
		return h(ctx, e, args)
	}
	return "", fmt.Errorf("unknown tool: %s", name)
}

// 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()
	}
}