toolsrv: add 9P-based tool server foundation
New files: - fs9p.go: FS9P struct with /proc-based execution model - /proc/new (rdwr): blocking tool execution - /proc/new.bg: background/detached execution - /proc/<pid>/out, wait, stat, ctl: process management - /tools, /ctl, /info: tool registry and host info - Payload format: key=value lines, newline delimited - exec9p.go: Decoupled tool execution - ExecuteTool: standalone function for tool execution - executeSandboxed: sandbox execution without Server deps - executeBypassDirect: bypass broker without Server deps - fs9p_test.go: Unit tests for FS9P Design notes: - Background execution is the internal model - Foreground (/proc/new rdwr) blocks until completion - No per-connection state (cwd/env passed per-call or at attach) - Session ID scoping via registry, not env propagation This is the foundation - actual 9P serving and integration pending.
This commit is contained in:
Levi Neely
parent
9c3e935a7e
commit
ced1389549
|
|
@ -276,7 +276,14 @@ func (s *Session) LoadTool(name string, ag *agent.Agent) error {
|
|||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
return LoadToolOnConn(s.ToolsConn(), name)
|
||||
conn := s.ToolsConn()
|
||||
if conn == nil {
|
||||
return nil
|
||||
}
|
||||
// Ensure session ID is set on toolsrv before loading - the server uses this
|
||||
// to scope the tool registry. Required because conn may have been reconnected.
|
||||
conn.SetEnv("OLLIE_SESSION_ID", s.ID)
|
||||
return LoadToolOnConn(conn, name)
|
||||
}
|
||||
|
||||
// LoadToolOnConn loads a tool on the given connection.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,318 @@
|
|||
// exec9p.go - Decoupled tool execution for 9P toolsrv
|
||||
package toolsrv
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"ollie/paths"
|
||||
"ollie/sandbox"
|
||||
)
|
||||
|
||||
// ExecConfig contains all context needed for tool execution.
|
||||
// This decouples execution from Server state.
|
||||
type ExecConfig struct {
|
||||
CWD string
|
||||
Env map[string]string
|
||||
Yolo bool
|
||||
Timeout int
|
||||
}
|
||||
|
||||
// ExecuteTool runs a tool script with the given args and returns the result.
|
||||
// This is the core execution function, decoupled from Server.
|
||||
func ExecuteTool(ctx context.Context, info ToolInfo, args json.RawMessage, cfg ExecConfig) (json.RawMessage, error) {
|
||||
// Resolve script path
|
||||
toolPath, err := ResolveTool(info.Name)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("resolve tool %s: %w", info.Name, err)
|
||||
}
|
||||
|
||||
// Extract dispatch-level flags from args
|
||||
bypassed := false
|
||||
timeout := cfg.Timeout
|
||||
if timeout <= 0 {
|
||||
timeout = 30
|
||||
}
|
||||
sandboxName := "default"
|
||||
|
||||
var argMap map[string]interface{}
|
||||
if err := json.Unmarshal(args, &argMap); err == nil {
|
||||
if e, ok := argMap["bypass"]; ok {
|
||||
switch v := e.(type) {
|
||||
case bool:
|
||||
bypassed = v
|
||||
case string:
|
||||
bypassed = v == "true" || v == "1"
|
||||
}
|
||||
delete(argMap, "bypass")
|
||||
}
|
||||
if t, ok := argMap["timeout"]; ok {
|
||||
switch v := t.(type) {
|
||||
case float64:
|
||||
timeout = int(v)
|
||||
case string:
|
||||
var n int
|
||||
if _, err := fmt.Sscanf(v, "%d", &n); err == nil {
|
||||
timeout = n
|
||||
}
|
||||
}
|
||||
delete(argMap, "timeout")
|
||||
}
|
||||
if s, ok := argMap["sandbox"]; ok {
|
||||
if v, ok := s.(string); ok && v != "" {
|
||||
sandboxName = v
|
||||
}
|
||||
delete(argMap, "sandbox")
|
||||
}
|
||||
args, _ = json.Marshal(argMap)
|
||||
}
|
||||
|
||||
// Check if tool requires sudo
|
||||
needsSudo := false
|
||||
if m, err := LoadMetaFile(info.Name); err == nil && m != nil {
|
||||
resolved := m.Resolve()
|
||||
if resolved != nil && resolved.Sudo {
|
||||
needsSudo = true
|
||||
bypassed = true
|
||||
}
|
||||
}
|
||||
|
||||
stdinData := string(args)
|
||||
cwd := cfg.CWD
|
||||
if cwd == "" {
|
||||
cwd, _ = os.Getwd()
|
||||
}
|
||||
|
||||
var result string
|
||||
if needsSudo {
|
||||
sudoCode := fmt.Sprintf("cat <<'OLLIE_EOF' | %s\n%s\nOLLIE_EOF", toolPath, stdinData)
|
||||
result, err = executeBypassDirect(ctx, sudoCode, cwd, cfg.Env, timeout, true)
|
||||
} else if bypassed {
|
||||
bypassCode := fmt.Sprintf("cat <<'OLLIE_EOF' | %s\n%s\nOLLIE_EOF", toolPath, stdinData)
|
||||
result, err = executeBypassDirect(ctx, bypassCode, cwd, cfg.Env, timeout, false)
|
||||
} else {
|
||||
result, err = executeSandboxed(ctx, toolPath, stdinData, cwd, cfg.Env, timeout, sandboxName, cfg.Yolo)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return json.Marshal(ToolResult{
|
||||
IsError: true,
|
||||
Content: []ToolResultContent{{Type: "text", Text: result + ": " + err.Error()}},
|
||||
})
|
||||
}
|
||||
return json.Marshal(ToolResult{
|
||||
Content: []ToolResultContent{{Type: "text", Text: result}},
|
||||
})
|
||||
}
|
||||
|
||||
// executeSandboxed runs a tool script inside the sandbox.
|
||||
func executeSandboxed(ctx context.Context, toolPath, stdinData, cwd string, envExtra map[string]string, timeout int, sandboxName string, yolo bool) (string, error) {
|
||||
if timeout <= 0 {
|
||||
timeout = 30
|
||||
}
|
||||
|
||||
var sandboxCfg *sandbox.Config
|
||||
if !yolo {
|
||||
cfgPath := filepath.Join(paths.CfgDir(), "sandbox", sandboxName+".yaml")
|
||||
f, err := os.Open(cfgPath)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("sandbox %q not found: %w", sandboxName, err)
|
||||
}
|
||||
defer f.Close()
|
||||
var loadErr error
|
||||
sandboxCfg, loadErr = sandbox.LoadSandbox(f)
|
||||
if loadErr != nil {
|
||||
return "", loadErr
|
||||
}
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
|
||||
defer cancel()
|
||||
|
||||
interpreter := []string{"bash", "-c", toolPath}
|
||||
|
||||
// Build environment
|
||||
envMap := make(map[string]string)
|
||||
for _, ev := range os.Environ() {
|
||||
k, v, _ := strings.Cut(ev, "=")
|
||||
envMap[k] = v
|
||||
}
|
||||
for k, v := range envExtra {
|
||||
envMap[k] = v
|
||||
}
|
||||
// Compute NAMESPACE if not set
|
||||
if _, ok := envMap["NAMESPACE"]; !ok {
|
||||
if ns := plan9Namespace(envMap); ns != "" {
|
||||
envMap["NAMESPACE"] = ns
|
||||
}
|
||||
}
|
||||
getenv := func(key string) string { return envMap[key] }
|
||||
|
||||
var cmd *exec.Cmd
|
||||
if yolo {
|
||||
cmd = exec.CommandContext(ctx, interpreter[0], interpreter[1:]...)
|
||||
} else {
|
||||
wrapped, wrapErr := sandbox.WrapCommand(sandboxCfg, interpreter, cwd, getenv)
|
||||
if wrapErr != nil {
|
||||
return "", wrapErr
|
||||
}
|
||||
cmd = exec.CommandContext(ctx, wrapped[0], wrapped[1:]...)
|
||||
}
|
||||
|
||||
cmd.Dir = cwd
|
||||
if stdinData != "" {
|
||||
cmd.Stdin = strings.NewReader(stdinData)
|
||||
}
|
||||
|
||||
// Set environment
|
||||
for k, v := range envMap {
|
||||
cmd.Env = append(cmd.Env, k+"="+v)
|
||||
}
|
||||
|
||||
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||||
cmd.Cancel = func() error {
|
||||
if cmd.Process != nil {
|
||||
syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
cmd.WaitDelay = time.Second
|
||||
|
||||
var outputBuf bytes.Buffer
|
||||
lw := &limitedWriter{
|
||||
w: &outputBuf,
|
||||
limit: 10 * 1024 * 1024,
|
||||
stream: StreamFunc(ctx),
|
||||
}
|
||||
cmd.Stdout = lw
|
||||
cmd.Stderr = lw
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
return "", fmt.Errorf("execution failed: %w", err)
|
||||
}
|
||||
|
||||
err := cmd.Wait()
|
||||
output := outputBuf.Bytes()
|
||||
if lw.truncated {
|
||||
output = append(output, []byte("\n[output truncated at 10MB]")...)
|
||||
}
|
||||
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
return string(output), fmt.Errorf("execution timeout after %d seconds", timeout)
|
||||
}
|
||||
if err != nil {
|
||||
errOutput := string(output)
|
||||
const maxErrOutput = 8192
|
||||
if len(errOutput) > maxErrOutput {
|
||||
errOutput = errOutput[:maxErrOutput] + fmt.Sprintf("\n[...truncated %d bytes in error]", len(output)-maxErrOutput)
|
||||
}
|
||||
return string(output), fmt.Errorf("execution failed: %w\nOutput: %s", err, errOutput)
|
||||
}
|
||||
return string(output), nil
|
||||
}
|
||||
|
||||
// executeBypassDirect runs a command via the bypass broker.
|
||||
func executeBypassDirect(ctx context.Context, cmd, cwd string, envExtra map[string]string, timeout int, sudo bool) (string, error) {
|
||||
xdg := os.Getenv("XDG_RUNTIME_DIR")
|
||||
if xdg == "" {
|
||||
return "", fmt.Errorf("bypass not available: no XDG_RUNTIME_DIR")
|
||||
}
|
||||
sockPath := filepath.Join(xdg, "ollie", "bypass.sock")
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// Build environment map
|
||||
envMap := make(map[string]string)
|
||||
for _, kv := range os.Environ() {
|
||||
if k, v, ok := strings.Cut(kv, "="); ok {
|
||||
envMap[k] = v
|
||||
}
|
||||
}
|
||||
for k, v := range envExtra {
|
||||
envMap[k] = v
|
||||
}
|
||||
|
||||
// Connect and send request
|
||||
conn, err := net.DialTimeout("unix", sockPath, 5*time.Second)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("bypass not available: %w", err)
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
reqJSON, _ := json.Marshal(struct {
|
||||
Cmd string `json:"cmd"`
|
||||
Cwd string `json:"cwd"`
|
||||
Env map[string]string `json:"env"`
|
||||
Sudo bool `json:"sudo,omitempty"`
|
||||
}{Cmd: cmd, Cwd: cwd, Env: envMap, Sudo: sudo})
|
||||
reqJSON = append(reqJSON, '\n')
|
||||
|
||||
if _, err := conn.Write(reqJSON); err != nil {
|
||||
return "", fmt.Errorf("bypass write failed: %w", err)
|
||||
}
|
||||
|
||||
// Read response frames
|
||||
var outputBuf bytes.Buffer
|
||||
exitCode := readBypassFrames(ctx, conn, &outputBuf, StreamFunc(ctx))
|
||||
|
||||
output := outputBuf.String()
|
||||
if exitCode != 0 {
|
||||
return output, fmt.Errorf("bypass execution failed (exit %d)", exitCode)
|
||||
}
|
||||
return output, nil
|
||||
}
|
||||
|
||||
// readBypassFrames reads framed output from bypass broker.
|
||||
func readBypassFrames(ctx context.Context, conn net.Conn, w *bytes.Buffer, stream func(string)) int {
|
||||
header := make([]byte, 5)
|
||||
for {
|
||||
conn.SetReadDeadline(time.Now().Add(1 * time.Second))
|
||||
_, err := io.ReadFull(conn, header)
|
||||
if err != nil {
|
||||
if ne, ok := err.(net.Error); ok && ne.Timeout() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return -1
|
||||
default:
|
||||
continue
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
frameType := header[0]
|
||||
length := binary.BigEndian.Uint32(header[1:5])
|
||||
|
||||
payload := make([]byte, length)
|
||||
if length > 0 {
|
||||
conn.SetReadDeadline(time.Now().Add(30 * time.Second))
|
||||
if _, err := io.ReadFull(conn, payload); err != nil {
|
||||
return -1
|
||||
}
|
||||
}
|
||||
|
||||
switch frameType {
|
||||
case 'd':
|
||||
w.Write(payload)
|
||||
if stream != nil {
|
||||
stream(string(payload))
|
||||
}
|
||||
case 'x':
|
||||
var code int
|
||||
fmt.Sscanf(string(payload), "%d", &code)
|
||||
return code
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,561 @@
|
|||
// 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 toolsrv
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 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() []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 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 \\
|
||||
value = strings.ReplaceAll(value, "\\n", "\n")
|
||||
value = strings.ReplaceAll(value, "\\\\", "\\")
|
||||
args[key] = value
|
||||
}
|
||||
return args
|
||||
}
|
||||
|
||||
// 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])
|
||||
default:
|
||||
return fmt.Errorf("unknown command: %s", parts[0])
|
||||
}
|
||||
}
|
||||
|
||||
// HandleTools processes reads/writes to /tools.
|
||||
func (fs *FS9P) HandleToolsRead() string {
|
||||
var sb strings.Builder
|
||||
for _, t := range fs.ListTools() {
|
||||
if t.Description != "" {
|
||||
fmt.Fprintf(&sb, "%s\t%s\n", t.Name, t.Description)
|
||||
} else {
|
||||
sb.WriteString(t.Name + "\n")
|
||||
}
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// 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)
|
||||
}
|
||||
|
|
@ -0,0 +1,216 @@
|
|||
package toolsrv
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestFS9P_Basic(t *testing.T) {
|
||||
fs := NewFS9P("/tmp")
|
||||
|
||||
// Test CWD
|
||||
if fs.CWD() != "/tmp" {
|
||||
t.Errorf("CWD() = %q, want /tmp", fs.CWD())
|
||||
}
|
||||
|
||||
fs.SetCWD("/home")
|
||||
if fs.CWD() != "/home" {
|
||||
t.Errorf("CWD() = %q, want /home", fs.CWD())
|
||||
}
|
||||
|
||||
// Test Env
|
||||
fs.SetEnv("FOO", "bar")
|
||||
if fs.GetEnv("FOO") != "bar" {
|
||||
t.Errorf("GetEnv(FOO) = %q, want bar", fs.GetEnv("FOO"))
|
||||
}
|
||||
|
||||
// Test HostInfo
|
||||
info := fs.HostInfo9P()
|
||||
if !strings.Contains(info, "platform=") {
|
||||
t.Errorf("HostInfo9P() missing platform")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFS9P_ToolsWithoutRegistry(t *testing.T) {
|
||||
fs := NewFS9P("/tmp")
|
||||
|
||||
// Without registry, ListTools should return nil
|
||||
tools := fs.ListTools()
|
||||
if tools != nil {
|
||||
t.Errorf("ListTools() without registry = %v, want nil", tools)
|
||||
}
|
||||
|
||||
// Load should fail without registry
|
||||
err := fs.LoadTool("shell")
|
||||
if err == nil {
|
||||
t.Error("LoadTool without registry should fail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFS9P_HandleCtl(t *testing.T) {
|
||||
fs := NewFS9P("/tmp")
|
||||
|
||||
// Unknown command
|
||||
err := fs.HandleCtl("unknown")
|
||||
if err == nil {
|
||||
t.Error("HandleCtl(unknown) should fail")
|
||||
}
|
||||
|
||||
// load without name
|
||||
err = fs.HandleCtl("load")
|
||||
if err == nil {
|
||||
t.Error("HandleCtl(load) without name should fail")
|
||||
}
|
||||
|
||||
// unload without name
|
||||
err = fs.HandleCtl("unload")
|
||||
if err == nil {
|
||||
t.Error("HandleCtl(unload) without name should fail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFS9P_ProcLimit(t *testing.T) {
|
||||
fs := NewFS9P("/tmp")
|
||||
fs.procLimit = 2
|
||||
|
||||
// Create a fake registry and session ID
|
||||
reg, err := NewRegistry()
|
||||
if err != nil {
|
||||
t.Skip("no tools directory available")
|
||||
}
|
||||
fs.SetRegistry(reg, "test-session")
|
||||
|
||||
// Load a simple tool if available
|
||||
if err := reg.Load("test-session", "shell"); err != nil {
|
||||
t.Skip("shell tool not available")
|
||||
}
|
||||
|
||||
// Proc limit test would require actual tool execution
|
||||
// For now just verify the limit is set
|
||||
if fs.procLimit != 2 {
|
||||
t.Errorf("procLimit = %d, want 2", fs.procLimit)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFS9P_ParsePayload(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
want map[string]string
|
||||
}{
|
||||
{
|
||||
input: "tool=shell\ncmd=ls -la\n",
|
||||
want: map[string]string{"tool": "shell", "cmd": "ls -la"},
|
||||
},
|
||||
{
|
||||
input: "key=value with spaces\n",
|
||||
want: map[string]string{"key": "value with spaces"},
|
||||
},
|
||||
{
|
||||
input: "multi=line\\none\\ntwo\n",
|
||||
want: map[string]string{"multi": "line\none\ntwo"},
|
||||
},
|
||||
{
|
||||
input: "escaped=back\\\\slash\n",
|
||||
want: map[string]string{"escaped": "back\\slash"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
got := parsePayload(tt.input)
|
||||
for k, v := range tt.want {
|
||||
if got[k] != v {
|
||||
t.Errorf("parsePayload(%q)[%q] = %q, want %q", tt.input, k, got[k], v)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFS9P_ProcLifecycle(t *testing.T) {
|
||||
fs := NewFS9P("/tmp")
|
||||
|
||||
// No procs initially
|
||||
pids := fs.ListProcs()
|
||||
if len(pids) != 0 {
|
||||
t.Errorf("ListProcs() = %v, want empty", pids)
|
||||
}
|
||||
|
||||
// GetProc for non-existent
|
||||
if fs.GetProc(999) != nil {
|
||||
t.Error("GetProc(999) should return nil")
|
||||
}
|
||||
|
||||
// DismissProc for non-existent
|
||||
if fs.DismissProc(999) {
|
||||
t.Error("DismissProc(999) should return false")
|
||||
}
|
||||
|
||||
// SignalProc for non-existent
|
||||
if err := fs.SignalProc(999, 15); err == nil {
|
||||
t.Error("SignalProc(999) should fail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFS9P_HandleProcCtl(t *testing.T) {
|
||||
fs := NewFS9P("/tmp")
|
||||
|
||||
// Unknown command
|
||||
err := fs.HandleProcCtl(1, "unknown")
|
||||
if err == nil {
|
||||
t.Error("HandleProcCtl(unknown) should fail")
|
||||
}
|
||||
|
||||
// signal without number
|
||||
err = fs.HandleProcCtl(1, "signal")
|
||||
if err == nil {
|
||||
t.Error("HandleProcCtl(signal) without number should fail")
|
||||
}
|
||||
|
||||
// signal with invalid number
|
||||
err = fs.HandleProcCtl(1, "signal abc")
|
||||
if err == nil {
|
||||
t.Error("HandleProcCtl(signal abc) should fail")
|
||||
}
|
||||
|
||||
// dismiss on non-existent (should succeed silently)
|
||||
err = fs.HandleProcCtl(999, "dismiss")
|
||||
if err != nil {
|
||||
t.Errorf("HandleProcCtl(dismiss) on non-existent = %v, want nil", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProc9P_Stat(t *testing.T) {
|
||||
proc := &Proc9P{
|
||||
PID: 1,
|
||||
Tool: "shell",
|
||||
StartTime: time.Now(),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
|
||||
// Running state
|
||||
stat := proc.Stat()
|
||||
if !strings.Contains(stat, "running") {
|
||||
t.Errorf("Stat() for running proc should contain 'running', got %q", stat)
|
||||
}
|
||||
|
||||
// Mark as exited
|
||||
proc.Exited = true
|
||||
proc.ExitCode = 0
|
||||
proc.EndTime = time.Now()
|
||||
|
||||
stat = proc.Stat()
|
||||
if !strings.Contains(stat, "exited 0") {
|
||||
t.Errorf("Stat() for exited proc should contain 'exited 0', got %q", stat)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFS9P_NewProcWithoutRegistry(t *testing.T) {
|
||||
fs := NewFS9P("/tmp")
|
||||
|
||||
ctx := context.Background()
|
||||
_, _, err := fs.NewProc(ctx, "tool=shell\ncmd=ls\n", false)
|
||||
if err == nil {
|
||||
t.Error("NewProc without registry should fail")
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue