diff --git a/cmd/ollie-remote/main.go b/cmd/ollie-remote/main.go deleted file mode 100644 index 89ebacf..0000000 --- a/cmd/ollie-remote/main.go +++ /dev/null @@ -1,93 +0,0 @@ -// ollie-remote — tool execution server. -// -// Accepts JSON-RPC 2.0 connections. Executes tool calls in a sandboxed -// environment. Used both locally (spawned by olliesrv) and remotely -// (deployed via SSH bootstrap). -// -// Usage: -// -// ollie-remote serve --cwd /path/to/project [--listen /path/to/sock] [--yolo] -// -// When --listen is provided, accepts connections on a Unix socket. -// Otherwise, serves a single session over stdin/stdout (bootstrap mode). -package main - -import ( - "context" - "flag" - "fmt" - "os" - "os/signal" - "path/filepath" - "syscall" - - "ollie/env" - "ollie/toolsrv" -) - -var ( - cwd = flag.String("cwd", ".", "working directory for execution") - listenPath = flag.String("listen", "", "Unix socket path to listen on (default: stdio mode)") - yolo = flag.Bool("yolo", false, "skip sandbox enforcement") -) - -func main() { - if len(os.Args) < 2 || os.Args[1] != "serve" { - fmt.Fprintln(os.Stderr, "usage: ollie-remote serve --cwd [--listen ]") - os.Exit(1) - } - flag.CommandLine.Parse(os.Args[2:]) - - env.EnsureDefaults() - - // Prepend our bin dir to PATH so landrun is found (deployed alongside us). - home, _ := os.UserHomeDir() - binDir := filepath.Join(home, ".config", "ollie", "bin") - os.Setenv("PATH", binDir+":"+os.Getenv("PATH")) - - toolReg, _ := toolsrv.NewRegistry() - sessionID := os.Getenv("OLLIE_SESSION_ID") - - var opts []toolsrv.Option - if *yolo { - opts = append(opts, toolsrv.WithYolo()) - } - if toolReg != nil && sessionID != "" { - opts = append(opts, toolsrv.WithToolRegistry(toolReg, sessionID)) - } - server := toolsrv.New(*cwd) - for _, o := range opts { - o(server) - } - - // Wire up environment propagation for late-arriving session IDs. - server.OnEnvSet = func(key, value string) { - switch key { - case "OLLIE_SESSION_ID": - os.Setenv("OLLIE_SESSION_ID", value) - if toolReg != nil { - server.SetToolRegistry(toolReg, value) - } - case "OLLIE_UNAME": - os.Setenv("OLLIE_UNAME", value) - } - } - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - sigCh := make(chan os.Signal, 1) - signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) - go func() { - <-sigCh - cancel() - }() - - if *listenPath != "" { - if err := server.ServeSocket(ctx, *listenPath); err != nil { - fmt.Fprintf(os.Stderr, "%v\n", err) - os.Exit(1) - } - } else { - server.ServeRPC(ctx, os.Stdin, os.Stdout) - } -} diff --git a/cmd/olliesrv/bypass_notify.go b/cmd/olliesrv/bypass_notify.go index c025f8e..22780b5 100644 --- a/cmd/olliesrv/bypass_notify.go +++ b/cmd/olliesrv/bypass_notify.go @@ -5,7 +5,7 @@ import ( "sync" "github.com/godbus/dbus/v5" - "ollie/bypass" + "ollie/cmd/olliesrv/internal/bypass" ) // bypassNotifier handles desktop notifications for bypass requests diff --git a/agent/agent.go b/cmd/olliesrv/internal/agent/agent.go similarity index 99% rename from agent/agent.go rename to cmd/olliesrv/internal/agent/agent.go index e483ce4..d440072 100644 --- a/agent/agent.go +++ b/cmd/olliesrv/internal/agent/agent.go @@ -9,7 +9,7 @@ import ( "sync/atomic" "syscall" - "ollie/backend" + "ollie/cmd/olliesrv/internal/backend" olog "ollie/log" "ollie/paths" "ollie/toolsrv" @@ -387,7 +387,7 @@ const ( // SetSessionEnv injects session env vars into the execute server. func (ag *Agent) SetSessionEnv(sessionID string) { - if ag.runtime.ToolServer == nil { + if ag.runtime == nil || ag.runtime.ToolServer == nil { return } ag.runtime.ToolServer.SetEnv("OLLIE_SESSION_ID", sessionID) diff --git a/agent/agent_config.go b/cmd/olliesrv/internal/agent/agent_config.go similarity index 97% rename from agent/agent_config.go rename to cmd/olliesrv/internal/agent/agent_config.go index a75688a..e362c1c 100644 --- a/agent/agent_config.go +++ b/cmd/olliesrv/internal/agent/agent_config.go @@ -4,7 +4,7 @@ import ( "encoding/json" "io" - "ollie/backend" + "ollie/cmd/olliesrv/internal/backend" ) // Prompt holds the agent prompt — a list of file paths to read and concatenate. diff --git a/agent/agent_test.go b/cmd/olliesrv/internal/agent/agent_test.go similarity index 50% rename from agent/agent_test.go rename to cmd/olliesrv/internal/agent/agent_test.go index 860e235..693b944 100644 --- a/agent/agent_test.go +++ b/cmd/olliesrv/internal/agent/agent_test.go @@ -1,52 +1,12 @@ package agent import ( - "encoding/json" - "io" "strings" "testing" "ollie/toolsrv" ) -// mockToolServer creates a Conn backed by a goroutine that responds to -// list_tools with the given ToolInfo slice. Simulates the ollie-remote RPC. -func mockToolServer(t *testing.T, infos []toolsrv.ToolInfo) *toolsrv.Conn { - t.Helper() - cr, sw := io.Pipe() // client reads, server writes - sr, cw := io.Pipe() // server reads, client writes - - go func() { - dec := json.NewDecoder(sr) - enc := json.NewEncoder(sw) - for { - var req struct { - JSONRPC string `json:"jsonrpc"` - ID int64 `json:"id"` - Method string `json:"method"` - Params json.RawMessage `json:"params,omitempty"` - } - if err := dec.Decode(&req); err != nil { - return - } - switch req.Method { - case "list_tools": - result, _ := json.Marshal(infos) - enc.Encode(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": json.RawMessage(result)}) - default: - enc.Encode(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": true}) - } - } - }() - - rwc := struct { - io.Reader - io.Writer - io.Closer - }{cr, cw, cw} - return toolsrv.NewConn(rwc, nil) -} - func TestRefreshToolListing(t *testing.T) { infos := []toolsrv.ToolInfo{ { @@ -56,20 +16,17 @@ func TestRefreshToolListing(t *testing.T) { }, } - conn := mockToolServer(t, infos) - defer conn.Close() - p := &Preamble{} p.Set(SectionSystem, "# System") p.Set(SectionEnv, "# Env") p.Set(SectionAgent, "# Agent") p.Set(SectionTools, "# Tools\n\n## old_tool\n\nOld description\n\n") - ag := &Agent{runtime: &Runtime{Preamble: p, ToolServer: conn}} + // Test renderTools directly since refreshToolListing depends on ToolServer + rendered := renderTools(infos) + p.Set(SectionTools, rendered) - ag.refreshToolListing() - - preamble := ag.runtime.PreambleString() + preamble := p.String() if !strings.Contains(preamble, "## new_tool") { t.Error("should contain new tool") diff --git a/agent/commands.go b/cmd/olliesrv/internal/agent/commands.go similarity index 100% rename from agent/commands.go rename to cmd/olliesrv/internal/agent/commands.go diff --git a/agent/config_paths.go b/cmd/olliesrv/internal/agent/config_paths.go similarity index 100% rename from agent/config_paths.go rename to cmd/olliesrv/internal/agent/config_paths.go diff --git a/agent/cost.go b/cmd/olliesrv/internal/agent/cost.go similarity index 98% rename from agent/cost.go rename to cmd/olliesrv/internal/agent/cost.go index d0686ac..0ad48a1 100644 --- a/agent/cost.go +++ b/cmd/olliesrv/internal/agent/cost.go @@ -3,7 +3,7 @@ package agent import ( "strings" - "ollie/backend" + "ollie/cmd/olliesrv/internal/backend" ) // auditTruncate trims s to 200 runes for log output. diff --git a/agent/fifo.go b/cmd/olliesrv/internal/agent/fifo.go similarity index 100% rename from agent/fifo.go rename to cmd/olliesrv/internal/agent/fifo.go diff --git a/agent/fifo_test.go b/cmd/olliesrv/internal/agent/fifo_test.go similarity index 100% rename from agent/fifo_test.go rename to cmd/olliesrv/internal/agent/fifo_test.go diff --git a/agent/history.go b/cmd/olliesrv/internal/agent/history.go similarity index 99% rename from agent/history.go rename to cmd/olliesrv/internal/agent/history.go index 7d5c32e..0010c87 100644 --- a/agent/history.go +++ b/cmd/olliesrv/internal/agent/history.go @@ -7,7 +7,7 @@ import ( "os" "strings" - "ollie/backend" + "ollie/cmd/olliesrv/internal/backend" "ollie/paths" "gopkg.in/yaml.v3" diff --git a/agent/loop.go b/cmd/olliesrv/internal/agent/loop.go similarity index 95% rename from agent/loop.go rename to cmd/olliesrv/internal/agent/loop.go index 376399c..c554bd2 100644 --- a/agent/loop.go +++ b/cmd/olliesrv/internal/agent/loop.go @@ -10,7 +10,7 @@ import ( "sync" "time" - "ollie/backend" + "ollie/cmd/olliesrv/internal/backend" "ollie/paths" "ollie/toolsrv" ) @@ -481,26 +481,34 @@ func execOne(rt *Runtime, ctx TurnCtx, call backend.ToolCall, resultCache *sync. var result string var resultBlocks []backend.ContentBlock var isErr bool - streamed := false + + // Built-in: tool_load - must be handled specially because tools run inside + // toolsrv and can't load other tools into themselves. + if call.Name == "tool_load" { + var args struct { + Name string `json:"name"` + } + if err := json.Unmarshal(call.Arguments, &args); err != nil { + result = "error: " + err.Error() + isErr = true + } else if args.Name == "" { + result = "error: name is required" + isErr = true + } else if rt.ToolServer == nil { + result = "error: no tool server" + isErr = true + } else if err := rt.ToolServer.LoadTool(args.Name); err != nil { + result = "error: " + err.Error() + isErr = true + } else { + result = "loaded: " + args.Name + } + emit(ctx, Event{Role: "tool", Name: call.Name, Content: result, OutputFormat: toolOutputFormat(rt, call.Name)}) + return toolResult{ToolCallID: call.ID, Name: call.Name, Content: result, IsError: isErr}, false + } if rt.Exec != nil { - streamBytes := 0 - streamCtx := toolsrv.WithOutputStream(ctx, func(data string) { - if streamBytes >= defaultToolResultMaxBytes { - return - } - if !streamed { - emit(ctx, Event{Role: "tool", Name: call.Name, OutputFormat: toolOutputFormat(rt, call.Name)}) - streamed = true - } - streamBytes += len(data) - if streamBytes > defaultToolResultMaxBytes { - excess := streamBytes - defaultToolResultMaxBytes - data = data[:len(data)-excess] - } - emit(ctx, Event{Role: "tool", Name: call.Name, Content: data, OutputFormat: toolOutputFormat(rt, call.Name)}) - }) - out, blocks, err := rt.Exec(streamCtx, call.Name, call.Arguments) + out, blocks, err := rt.Exec(ctx, call.Name, call.Arguments) if err != nil { isErr = true if ctx.Err() != nil { @@ -555,14 +563,7 @@ func execOne(rt *Runtime, ctx TurnCtx, call backend.ToolCall, resultCache *sync. }) } - if !streamed { - emit(ctx, Event{Role: "tool", Name: call.Name, Content: result, OutputFormat: toolOutputFormat(rt, call.Name)}) - } else { - if suffix != "" { - emit(ctx, Event{Role: "tool", Name: call.Name, Content: suffix, OutputFormat: toolOutputFormat(rt, call.Name)}) - } - emit(ctx, Event{Role: "tool", Name: call.Name, OutputFormat: toolOutputFormat(rt, call.Name)}) - } + emit(ctx, Event{Role: "tool", Name: call.Name, Content: result, OutputFormat: toolOutputFormat(rt, call.Name)}) tier := MemoryHot if !isErr { diff --git a/agent/prompt_resolver.go b/cmd/olliesrv/internal/agent/prompt_resolver.go similarity index 100% rename from agent/prompt_resolver.go rename to cmd/olliesrv/internal/agent/prompt_resolver.go diff --git a/agent/runtime.go b/cmd/olliesrv/internal/agent/runtime.go similarity index 98% rename from agent/runtime.go rename to cmd/olliesrv/internal/agent/runtime.go index ef717ca..b69e6bc 100644 --- a/agent/runtime.go +++ b/cmd/olliesrv/internal/agent/runtime.go @@ -8,8 +8,8 @@ import ( "strings" "time" - "ollie/backend" - "ollie/prompts" + "ollie/cmd/olliesrv/internal/backend" + "ollie/cmd/olliesrv/internal/prompts" "ollie/toolsrv" ) diff --git a/agent/runtime_test.go b/cmd/olliesrv/internal/agent/runtime_test.go similarity index 100% rename from agent/runtime_test.go rename to cmd/olliesrv/internal/agent/runtime_test.go diff --git a/agent/text_parse.go b/cmd/olliesrv/internal/agent/text_parse.go similarity index 99% rename from agent/text_parse.go rename to cmd/olliesrv/internal/agent/text_parse.go index 3c198eb..e80ab47 100644 --- a/agent/text_parse.go +++ b/cmd/olliesrv/internal/agent/text_parse.go @@ -5,7 +5,7 @@ import ( "fmt" "strings" - "ollie/backend" + "ollie/cmd/olliesrv/internal/backend" ) // parseTextToolCalls extracts tool calls from assistant text when the model diff --git a/agent/turn.go b/cmd/olliesrv/internal/agent/turn.go similarity index 99% rename from agent/turn.go rename to cmd/olliesrv/internal/agent/turn.go index 83b47ae..cca9a2f 100644 --- a/agent/turn.go +++ b/cmd/olliesrv/internal/agent/turn.go @@ -9,7 +9,7 @@ import ( "strconv" "strings" - "ollie/backend" + "ollie/cmd/olliesrv/internal/backend" "ollie/toolsrv" ) diff --git a/backend/anthropic.go b/cmd/olliesrv/internal/backend/anthropic.go similarity index 100% rename from backend/anthropic.go rename to cmd/olliesrv/internal/backend/anthropic.go diff --git a/backend/anthropic_stream_test.go b/cmd/olliesrv/internal/backend/anthropic_stream_test.go similarity index 100% rename from backend/anthropic_stream_test.go rename to cmd/olliesrv/internal/backend/anthropic_stream_test.go diff --git a/backend/backend.go b/cmd/olliesrv/internal/backend/backend.go similarity index 100% rename from backend/backend.go rename to cmd/olliesrv/internal/backend/backend.go diff --git a/backend/backend_test.go b/cmd/olliesrv/internal/backend/backend_test.go similarity index 99% rename from backend/backend_test.go rename to cmd/olliesrv/internal/backend/backend_test.go index 9c5c9fa..d555995 100644 --- a/backend/backend_test.go +++ b/cmd/olliesrv/internal/backend/backend_test.go @@ -12,7 +12,7 @@ import ( "testing" "time" - "ollie/backend" + "ollie/cmd/olliesrv/internal/backend" ) func mustNewOpenAI(t *testing.T, name, baseURL, apiKey string) *backend.OpenAIBackend { diff --git a/backend/base.go b/cmd/olliesrv/internal/backend/base.go similarity index 100% rename from backend/base.go rename to cmd/olliesrv/internal/backend/base.go diff --git a/backend/codewhisperer_test.go b/cmd/olliesrv/internal/backend/codewhisperer_test.go similarity index 100% rename from backend/codewhisperer_test.go rename to cmd/olliesrv/internal/backend/codewhisperer_test.go diff --git a/backend/config.go b/cmd/olliesrv/internal/backend/config.go similarity index 100% rename from backend/config.go rename to cmd/olliesrv/internal/backend/config.go diff --git a/backend/copilot.go b/cmd/olliesrv/internal/backend/copilot.go similarity index 100% rename from backend/copilot.go rename to cmd/olliesrv/internal/backend/copilot.go diff --git a/backend/export_test.go b/cmd/olliesrv/internal/backend/export_test.go similarity index 100% rename from backend/export_test.go rename to cmd/olliesrv/internal/backend/export_test.go diff --git a/backend/gemini.go b/cmd/olliesrv/internal/backend/gemini.go similarity index 100% rename from backend/gemini.go rename to cmd/olliesrv/internal/backend/gemini.go diff --git a/backend/kiro.go b/cmd/olliesrv/internal/backend/kiro.go similarity index 100% rename from backend/kiro.go rename to cmd/olliesrv/internal/backend/kiro.go diff --git a/backend/kiro_internal.go b/cmd/olliesrv/internal/backend/kiro_internal.go similarity index 100% rename from backend/kiro_internal.go rename to cmd/olliesrv/internal/backend/kiro_internal.go diff --git a/backend/new.go b/cmd/olliesrv/internal/backend/new.go similarity index 100% rename from backend/new.go rename to cmd/olliesrv/internal/backend/new.go diff --git a/backend/new_test.go b/cmd/olliesrv/internal/backend/new_test.go similarity index 100% rename from backend/new_test.go rename to cmd/olliesrv/internal/backend/new_test.go diff --git a/backend/ollama.go b/cmd/olliesrv/internal/backend/ollama.go similarity index 100% rename from backend/ollama.go rename to cmd/olliesrv/internal/backend/ollama.go diff --git a/backend/ollama_stream_test.go b/cmd/olliesrv/internal/backend/ollama_stream_test.go similarity index 100% rename from backend/ollama_stream_test.go rename to cmd/olliesrv/internal/backend/ollama_stream_test.go diff --git a/backend/oneshot.go b/cmd/olliesrv/internal/backend/oneshot.go similarity index 100% rename from backend/oneshot.go rename to cmd/olliesrv/internal/backend/oneshot.go diff --git a/backend/openai.go b/cmd/olliesrv/internal/backend/openai.go similarity index 100% rename from backend/openai.go rename to cmd/olliesrv/internal/backend/openai.go diff --git a/backend/openai_encode_test.go b/cmd/olliesrv/internal/backend/openai_encode_test.go similarity index 100% rename from backend/openai_encode_test.go rename to cmd/olliesrv/internal/backend/openai_encode_test.go diff --git a/backend/openai_stream_test.go b/cmd/olliesrv/internal/backend/openai_stream_test.go similarity index 100% rename from backend/openai_stream_test.go rename to cmd/olliesrv/internal/backend/openai_stream_test.go diff --git a/backend/openai_test.go b/cmd/olliesrv/internal/backend/openai_test.go similarity index 100% rename from backend/openai_test.go rename to cmd/olliesrv/internal/backend/openai_test.go diff --git a/backend/pool.go b/cmd/olliesrv/internal/backend/pool.go similarity index 100% rename from backend/pool.go rename to cmd/olliesrv/internal/backend/pool.go diff --git a/backend/testhelper_test.go b/cmd/olliesrv/internal/backend/testhelper_test.go similarity index 100% rename from backend/testhelper_test.go rename to cmd/olliesrv/internal/backend/testhelper_test.go diff --git a/bypass/broker.go b/cmd/olliesrv/internal/bypass/broker.go similarity index 100% rename from bypass/broker.go rename to cmd/olliesrv/internal/bypass/broker.go diff --git a/bypass/broker_test.go b/cmd/olliesrv/internal/bypass/broker_test.go similarity index 100% rename from bypass/broker_test.go rename to cmd/olliesrv/internal/bypass/broker_test.go diff --git a/bypass/policy.go b/cmd/olliesrv/internal/bypass/policy.go similarity index 100% rename from bypass/policy.go rename to cmd/olliesrv/internal/bypass/policy.go diff --git a/bypass/policy_test.go b/cmd/olliesrv/internal/bypass/policy_test.go similarity index 100% rename from bypass/policy_test.go rename to cmd/olliesrv/internal/bypass/policy_test.go diff --git a/bypass/request.go b/cmd/olliesrv/internal/bypass/request.go similarity index 100% rename from bypass/request.go rename to cmd/olliesrv/internal/bypass/request.go diff --git a/fs/agent_log.go b/cmd/olliesrv/internal/fs/agent_log.go similarity index 98% rename from fs/agent_log.go rename to cmd/olliesrv/internal/fs/agent_log.go index 5df7be6..c6444be 100644 --- a/fs/agent_log.go +++ b/cmd/olliesrv/internal/fs/agent_log.go @@ -3,8 +3,8 @@ package fs import ( "sync" - "ollie/agent" - "ollie/backend" + "ollie/cmd/olliesrv/internal/agent" + "ollie/cmd/olliesrv/internal/backend" "ollie/format" ) diff --git a/fs/cache.go b/cmd/olliesrv/internal/fs/cache.go similarity index 96% rename from fs/cache.go rename to cmd/olliesrv/internal/fs/cache.go index f63cfea..eca40ea 100644 --- a/fs/cache.go +++ b/cmd/olliesrv/internal/fs/cache.go @@ -5,7 +5,7 @@ import ( "strings" "sync" - "ollie/backend" + "ollie/cmd/olliesrv/internal/backend" ) // ModelCache caches the model list from all backends. diff --git a/fs/ctl.go b/cmd/olliesrv/internal/fs/ctl.go similarity index 100% rename from fs/ctl.go rename to cmd/olliesrv/internal/fs/ctl.go diff --git a/fs/handlers_agent.go b/cmd/olliesrv/internal/fs/handlers_agent.go similarity index 94% rename from fs/handlers_agent.go rename to cmd/olliesrv/internal/fs/handlers_agent.go index ae27492..65dea99 100644 --- a/fs/handlers_agent.go +++ b/cmd/olliesrv/internal/fs/handlers_agent.go @@ -8,9 +8,9 @@ import ( "strconv" "strings" - "ollie/agent" + "ollie/cmd/olliesrv/internal/agent" "ollie/paths" - "ollie/session" + "ollie/cmd/olliesrv/internal/session" ) func requestAgentNew(ctx HandlerCtx, data []byte) ([]byte, error) { @@ -412,6 +412,38 @@ func readAgentStats(ctx HandlerCtx) ([]byte, error) { return []byte(s), nil } +func readAgentTools(ctx HandlerCtx) ([]byte, error) { + ts := ctx.Agent.ToolServer() + if ts == nil { + return []byte("(no tool server)\n"), nil + } + loaded, err := ts.ListTools() + if err != nil { + return nil, err + } + var sb strings.Builder + for _, ti := range loaded { + if ti.Description != "" { + fmt.Fprintf(&sb, "%-20s %s\n", ti.Name, ti.Description) + } else { + sb.WriteString(ti.Name + "\n") + } + } + return []byte(sb.String()), nil +} + +func writeAgentTools(ctx HandlerCtx, data []byte) error { + name := strings.TrimSpace(string(data)) + if name == "" { + return nil + } + ts := ctx.Agent.ToolServer() + if ts == nil { + return fmt.Errorf("no tool server") + } + return ts.LoadTool(name) +} + func readAgentName(ctx HandlerCtx) ([]byte, error) { return []byte(ctx.Agent.Name() + "\n"), nil } diff --git a/fs/handlers_root.go b/cmd/olliesrv/internal/fs/handlers_root.go similarity index 97% rename from fs/handlers_root.go rename to cmd/olliesrv/internal/fs/handlers_root.go index 7f1ffb6..15ee4fa 100644 --- a/fs/handlers_root.go +++ b/cmd/olliesrv/internal/fs/handlers_root.go @@ -8,11 +8,11 @@ import ( "strings" "sync" - "ollie/agent" - "ollie/backend" - "ollie/bypass" + "ollie/cmd/olliesrv/internal/agent" + "ollie/cmd/olliesrv/internal/backend" + "ollie/cmd/olliesrv/internal/bypass" "ollie/fsedsl" - "ollie/session" + "ollie/cmd/olliesrv/internal/session" "ollie/toolsrv" ) diff --git a/fs/handlers_session.go b/cmd/olliesrv/internal/fs/handlers_session.go similarity index 98% rename from fs/handlers_session.go rename to cmd/olliesrv/internal/fs/handlers_session.go index 6befa0f..34b7511 100644 --- a/fs/handlers_session.go +++ b/cmd/olliesrv/internal/fs/handlers_session.go @@ -5,8 +5,8 @@ import ( "os" "strings" - "ollie/bypass" - "ollie/session" + "ollie/cmd/olliesrv/internal/bypass" + "ollie/cmd/olliesrv/internal/session" ) func readSessionIdx(_ HandlerCtx) ([]byte, error) { diff --git a/fs/lifecycle.go b/cmd/olliesrv/internal/fs/lifecycle.go similarity index 97% rename from fs/lifecycle.go rename to cmd/olliesrv/internal/fs/lifecycle.go index a6f5cca..23ea705 100644 --- a/fs/lifecycle.go +++ b/cmd/olliesrv/internal/fs/lifecycle.go @@ -3,7 +3,7 @@ package fs import ( "sync" - "ollie/session" + "ollie/cmd/olliesrv/internal/session" ) // RootState holds the root-level state for the 9P filesystem. diff --git a/fs/newroot.go b/cmd/olliesrv/internal/fs/newroot.go similarity index 96% rename from fs/newroot.go rename to cmd/olliesrv/internal/fs/newroot.go index a90b55f..b71172a 100644 --- a/fs/newroot.go +++ b/cmd/olliesrv/internal/fs/newroot.go @@ -4,10 +4,10 @@ import ( "context" "os/user" - "ollie/bypass" + "ollie/cmd/olliesrv/internal/bypass" "ollie/fsedsl" olog "ollie/log" - "ollie/session" + "ollie/cmd/olliesrv/internal/session" ) // Config carries all runtime dependencies for the filesystem tree. diff --git a/fs/replay_test.go b/cmd/olliesrv/internal/fs/replay_test.go similarity index 99% rename from fs/replay_test.go rename to cmd/olliesrv/internal/fs/replay_test.go index 4e68d9d..155b40b 100644 --- a/fs/replay_test.go +++ b/cmd/olliesrv/internal/fs/replay_test.go @@ -5,7 +5,7 @@ import ( "strings" "testing" - "ollie/backend" + "ollie/cmd/olliesrv/internal/backend" "ollie/format" ) diff --git a/fs/session_node.go b/cmd/olliesrv/internal/fs/session_node.go similarity index 97% rename from fs/session_node.go rename to cmd/olliesrv/internal/fs/session_node.go index bd8f72e..4ea4853 100644 --- a/fs/session_node.go +++ b/cmd/olliesrv/internal/fs/session_node.go @@ -6,8 +6,8 @@ import ( "sync" "time" - "ollie/agent" - "ollie/session" + "ollie/cmd/olliesrv/internal/agent" + "ollie/cmd/olliesrv/internal/session" "ollie/toolsrv" ) diff --git a/fs/spec.go b/cmd/olliesrv/internal/fs/spec.go similarity index 97% rename from fs/spec.go rename to cmd/olliesrv/internal/fs/spec.go index 6b10b31..db79757 100644 --- a/fs/spec.go +++ b/cmd/olliesrv/internal/fs/spec.go @@ -7,8 +7,8 @@ package fs import ( "context" - "ollie/agent" - "ollie/bypass" + "ollie/cmd/olliesrv/internal/agent" + "ollie/cmd/olliesrv/internal/bypass" "ollie/fsedsl" olog "ollie/log" ) @@ -273,6 +273,12 @@ FileNode("generate", 0666, Read(readAgentStats), ), + FileNode("tools", 0666, + Doc("Loaded tools. Read: list. Write: tool name to load."), + Read(readAgentTools), + Write(writeAgentTools), + ), + FileNode("name", 0666, Doc("Agent display name. Read: current name. Write: rename agent."), Read(readAgentName), diff --git a/prompts/prompts.go b/cmd/olliesrv/internal/prompts/prompts.go similarity index 100% rename from prompts/prompts.go rename to cmd/olliesrv/internal/prompts/prompts.go diff --git a/prompts/prompts_test.go b/cmd/olliesrv/internal/prompts/prompts_test.go similarity index 100% rename from prompts/prompts_test.go rename to cmd/olliesrv/internal/prompts/prompts_test.go diff --git a/prompts/system_prompt.md b/cmd/olliesrv/internal/prompts/system_prompt.md similarity index 100% rename from prompts/system_prompt.md rename to cmd/olliesrv/internal/prompts/system_prompt.md diff --git a/session/event.go b/cmd/olliesrv/internal/session/event.go similarity index 100% rename from session/event.go rename to cmd/olliesrv/internal/session/event.go diff --git a/session/persist.go b/cmd/olliesrv/internal/session/persist.go similarity index 98% rename from session/persist.go rename to cmd/olliesrv/internal/session/persist.go index 4b489d6..1c0d192 100644 --- a/session/persist.go +++ b/cmd/olliesrv/internal/session/persist.go @@ -9,8 +9,8 @@ import ( "strings" "sync" - "ollie/agent" - "ollie/backend" + "ollie/cmd/olliesrv/internal/agent" + "ollie/cmd/olliesrv/internal/backend" ) // --- Persistence --- diff --git a/session/registry.go b/cmd/olliesrv/internal/session/registry.go similarity index 98% rename from session/registry.go rename to cmd/olliesrv/internal/session/registry.go index 63a3981..b463a7d 100644 --- a/session/registry.go +++ b/cmd/olliesrv/internal/session/registry.go @@ -17,8 +17,8 @@ import ( "sync" "time" - "ollie/agent" - "ollie/backend" + "ollie/cmd/olliesrv/internal/agent" + "ollie/cmd/olliesrv/internal/backend" olog "ollie/log" "ollie/paths" ) @@ -304,6 +304,7 @@ func CreateEmpty(name, remote string) (*Session, error) { Ctx: ctx, CWD: cwd, RemoteTarget: remote, + SessionID: sessID, Yolo: pkgYolo, }) if err != nil { @@ -376,6 +377,7 @@ func buildAgent(sess *Session, p AgentParams) (*agent.Agent, error) { Ctx: sess.Ctx, CWD: p.CWD, RemoteTarget: remote, + SessionID: sessID, Yolo: pkgYolo, ReuseFrom: reuseFrom, }) diff --git a/session/session.go b/cmd/olliesrv/internal/session/session.go similarity index 93% rename from session/session.go rename to cmd/olliesrv/internal/session/session.go index e95d55f..c023e03 100644 --- a/session/session.go +++ b/cmd/olliesrv/internal/session/session.go @@ -2,7 +2,6 @@ package session import ( "context" - "encoding/json" "fmt" "os" "path/filepath" @@ -10,7 +9,8 @@ import ( "sync" "time" - "ollie/agent" + "ollie/cmd/olliesrv/internal/agent" + "ollie/cmd/olliesrv/internal/toolclient" olog "ollie/log" "ollie/paths" "ollie/toolsrv" @@ -27,8 +27,8 @@ type Session struct { Cancel context.CancelFunc // Tool server lifecycle (set during setup/teardown, not raced) - Proc *toolsrv.Process - Keeper *toolsrv.ProcessKeeper + Proc *toolclient.Process + Keeper *toolclient.ProcessKeeper Remote string // Agent state @@ -242,6 +242,7 @@ func (s *Session) Resume() error { Ctx: ctx, CWD: cwd, RemoteTarget: s.Remote, + SessionID: s.ID, Yolo: pkgYolo, }) if err != nil { @@ -276,7 +277,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. @@ -284,13 +292,7 @@ func LoadToolOnConn(conn *toolsrv.Conn, name string) error { if conn == nil { return nil } - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - args, _ := json.Marshal(map[string]string{"name": name}) - if _, err := conn.CallTool(ctx, "tool_load", json.RawMessage(args)); err != nil { - return fmt.Errorf("remote load: %w", err) - } - return nil + return conn.LoadTool(name) } // --- Autosave --- diff --git a/session/setup.go b/cmd/olliesrv/internal/session/setup.go similarity index 79% rename from session/setup.go rename to cmd/olliesrv/internal/session/setup.go index 89912aa..b49840f 100644 --- a/session/setup.go +++ b/cmd/olliesrv/internal/session/setup.go @@ -4,17 +4,18 @@ import ( "context" "fmt" - "ollie/agent" + "ollie/cmd/olliesrv/internal/agent" + "ollie/cmd/olliesrv/internal/toolclient" "ollie/paths" - "ollie/prompts" + "ollie/cmd/olliesrv/internal/prompts" "ollie/toolsrv" ) // SessionInfra holds the tool server infrastructure for a session. // Created by SetupToolServer and used by both CreateAgent and restoreSession. type SessionInfra struct { - Proc *toolsrv.Process - Keeper *toolsrv.ProcessKeeper + Proc *toolclient.Process + Keeper *toolclient.ProcessKeeper ToolsConn *toolsrv.Conn NewToolServer func() *toolsrv.Conn Platform string @@ -26,6 +27,7 @@ type ToolServerConfig struct { Ctx context.Context CWD string RemoteTarget string + SessionID string // session ID for tool registry Yolo bool // ReuseFrom, if non-nil, reuses existing infrastructure instead of spawning. @@ -35,8 +37,8 @@ type ToolServerConfig struct { // InfraConfig describes existing infrastructure to reuse. type InfraConfig struct { - Proc *toolsrv.Process - Keeper *toolsrv.ProcessKeeper + Proc *toolclient.Process + Keeper *toolclient.ProcessKeeper ToolsConn *toolsrv.Conn } @@ -75,25 +77,29 @@ func SetupToolServer(cfg ToolServerConfig) (*SessionInfra, error) { IsGitRepo: paths.IsGitRepo(cfg.CWD), }, nil } - var proc *toolsrv.Process - var keeper *toolsrv.ProcessKeeper + var proc *toolclient.Process + var keeper *toolclient.ProcessKeeper var newToolServer func() *toolsrv.Conn var err error var platform string var isGitRepo bool if cfg.RemoteTarget != "" { - proc, err = toolsrv.SpawnRemote(cfg.Ctx, toolsrv.RemoteConfig{ + proc, err = toolclient.SpawnRemote(cfg.Ctx, toolclient.RemoteConfig{ SSHTarget: cfg.RemoteTarget, CWD: cfg.CWD, + SessionID: cfg.SessionID, + Yolo: cfg.Yolo, }) if err != nil { return nil, fmt.Errorf("remote spawn: %w", err) } - keeper = toolsrv.NewProcessKeeper(cfg.Ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) { - return toolsrv.SpawnRemote(ctx, toolsrv.RemoteConfig{ + keeper = toolclient.NewProcessKeeper(cfg.Ctx, proc, func(ctx context.Context) (*toolclient.Process, error) { + return toolclient.SpawnRemote(ctx, toolclient.RemoteConfig{ SSHTarget: cfg.RemoteTarget, CWD: cfg.CWD, + SessionID: cfg.SessionID, + Yolo: cfg.Yolo, }) }) newToolServer = func() *toolsrv.Conn { @@ -106,16 +112,19 @@ func SetupToolServer(cfg ToolServerConfig) (*SessionInfra, error) { platform = proc.Info.Platform isGitRepo = proc.Info.IsGitRepo } else { - var dialOpts []toolsrv.Option + var dialOpts []toolclient.Option if cfg.Yolo { - dialOpts = append(dialOpts, toolsrv.WithYolo()) + dialOpts = append(dialOpts, toolclient.WithYolo()) } - proc, err = toolsrv.Spawn(cfg.Ctx, cfg.CWD, dialOpts...) + if cfg.SessionID != "" { + dialOpts = append(dialOpts, toolclient.WithSessionID(cfg.SessionID)) + } + proc, err = toolclient.Spawn(cfg.Ctx, cfg.CWD, dialOpts...) if err != nil { return nil, fmt.Errorf("local spawn: %w", err) } - keeper = toolsrv.NewProcessKeeper(cfg.Ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) { - return toolsrv.Spawn(ctx, cfg.CWD, dialOpts...) + keeper = toolclient.NewProcessKeeper(cfg.Ctx, proc, func(ctx context.Context) (*toolclient.Process, error) { + return toolclient.Spawn(ctx, cfg.CWD, dialOpts...) }) newToolServer = func() *toolsrv.Conn { conn, err := keeper.Dial() diff --git a/cmd/olliesrv/internal/toolclient/spawn.go b/cmd/olliesrv/internal/toolclient/spawn.go new file mode 100644 index 0000000..a7754ce --- /dev/null +++ b/cmd/olliesrv/internal/toolclient/spawn.go @@ -0,0 +1,403 @@ +// spawn.go - Process spawning and lifecycle management for toolsrv. +package toolclient + +import ( + "bufio" + "context" + "crypto/rand" + "encoding/hex" + "fmt" + "io" + "net" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "syscall" + "time" + + "ollie/paths" + toolsrvclient "ollie/toolsrv" +) + +// Process represents a running toolsrv process. +type Process struct { + Cmd *exec.Cmd + SocketPath string + Secret string + Info ProcessInfo + cancel context.CancelFunc +} + +// ProcessInfo contains metadata about the toolsrv process. +type ProcessInfo struct { + Platform string + IsGitRepo bool +} + +// Option configures spawning behavior. +type Option func(*spawnConfig) + +type spawnConfig struct { + yolo bool + sessionID string +} + +// WithYolo disables sandbox enforcement. +func WithYolo() Option { + return func(c *spawnConfig) { c.yolo = true } +} + +// WithSessionID sets the session ID for tool registry. +func WithSessionID(id string) Option { + return func(c *spawnConfig) { c.sessionID = id } +} + +// Spawn starts a local toolsrv process. +func Spawn(ctx context.Context, cwd string, opts ...Option) (*Process, error) { + cfg := &spawnConfig{} + for _, opt := range opts { + opt(cfg) + } + + // Generate socket path and secret + socketPath := filepath.Join(os.TempDir(), fmt.Sprintf("toolsrv-%d-%s.sock", os.Getpid(), randomHex(8))) + secret := randomHex(32) + + // Find toolsrv binary + toolsrvPath, err := findToolsrv() + if err != nil { + return nil, fmt.Errorf("find toolsrv: %w", err) + } + + // Build command + args := []string{"serve", "--cwd", cwd, "--listen", socketPath} + if cfg.yolo { + args = append(args, "--yolo") + } + if cfg.sessionID != "" { + args = append(args, "--session-id", cfg.sessionID) + } + + procCtx, cancel := context.WithCancel(ctx) + cmd := exec.CommandContext(procCtx, toolsrvPath, args...) + // Note: secret is NOT passed via env var. It's established on first + // connection via 9P Tauth (socket permissions + SSH handle security). + cmd.Stderr = os.Stderr // Let server errors go to stderr for debugging + + if err := cmd.Start(); err != nil { + cancel() + return nil, fmt.Errorf("start toolsrv: %w", err) + } + + // Wait for socket to be ready + if err := waitForSocket(socketPath, 5*time.Second); err != nil { + cancel() + cmd.Process.Kill() + return nil, fmt.Errorf("wait for socket: %w", err) + } + + return &Process{ + Cmd: cmd, + SocketPath: socketPath, + Secret: secret, + Info: ProcessInfo{ + Platform: "linux", + IsGitRepo: false, // Caller can override + }, + cancel: cancel, + }, nil +} + +// Kill terminates the process. +func (p *Process) Kill() error { + if p.cancel != nil { + p.cancel() + } + if p.Cmd != nil && p.Cmd.Process != nil { + return p.Cmd.Process.Kill() + } + return nil +} + +// Close terminates the process (alias for Kill). +func (p *Process) Close() error { + return p.Kill() +} + +// Wait waits for the process to exit. +func (p *Process) Wait() error { + if p.Cmd != nil { + return p.Cmd.Wait() + } + return nil +} + +// RemoteConfig configures remote toolsrv spawning. +type RemoteConfig struct { + SSHTarget string // user@host or host + CWD string + SessionID string // session ID for tool registry + Yolo bool // skip sandbox enforcement +} + +// SpawnRemote starts a remote toolsrv process via SSH. +// Deploys the local ~/.config/ollie directory (including tools/*.meta) to the +// remote host, then starts toolsrv with SSH Unix socket forwarding. +func SpawnRemote(ctx context.Context, cfg RemoteConfig) (*Process, error) { + // Verify local config dir exists + cfgDir := paths.CfgDir() + if _, err := os.Stat(cfgDir); err != nil { + return nil, fmt.Errorf("local config dir not found at %s", cfgDir) + } + + // Generate socket paths + sockDir := filepath.Join(paths.RuntimeDir(), "ollie") + os.MkdirAll(sockDir, 0700) + localSock := filepath.Join(sockDir, fmt.Sprintf("remote-%d-%s.sock", os.Getpid(), randomHex(8))) + remoteSock := fmt.Sprintf("/tmp/ollie-toolsrv-%d-%s.sock", os.Getpid(), randomHex(8)) + + // Build remote bootstrap script + // Extracts tarball to ~/.cache/ollie, sets XDG_CONFIG_HOME so tools/*.meta are found + var args []string + args = append(args, "serve", "--cwd", shellEscape(cfg.CWD), "--listen", shellEscape(remoteSock)) + if cfg.SessionID != "" { + args = append(args, "--session-id", shellEscape(cfg.SessionID)) + } + if cfg.Yolo { + args = append(args, "--yolo") + } + + bootstrap := fmt.Sprintf(`#!/bin/sh +set -e +CACHE_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/ollie" +mkdir -p "$CACHE_DIR" +tar xzf - -C "$CACHE_DIR" +export XDG_CONFIG_HOME="$(dirname "$CACHE_DIR")" +export PATH="$CACHE_DIR/bin:$PATH" +exec "$CACHE_DIR/bin/toolsrv" %s +`, strings.Join(args, " ")) + + // Build SSH command with socket forwarding + sshArgs := []string{ + "-o", "BatchMode=yes", + "-o", "StrictHostKeyChecking=accept-new", + "-o", "StreamLocalBindUnlink=yes", + "-L", localSock + ":" + remoteSock, + } + sshTarget := cfg.SSHTarget + if host, port, ok := strings.Cut(sshTarget, ":"); ok && port != "" { + sshTarget = host + sshArgs = append(sshArgs, "-p", port) + } + sshArgs = append(sshArgs, sshTarget, "bash -s") + + procCtx, cancel := context.WithCancel(ctx) + cmd := exec.CommandContext(procCtx, "ssh", sshArgs...) + cmd.SysProcAttr = &syscall.SysProcAttr{ + Setpgid: true, + Pdeathsig: syscall.SIGTERM, + } + + stdin, err := cmd.StdinPipe() + if err != nil { + cancel() + return nil, fmt.Errorf("ssh stdin pipe: %w", err) + } + stderrPipe, err := cmd.StderrPipe() + if err != nil { + cancel() + return nil, fmt.Errorf("ssh stderr pipe: %w", err) + } + + if err := cmd.Start(); err != nil { + cancel() + return nil, fmt.Errorf("ssh start: %w", err) + } + + // Send tarball of config dir over stdin + if err := writeTarball(stdin, cfgDir); err != nil { + cancel() + cmd.Process.Kill() + return nil, fmt.Errorf("send tarball: %w", err) + } + // Send bootstrap script + if _, err := stdin.Write([]byte(bootstrap)); err != nil { + cancel() + cmd.Process.Kill() + return nil, fmt.Errorf("send bootstrap: %w", err) + } + stdin.Close() + + // Wait for toolsrv to signal ready + readyCh := make(chan struct{}) + go func() { + scanner := bufio.NewScanner(stderrPipe) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + // toolsrv prints "toolsrv: listening on " + if strings.Contains(line, "listening on") { + close(readyCh) + return + } + } + }() + + select { + case <-readyCh: + case <-time.After(30 * time.Second): + cancel() + cmd.Process.Kill() + return nil, fmt.Errorf("remote toolsrv startup timeout") + } + + // Wait for local socket forwarding to be ready + if err := waitForSocket(localSock, 5*time.Second); err != nil { + cancel() + cmd.Process.Kill() + return nil, fmt.Errorf("wait for forwarded socket: %w", err) + } + + return &Process{ + Cmd: cmd, + SocketPath: localSock, + Secret: "", // Will be established on first Tauth + Info: ProcessInfo{ + Platform: "linux", // TODO: detect from remote + IsGitRepo: false, + }, + cancel: cancel, + }, nil +} + +// writeTarball creates a gzipped tarball of the directory and writes it to w. +func writeTarball(w io.Writer, dir string) error { + cmd := exec.Command("tar", "czf", "-", "-C", dir, ".") + cmd.Stdout = w + cmd.Stderr = os.Stderr + return cmd.Run() +} + +// shellEscape escapes a string for safe use in a shell command. +func shellEscape(s string) string { + if s == "" { + return "''" + } + if !strings.ContainsAny(s, " \t\n\r'\"\\$`!#&|;(){}") { + return s + } + return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'" +} + +// ProcessKeeper manages process lifecycle with automatic respawning. +type ProcessKeeper struct { + ctx context.Context + proc *Process + respawn func(context.Context) (*Process, error) + mu sync.Mutex +} + +// NewProcessKeeper creates a new process keeper. +func NewProcessKeeper(ctx context.Context, proc *Process, respawn func(context.Context) (*Process, error)) *ProcessKeeper { + return &ProcessKeeper{ + ctx: ctx, + proc: proc, + respawn: respawn, + } +} + +// SetContext updates the context used for respawning. +func (pk *ProcessKeeper) SetContext(ctx context.Context) { + pk.mu.Lock() + defer pk.mu.Unlock() + pk.ctx = ctx +} + +// Dial connects to the managed process, respawning if necessary. +func (pk *ProcessKeeper) Dial() (*toolsrvclient.Conn, error) { + pk.mu.Lock() + defer pk.mu.Unlock() + + // Try to connect to existing process + if pk.proc != nil { + conn, err := toolsrvclient.Dial(pk.proc.SocketPath, pk.proc.Secret) + if err == nil { + // Save the secret if this was first connection + if pk.proc.Secret == "" { + pk.proc.Secret = conn.Secret() + } + return conn, nil + } + // Process may have died, try to respawn + } + + if pk.respawn != nil { + proc, err := pk.respawn(pk.ctx) + if err != nil { + return nil, fmt.Errorf("respawn: %w", err) + } + pk.proc = proc + conn, err := toolsrvclient.Dial(proc.SocketPath, proc.Secret) + if err != nil { + return nil, err + } + return conn, nil + } + + return nil, fmt.Errorf("no process available") +} + +// Close terminates the managed process. +func (pk *ProcessKeeper) Close() error { + pk.mu.Lock() + defer pk.mu.Unlock() + if pk.proc != nil { + return pk.proc.Kill() + } + return nil +} + +// --- Helpers --- + +func randomHex(n int) string { + b := make([]byte, n) + rand.Read(b) + return hex.EncodeToString(b) +} + +func findToolsrv() (string, error) { + // Check if toolsrv is in PATH + if path, err := exec.LookPath("toolsrv"); err == nil { + return path, nil + } + + // Check common locations + home, _ := os.UserHomeDir() + candidates := []string{ + filepath.Join(home, ".config", "ollie", "bin", "toolsrv"), + filepath.Join(home, "go", "bin", "toolsrv"), + "/usr/local/bin/toolsrv", + } + + for _, p := range candidates { + if _, err := os.Stat(p); err == nil { + return p, nil + } + } + + return "", fmt.Errorf("toolsrv binary not found in PATH or common locations") +} + +func waitForSocket(path string, timeout time.Duration) error { + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + conn, err := net.DialTimeout("unix", path, 100*time.Millisecond) + if err == nil { + conn.Close() + return nil + } + time.Sleep(50 * time.Millisecond) + } + return fmt.Errorf("timeout waiting for socket %s", path) +} diff --git a/cmd/olliesrv/main.go b/cmd/olliesrv/main.go index 2302ad8..e03373b 100644 --- a/cmd/olliesrv/main.go +++ b/cmd/olliesrv/main.go @@ -14,13 +14,13 @@ import ( "syscall" "9fans.net/go/plan9/client" - "ollie/agent" - "ollie/bypass" + "ollie/cmd/olliesrv/internal/agent" + "ollie/cmd/olliesrv/internal/bypass" "ollie/env" - "ollie/fs" + "ollie/cmd/olliesrv/internal/fs" olog "ollie/log" "ollie/paths" - "ollie/session" + "ollie/cmd/olliesrv/internal/session" "github.com/godbus/dbus/v5" ) diff --git a/cmd/olliesrv/server.go b/cmd/olliesrv/server.go index 32c9291..8ec6ed1 100644 --- a/cmd/olliesrv/server.go +++ b/cmd/olliesrv/server.go @@ -14,9 +14,9 @@ import ( "sync" "time" - "ollie/fs" + "ollie/cmd/olliesrv/internal/fs" olog "ollie/log" - "ollie/session" + "ollie/cmd/olliesrv/internal/session" "9fans.net/go/plan9" ) diff --git a/cmd/ollie-remote/.gitignore b/cmd/toolsrv/.gitignore similarity index 100% rename from cmd/ollie-remote/.gitignore rename to cmd/toolsrv/.gitignore diff --git a/cmd/toolsrv/auth9p.go b/cmd/toolsrv/auth9p.go new file mode 100644 index 0000000..971a793 --- /dev/null +++ b/cmd/toolsrv/auth9p.go @@ -0,0 +1,17 @@ +// auth9p.go - Authentication helpers for 9P toolsrv +package main + +import ( + "crypto/rand" + "encoding/hex" +) + +// GenerateSecret generates a random shared secret for toolsrv auth. +// The secret is transmitted over the socket via 9P Tauth. +func GenerateSecret() (string, error) { + secret := make([]byte, 32) + if _, err := rand.Read(secret); err != nil { + return "", err + } + return hex.EncodeToString(secret), nil +} diff --git a/cmd/toolsrv/exec9p.go b/cmd/toolsrv/exec9p.go new file mode 100644 index 0000000..6e3f2d8 --- /dev/null +++ b/cmd/toolsrv/exec9p.go @@ -0,0 +1,380 @@ +// exec9p.go - Decoupled tool execution for 9P toolsrv +package main + +import ( + "bytes" + "context" + "encoding/binary" + "encoding/json" + "fmt" + "io" + "net" + "os" + "os/exec" + "path/filepath" + "strings" + "syscall" + "time" + + "ollie/cmd/toolsrv/internal/sandbox" + "ollie/paths" + "ollie/toolsrv" +) + +// 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 +} + +// StreamFunc returns the streaming output function from context, if any. +func StreamFunc(ctx context.Context) func(string) { + if fn, ok := ctx.Value(streamKey{}).(func(string)); ok { + return fn + } + return nil +} + +type streamKey struct{} + +// WithStreamFunc attaches a streaming output function to context. +func WithStreamFunc(ctx context.Context, fn func(string)) context.Context { + return context.WithValue(ctx, streamKey{}, fn) +} + +// 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 toolsrv.ToolInfo, args json.RawMessage, cfg ExecConfig) (json.RawMessage, error) { + // Resolve script path + toolPath, err := toolsrv.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 := toolsrv.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(toolsrv.ToolResult{ + IsError: true, + Content: []toolsrv.ToolResultContent{{Type: "text", Text: result + ": " + err.Error()}}, + }) + } + return json.Marshal(toolsrv.ToolResult{ + Content: []toolsrv.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 + } + } +} + +// plan9Namespace computes the Plan 9 namespace directory. +func plan9Namespace(env map[string]string) string { + if ns := env["NAMESPACE"]; ns != "" { + return ns + } + disp := env["DISPLAY"] + if disp == "" { + return "/tmp/ns." + env["USER"] + "." + "unix" + } + // Strip leading "localhost" if present + disp = strings.TrimPrefix(disp, "localhost") + disp = strings.ReplaceAll(disp, "/", "_") + return "/tmp/ns." + env["USER"] + "." + disp +} + +// limitedWriter wraps a writer with a size limit and optional streaming. +type limitedWriter struct { + w io.Writer + limit int + written int + truncated bool + stream func(string) +} + +func (lw *limitedWriter) Write(p []byte) (int, error) { + if lw.truncated { + return len(p), nil // discard but report success + } + remaining := lw.limit - lw.written + if remaining <= 0 { + lw.truncated = true + return len(p), nil + } + toWrite := p + if len(p) > remaining { + toWrite = p[:remaining] + lw.truncated = true + } + n, err := lw.w.Write(toWrite) + lw.written += n + if lw.stream != nil && n > 0 { + lw.stream(string(toWrite[:n])) + } + return len(p), err +} diff --git a/cmd/toolsrv/fs9p.go b/cmd/toolsrv/fs9p.go new file mode 100644 index 0000000..353655f --- /dev/null +++ b/cmd/toolsrv/fs9p.go @@ -0,0 +1,596 @@ +// Package toolsrv provides a 9P filesystem interface to the tool execution server. +// +// Namespace: +// +// /ctl write: load , unload +// /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// dir: process directory +// /proc//out read: output stream +// /proc//wait read: blocks until exit, returns exit code +// /proc//stat read: status (running, exited N, runtime) +// /proc//ctl write: signal , 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//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) +} diff --git a/cmd/toolsrv/fs9p_test.go b/cmd/toolsrv/fs9p_test.go new file mode 100644 index 0000000..04974cb --- /dev/null +++ b/cmd/toolsrv/fs9p_test.go @@ -0,0 +1,216 @@ +package main + +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") + } +} diff --git a/cmd/toolsrv/integration_test.go b/cmd/toolsrv/integration_test.go new file mode 100644 index 0000000..bd0ce4d --- /dev/null +++ b/cmd/toolsrv/integration_test.go @@ -0,0 +1,469 @@ +// integration_test.go - End-to-end tests for 9P toolsrv with Tauth. +// +// These tests build and spawn a real toolsrv process, then connect via 9P +// to verify the authentication flow works correctly. +package main + +import ( + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" + + "ollie/toolsrv" +) + +var testToolsrvBinary string + +// TestMain builds the toolsrv binary once for all integration tests. +func TestMain(m *testing.M) { + // Build toolsrv binary to a temp location + tmpDir, err := os.MkdirTemp("", "toolsrv-test-*") + if err != nil { + fmt.Fprintf(os.Stderr, "failed to create temp dir: %v\n", err) + os.Exit(1) + } + defer os.RemoveAll(tmpDir) + + testToolsrvBinary = filepath.Join(tmpDir, "toolsrv") + + // Build the binary + cmd := exec.Command("go", "build", "-o", testToolsrvBinary, "ollie/cmd/toolsrv") + cmd.Stderr = os.Stderr + if err := cmd.Run(); err != nil { + fmt.Fprintf(os.Stderr, "failed to build toolsrv: %v\n", err) + os.Exit(1) + } + + os.Exit(m.Run()) +} + +// startTestServer starts a toolsrv process for testing. +// Returns socketPath and a cleanup function. +func startTestServer(t *testing.T) (socketPath string, cleanup func()) { + t.Helper() + + tmpDir, err := os.MkdirTemp("", "toolsrv-socket-*") + if err != nil { + t.Fatalf("failed to create temp dir: %v", err) + } + + socketPath = filepath.Join(tmpDir, "toolsrv.sock") + cwd, _ := os.Getwd() + + ctx, cancel := context.WithCancel(context.Background()) + cmd := exec.CommandContext(ctx, testToolsrvBinary, "serve", "--cwd", cwd, "--listen", socketPath) + cmd.Stderr = os.Stderr + // Set session ID so tool registry is configured + cmd.Env = append(os.Environ(), "OLLIE_SESSION_ID=test-session") + + if err := cmd.Start(); err != nil { + cancel() + os.RemoveAll(tmpDir) + t.Fatalf("failed to start toolsrv: %v", err) + } + + // Wait for socket to be ready + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if _, err := os.Stat(socketPath); err == nil { + break + } + time.Sleep(50 * time.Millisecond) + } + + if _, err := os.Stat(socketPath); err != nil { + cancel() + cmd.Wait() + os.RemoveAll(tmpDir) + t.Fatalf("toolsrv socket not ready: %v", err) + } + + return socketPath, func() { + cancel() + cmd.Wait() + os.RemoveAll(tmpDir) + } +} + +func TestIntegration_FirstConnectionSetsSecret(t *testing.T) { + socketPath, cleanup := startTestServer(t) + defer cleanup() + + // First connection with a secret should succeed + secret := "test-secret-first-connection" + conn, err := toolsrv.Dial(socketPath, secret) + if err != nil { + t.Fatalf("first Dial failed: %v", err) + } + defer conn.Close() + + // Should have received a token + if conn.Token() == "" { + t.Error("expected non-empty token") + } + + // Secret should be stored + if conn.Secret() != secret { + t.Errorf("Secret() = %q, want %q", conn.Secret(), secret) + } +} + +func TestIntegration_SameSecretReconnects(t *testing.T) { + socketPath, cleanup := startTestServer(t) + defer cleanup() + + secret := "test-secret-reconnect" + + // First connection establishes the secret + conn1, err := toolsrv.Dial(socketPath, secret) + if err != nil { + t.Fatalf("first Dial failed: %v", err) + } + token1 := conn1.Token() + conn1.Close() + + // Second connection with same secret should succeed + conn2, err := toolsrv.Dial(socketPath, secret) + if err != nil { + t.Fatalf("second Dial failed: %v", err) + } + defer conn2.Close() + + // Should get the same token + if conn2.Token() != token1 { + t.Errorf("token changed: %q -> %q", token1, conn2.Token()) + } +} + +func TestIntegration_WrongSecretFails(t *testing.T) { + socketPath, cleanup := startTestServer(t) + defer cleanup() + + // First connection establishes the secret + conn1, err := toolsrv.Dial(socketPath, "correct-secret") + if err != nil { + t.Fatalf("first Dial failed: %v", err) + } + conn1.Close() + + // Second connection with wrong secret should fail + conn2, err := toolsrv.Dial(socketPath, "wrong-secret") + if err == nil { + conn2.Close() + t.Fatal("Dial with wrong secret should fail") + } +} + +func TestIntegration_EmptySecretGeneratesOne(t *testing.T) { + socketPath, cleanup := startTestServer(t) + defer cleanup() + + // Empty secret should generate a random one + conn, err := toolsrv.Dial(socketPath, "") + if err != nil { + t.Fatalf("Dial with empty secret failed: %v", err) + } + defer conn.Close() + + // Should have generated and stored a secret + if conn.Secret() == "" { + t.Error("expected non-empty generated secret") + } + + // Token should be set + if conn.Token() == "" { + t.Error("expected non-empty token") + } +} + +func TestIntegration_BasicOperations(t *testing.T) { + socketPath, cleanup := startTestServer(t) + defer cleanup() + + conn, err := toolsrv.Dial(socketPath, "test-secret-ops") + if err != nil { + t.Fatalf("Dial failed: %v", err) + } + defer conn.Close() + + // Test Ping + if err := conn.Ping(); err != nil { + t.Errorf("Ping failed: %v", err) + } + + // Test HostInfo + info, err := conn.HostInfo() + if err != nil { + t.Errorf("HostInfo failed: %v", err) + } + if info.Platform == "" { + t.Error("HostInfo.Platform is empty") + } + + // Test ListTools (should be empty initially, but shouldn't error) + tools, err := conn.ListTools() + if err != nil { + t.Errorf("ListTools failed: %v", err) + } + t.Logf("ListTools returned %d tools", len(tools)) +} + +func TestIntegration_ToolExecution(t *testing.T) { + // Skip if tools directory doesn't exist (CI environment) + if _, err := os.ReadDir(toolsrv.ToolsPath()); err != nil { + t.Skipf("skipping: tools directory not available: %v", err) + } + + socketPath, cleanup := startTestServer(t) + defer cleanup() + + conn, err := toolsrv.Dial(socketPath, "test-secret-exec") + if err != nil { + t.Fatalf("Dial failed: %v", err) + } + defer conn.Close() + + // Load the shell tool + if err := conn.LoadTool("shell"); err != nil { + t.Skipf("skipping: LoadTool failed (shell tool not available): %v", err) + } + + // Verify it's loaded + tools, err := conn.ListTools() + if err != nil { + t.Fatalf("ListTools failed: %v", err) + } + found := false + for _, tool := range tools { + if tool.Name == "shell" { + found = true + break + } + } + if !found { + t.Fatal("shell tool not found after loading") + } + + // Execute a tool - this mimics what agent/runtime.go does + ctx := context.Background() + args := json.RawMessage(`{"cmd": "echo hello from test"}`) + result, err := conn.CallTool(ctx, "shell", args) + if err != nil { + t.Fatalf("CallTool failed: %v", err) + } + + // Parse the result the same way agent/runtime.go does + var toolResult toolsrv.ToolResult + if err := json.Unmarshal(result, &toolResult); err != nil { + t.Fatalf("failed to parse toolsrv.ToolResult: %v (raw: %s)", err, string(result)) + } + + // Verify structure + if len(toolResult.Content) == 0 { + t.Fatal("toolsrv.ToolResult.Content is empty") + } + if toolResult.Content[0].Type != "text" { + t.Errorf("expected type 'text', got %q", toolResult.Content[0].Type) + } + if !strings.Contains(toolResult.Content[0].Text, "hello from test") { + t.Errorf("output doesn't contain expected text: %q", toolResult.Content[0].Text) + } + if toolResult.IsError { + t.Error("toolsrv.ToolResult.IsError should be false") + } + + t.Logf("Tool result: %+v", toolResult) +} + +func TestIntegration_MultilineContent(t *testing.T) { + // Skip if tools directory doesn't exist (CI environment) + if _, err := os.ReadDir(toolsrv.ToolsPath()); err != nil { + t.Skipf("skipping: tools directory not available: %v", err) + } + + socketPath, cleanup := startTestServer(t) + defer cleanup() + + conn, err := toolsrv.Dial(socketPath, "test-secret-multiline") + if err != nil { + t.Fatalf("Dial failed: %v", err) + } + defer conn.Close() + + // Load the file_write tool + if err := conn.LoadTool("file_write"); err != nil { + t.Skipf("skipping: LoadTool failed (file_write tool not available): %v", err) + } + + // Create a temp file path + tmpFile := filepath.Join(t.TempDir(), "multiline_test.txt") + + // Test with multiline content including special characters + multilineContent := "line one\nline two\nline three with \"quotes\"\nline four with 'single quotes'\nline five with \\backslash\n" + + ctx := context.Background() + args := json.RawMessage(fmt.Sprintf(`{"path": %q, "content": %q}`, tmpFile, multilineContent)) + result, err := conn.CallTool(ctx, "file_write", args) + if err != nil { + t.Fatalf("CallTool failed: %v", err) + } + + var toolResult toolsrv.ToolResult + if err := json.Unmarshal(result, &toolResult); err != nil { + t.Fatalf("failed to parse toolsrv.ToolResult: %v (raw: %s)", err, string(result)) + } + + if toolResult.IsError { + t.Fatalf("file_write returned error: %s", toolResult.Content[0].Text) + } + + // Read the file back and verify + written, err := os.ReadFile(tmpFile) + if err != nil { + t.Fatalf("failed to read written file: %v", err) + } + + if string(written) != multilineContent { + t.Errorf("content mismatch:\nexpected: %q\ngot: %q", multilineContent, string(written)) + } +} + +func TestIntegration_ToolExecutionError(t *testing.T) { + // Skip if tools directory doesn't exist (CI environment) + if _, err := os.ReadDir(toolsrv.ToolsPath()); err != nil { + t.Skipf("skipping: tools directory not available: %v", err) + } + + socketPath, cleanup := startTestServer(t) + defer cleanup() + + conn, err := toolsrv.Dial(socketPath, "test-secret-exec-err") + if err != nil { + t.Fatalf("Dial failed: %v", err) + } + defer conn.Close() + + // Load the shell tool + if err := conn.LoadTool("shell"); err != nil { + t.Skipf("skipping: LoadTool failed (shell tool not available): %v", err) + } + + // Execute a command that fails + ctx := context.Background() + args := json.RawMessage(`{"cmd": "exit 42"}`) + result, err := conn.CallTool(ctx, "shell", args) + if err != nil { + t.Fatalf("CallTool failed: %v", err) + } + + // Parse the result + var toolResult toolsrv.ToolResult + if err := json.Unmarshal(result, &toolResult); err != nil { + t.Fatalf("failed to parse toolsrv.ToolResult: %v (raw: %s)", err, string(result)) + } + + // Should have exit code in output + if len(toolResult.Content) == 0 { + t.Fatal("toolsrv.ToolResult.Content is empty") + } + if !strings.Contains(toolResult.Content[0].Text, "exit: 42") { + t.Errorf("output doesn't contain exit code: %q", toolResult.Content[0].Text) + } + + t.Logf("Error tool result: %+v", toolResult) +} + +func TestIntegration_ConcurrentConnections(t *testing.T) { + socketPath, cleanup := startTestServer(t) + defer cleanup() + + secret := "test-secret-concurrent" + + // Establish the secret + conn0, err := toolsrv.Dial(socketPath, secret) + if err != nil { + t.Fatalf("initial Dial failed: %v", err) + } + expectedToken := conn0.Token() + conn0.Close() + + // Open multiple concurrent connections + const numConns = 5 + conns := make([]*toolsrv.Conn, numConns) + errors := make([]error, numConns) + + for i := 0; i < numConns; i++ { + conns[i], errors[i] = toolsrv.Dial(socketPath, secret) + } + + // Check all succeeded with same token + for i := 0; i < numConns; i++ { + if errors[i] != nil { + t.Errorf("connection %d failed: %v", i, errors[i]) + continue + } + if conns[i].Token() != expectedToken { + t.Errorf("connection %d has different token", i) + } + conns[i].Close() + } +} + +func TestIntegration_ProcessKeeperRespawn(t *testing.T) { + // This test verifies that when a toolsrv process dies, + // the ProcessKeeper respawns it with a new secret. + + cwd, _ := os.Getwd() + + // We need to use the real Spawn which requires the binary in PATH + // For now, manually simulate what ProcessKeeper does + + // Start first server + socketPath1, cleanup1 := startTestServer(t) + + secret1 := "secret-for-first-server" + conn1, err := toolsrv.Dial(socketPath1, secret1) + if err != nil { + cleanup1() + t.Fatalf("first Dial failed: %v", err) + } + token1 := conn1.Token() + conn1.Close() + + // Kill the server + cleanup1() + + // Start second server (simulates respawn) + socketPath2, cleanup2 := startTestServer(t) + defer cleanup2() + + // New server should accept a new secret (any secret, since it's fresh) + secret2 := "secret-for-second-server" + conn2, err := toolsrv.Dial(socketPath2, secret2) + if err != nil { + t.Fatalf("second Dial failed: %v", err) + } + token2 := conn2.Token() + conn2.Close() + + // Tokens should be different (different servers) + if token1 == token2 { + t.Error("tokens should be different after respawn") + } + + // Old secret should NOT work on new server + conn3, err := toolsrv.Dial(socketPath2, secret1) + if err == nil { + conn3.Close() + t.Error("old secret should not work on new server") + } + + t.Logf("Successfully verified: old_token=%s new_token=%s", token1[:8], token2[:8]) + _ = cwd // silence unused warning +} diff --git a/sandbox/config.go b/cmd/toolsrv/internal/sandbox/config.go similarity index 100% rename from sandbox/config.go rename to cmd/toolsrv/internal/sandbox/config.go diff --git a/sandbox/default.yaml b/cmd/toolsrv/internal/sandbox/default.yaml similarity index 100% rename from sandbox/default.yaml rename to cmd/toolsrv/internal/sandbox/default.yaml diff --git a/sandbox/detect.go b/cmd/toolsrv/internal/sandbox/detect.go similarity index 100% rename from sandbox/detect.go rename to cmd/toolsrv/internal/sandbox/detect.go diff --git a/sandbox/restricted.yaml b/cmd/toolsrv/internal/sandbox/restricted.yaml similarity index 100% rename from sandbox/restricted.yaml rename to cmd/toolsrv/internal/sandbox/restricted.yaml diff --git a/sandbox/sandbox_test.go b/cmd/toolsrv/internal/sandbox/sandbox_test.go similarity index 100% rename from sandbox/sandbox_test.go rename to cmd/toolsrv/internal/sandbox/sandbox_test.go diff --git a/sandbox/wrapper.go b/cmd/toolsrv/internal/sandbox/wrapper.go similarity index 100% rename from sandbox/wrapper.go rename to cmd/toolsrv/internal/sandbox/wrapper.go diff --git a/cmd/toolsrv/main.go b/cmd/toolsrv/main.go new file mode 100644 index 0000000..6feaaa7 --- /dev/null +++ b/cmd/toolsrv/main.go @@ -0,0 +1,121 @@ +// toolsrv — 9P tool execution server. +// +// Executes tool calls in a sandboxed environment. Used both locally +// (spawned by olliesrv) and remotely (deployed via SSH bootstrap). +// +// Usage: +// +// toolsrv serve --cwd /path/to/project --listen /path/to/sock [--yolo] +// +// Listens on a Unix socket and serves a 9P filesystem for tool execution. +// Authentication happens via Tauth: first client to connect sets the secret, +// subsequent clients must provide the same secret. +package main + +import ( + "context" + "flag" + "fmt" + "net" + "os" + "os/signal" + "path/filepath" + "syscall" + + "ollie/env" + "ollie/fsedsl" +) + +var ( + cwd = flag.String("cwd", ".", "working directory for execution") + listenPath = flag.String("listen", "", "Unix socket path to listen on (required)") + sessionID = flag.String("session-id", "", "session ID for tool registry") + yolo = flag.Bool("yolo", false, "skip sandbox enforcement") + noAuth = flag.Bool("no-auth", false, "disable Tauth authentication (for debugging)") +) + +func main() { + if len(os.Args) < 2 || os.Args[1] != "serve" { + fmt.Fprintln(os.Stderr, "usage: toolsrv serve --cwd --listen [--yolo]") + os.Exit(1) + } + flag.CommandLine.Parse(os.Args[2:]) + + if *listenPath == "" { + fmt.Fprintln(os.Stderr, "error: --listen is required") + os.Exit(1) + } + + env.EnsureDefaults() + + // Prepend our bin dir to PATH so landrun is found (deployed alongside us). + home, _ := os.UserHomeDir() + binDir := filepath.Join(home, ".config", "ollie", "bin") + os.Setenv("PATH", binDir+":"+os.Getenv("PATH")) + + // Create tool registry + toolReg, _ := NewRegistry() + sid := *sessionID + if sid == "" { + sid = os.Getenv("OLLIE_SESSION_ID") // fallback to env var + } + + // Create 9P server (secret is established via Tauth, not env var) + srv := NewServer9P() + srv.SetYolo(*yolo) + if toolReg != nil && sid != "" { + srv.SetRegistry(toolReg, sid) + } + + // Build the filesystem tree + ctx := ToolsrvCtx{Server: srv} + tree := fsedsl.BuildTree(ToolsrvSpec(), ctx) + if tree == nil { + fmt.Fprintln(os.Stderr, "error: failed to build filesystem tree") + os.Exit(1) + } + + // Set up signal handling + runCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) + go func() { + <-sigCh + cancel() + }() + + // Remove stale socket if it exists + os.Remove(*listenPath) + + // Listen on Unix socket + ln, err := net.Listen("unix", *listenPath) + if err != nil { + fmt.Fprintf(os.Stderr, "error: listen: %v\n", err) + os.Exit(1) + } + defer ln.Close() + defer os.Remove(*listenPath) + + fmt.Fprintf(os.Stderr, "toolsrv: listening on %s\n", *listenPath) + + // Accept connections + go func() { + <-runCtx.Done() + ln.Close() + }() + + for { + conn, err := ln.Accept() + if err != nil { + select { + case <-runCtx.Done(): + return + default: + fmt.Fprintf(os.Stderr, "accept error: %v\n", err) + continue + } + } + go serve9P(runCtx, conn, tree, srv, *noAuth) + } +} diff --git a/cmd/toolsrv/registry.go b/cmd/toolsrv/registry.go new file mode 100644 index 0000000..4538fdb --- /dev/null +++ b/cmd/toolsrv/registry.go @@ -0,0 +1,114 @@ +// registry.go - Session-scoped tool registry for toolsrv. +package main + +import ( + "fmt" + "sort" + "sync" + + "ollie/toolsrv" +) + +// Registry tracks loaded tools per session. +type Registry struct { + mu sync.RWMutex + sessions map[string]map[string]toolsrv.ToolInfo + revisions map[string]uint64 +} + +// NewRegistry creates a new tool registry. +func NewRegistry() (*Registry, error) { + r := &Registry{ + sessions: make(map[string]map[string]toolsrv.ToolInfo), + revisions: make(map[string]uint64), + } + return r, nil +} + +// discover scans the tools directory and returns all valid tool infos. +func (r *Registry) discover() map[string]toolsrv.ToolInfo { + infos := toolsrv.DiscoverTools() + tools := make(map[string]toolsrv.ToolInfo, len(infos)) + for _, info := range infos { + tools[info.Name] = info + } + return tools +} + +// Load loads a tool into a session's registry. +func (r *Registry) Load(sessionID, name string) error { + tools := r.discover() + tool, exists := tools[name] + if !exists { + return fmt.Errorf("tool not found: %s", name) + } + + r.mu.Lock() + defer r.mu.Unlock() + + if r.sessions[sessionID] == nil { + r.sessions[sessionID] = make(map[string]toolsrv.ToolInfo) + } + + r.sessions[sessionID][name] = tool + r.revisions[sessionID]++ + return nil +} + +// Unload removes a tool from a session's registry. +func (r *Registry) Unload(sessionID, name string) error { + r.mu.Lock() + defer r.mu.Unlock() + + sessionTools, ok := r.sessions[sessionID] + if !ok { + return nil + } + + if _, exists := sessionTools[name]; !exists { + return nil + } + + delete(sessionTools, name) + r.revisions[sessionID]++ + return nil +} + +// Loaded returns all tools loaded for a session. +func (r *Registry) Loaded(sessionID string) []toolsrv.ToolInfo { + r.mu.RLock() + defer r.mu.RUnlock() + + sessionTools, ok := r.sessions[sessionID] + if !ok { + return nil + } + + var loaded []toolsrv.ToolInfo + for _, info := range sessionTools { + loaded = append(loaded, info) + } + sort.Slice(loaded, func(i, j int) bool { return loaded[i].Name < loaded[j].Name }) + return loaded +} + +// Lookup returns a tool by name from a session's registry. +func (r *Registry) Lookup(sessionID, name string) (toolsrv.ToolInfo, bool) { + r.mu.RLock() + defer r.mu.RUnlock() + + sessionTools, ok := r.sessions[sessionID] + if !ok { + return toolsrv.ToolInfo{}, false + } + + tool, exists := sessionTools[name] + return tool, exists +} + +// Revision returns the revision counter for a session's registry. +func (r *Registry) Revision(sessionID string) uint64 { + r.mu.RLock() + defer r.mu.RUnlock() + return r.revisions[sessionID] +} diff --git a/cmd/ollie-remote/sandbox-remote.yaml b/cmd/toolsrv/sandbox-remote.yaml similarity index 100% rename from cmd/ollie-remote/sandbox-remote.yaml rename to cmd/toolsrv/sandbox-remote.yaml diff --git a/cmd/toolsrv/server.go b/cmd/toolsrv/server.go new file mode 100644 index 0000000..8bb97ec --- /dev/null +++ b/cmd/toolsrv/server.go @@ -0,0 +1,517 @@ +// server.go - 9P filesystem server for toolsrv. +package main + +import ( + "context" + "fmt" + "io" + "net" + "os" + "sync" + + "ollie/fsedsl" + + "9fans.net/go/plan9" +) + +// fid represents an open file handle. +type fid struct { + path string + qid plan9.Qid + mode uint8 + opened bool + entry fsedsl.File + readBuf []byte // buffer for read data + readDone bool // true if file has been read + dirCache []byte +} + +// authFid represents an auth file handle for Tauth. +type authFid struct { + uname string + aname string + authenticated bool + token string +} + +// connState holds per-connection state. +type connState struct { + mu sync.RWMutex + fids map[uint32]*fid + authFids map[uint32]*authFid + ctx context.Context + uname string + srv *Server9P + noAuth bool // skip authentication (for debugging) +} + +// serve9P handles a single 9P connection using the fsedsl tree. +func serve9P(ctx context.Context, conn net.Conn, tree *fsedsl.Tree, srv *Server9P, noAuth bool) { + defer conn.Close() + + cs := &connState{ + fids: make(map[uint32]*fid), + authFids: make(map[uint32]*authFid), + ctx: ctx, + noAuth: noAuth, + srv: srv, + } + + for { + select { + case <-ctx.Done(): + return + default: + } + + fc, err := plan9.ReadFcall(conn) + if err != nil { + if err != io.EOF { + fmt.Fprintf(os.Stderr, "toolsrv: read error: %v\n", err) + } + return + } + + resp := handleFcall(cs, fc, tree) + if err := plan9.WriteFcall(conn, resp); err != nil { + fmt.Fprintf(os.Stderr, "toolsrv: write error: %v\n", err) + return + } + } +} + +func handleFcall(cs *connState, fc *plan9.Fcall, tree *fsedsl.Tree) *plan9.Fcall { + switch fc.Type { + case plan9.Tversion: + return handleVersion(fc) + case plan9.Tauth: + return handleAuth(cs, fc) + case plan9.Tattach: + return handleAttach(cs, fc, tree) + case plan9.Twalk: + return handleWalk(cs, fc, tree) + case plan9.Topen: + return handleOpen(cs, fc, tree) + case plan9.Tread: + return handleRead(cs, fc, tree) + case plan9.Twrite: + return handleWrite(cs, fc, tree) + case plan9.Tclunk: + return handleClunk(cs, fc) + case plan9.Tstat: + return handleStat(cs, fc, tree) + default: + return errFcall(fc, fmt.Sprintf("unsupported: %d", fc.Type)) + } +} + +func errFcall(fc *plan9.Fcall, msg string) *plan9.Fcall { + return &plan9.Fcall{Type: plan9.Rerror, Tag: fc.Tag, Ename: msg} +} + +func handleVersion(fc *plan9.Fcall) *plan9.Fcall { + msize := fc.Msize + if msize > 65536 { + msize = 65536 + } + return &plan9.Fcall{Type: plan9.Rversion, Tag: fc.Tag, Msize: msize, Version: "9P2000"} +} + +func handleAuth(cs *connState, fc *plan9.Fcall) *plan9.Fcall { + cs.mu.Lock() + cs.authFids[fc.Afid] = &authFid{ + uname: fc.Uname, + aname: fc.Aname, + } + cs.mu.Unlock() + + return &plan9.Fcall{ + Type: plan9.Rauth, + Tag: fc.Tag, + Aqid: plan9.Qid{Type: plan9.QTAUTH, Vers: 0, Path: uint64(fc.Afid)}, + } +} + +func handleAttach(cs *connState, fc *plan9.Fcall, tree *fsedsl.Tree) *plan9.Fcall { + // Skip auth check if noAuth is set (for debugging) + if cs.noAuth { + cs.mu.Lock() + cs.uname = fc.Uname + cs.fids[fc.Fid] = &fid{ + path: "/", + qid: plan9.Qid{Type: plan9.QTDIR, Path: 0}, + } + cs.mu.Unlock() + return &plan9.Fcall{Type: plan9.Rattach, Tag: fc.Tag, Qid: plan9.Qid{Type: plan9.QTDIR, Path: 0}} + } + + // Check auth - afid must exist and be authenticated + cs.mu.RLock() + af, hasAuth := cs.authFids[fc.Afid] + cs.mu.RUnlock() + + // NOFID means no auth required (we require auth) + if fc.Afid == plan9.NOFID { + return errFcall(fc, "authentication required") + } + + if !hasAuth { + return errFcall(fc, "invalid afid") + } + + if !af.authenticated { + return errFcall(fc, "not authenticated") + } + + cs.mu.Lock() + cs.uname = fc.Uname + cs.fids[fc.Fid] = &fid{ + path: "/", + qid: plan9.Qid{Type: plan9.QTDIR, Path: 0}, + } + cs.mu.Unlock() + + return &plan9.Fcall{Type: plan9.Rattach, Tag: fc.Tag, Qid: plan9.Qid{Type: plan9.QTDIR, Path: 0}} +} + +// handleAuthWrite handles writing the secret to an auth fid. +func handleAuthWrite(cs *connState, fc *plan9.Fcall, af *authFid) *plan9.Fcall { + clientSecret := string(fc.Data) + token, err := cs.srv.Authenticate(clientSecret) + if err != nil { + return errFcall(fc, err.Error()) + } + + af.authenticated = true + af.token = token + + return &plan9.Fcall{Type: plan9.Rwrite, Tag: fc.Tag, Count: uint32(len(fc.Data))} +} + +// handleAuthRead handles reading the token from an auth fid. +func handleAuthRead(cs *connState, fc *plan9.Fcall, af *authFid) *plan9.Fcall { + if !af.authenticated { + return errFcall(fc, "write secret first") + } + + data := []byte(af.token + "\n") + + offset := int(fc.Offset) + if offset >= len(data) { + return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: nil} + } + + end := offset + int(fc.Count) + if end > len(data) { + end = len(data) + } + + return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: data[offset:end]} +} + +func handleWalk(cs *connState, fc *plan9.Fcall, tree *fsedsl.Tree) *plan9.Fcall { + cs.mu.RLock() + f, ok := cs.fids[fc.Fid] + cs.mu.RUnlock() + if !ok { + return errFcall(fc, "unknown fid") + } + + newPath := f.path + var qids []plan9.Qid + + for _, name := range fc.Wname { + if name == ".." { + if newPath != "/" { + lastSlash := len(newPath) - 1 + for lastSlash > 0 && newPath[lastSlash] != '/' { + lastSlash-- + } + if lastSlash == 0 { + newPath = "/" + } else { + newPath = newPath[:lastSlash] + } + } + } else { + if newPath == "/" { + newPath = "/" + name + } else { + newPath = newPath + "/" + name + } + } + + // Check if path exists in tree + treePath := newPath + if treePath == "/" { + treePath = "" + } else { + treePath = treePath[1:] // strip leading / + } + + info, err := tree.Stat(treePath) + if err != nil { + if len(qids) == 0 { + return errFcall(fc, "not found") + } + break + } + + var qt uint8 = plan9.QTFILE + if info.IsDir() { + qt = plan9.QTDIR + } + qids = append(qids, plan9.Qid{Type: qt, Path: pathHash(newPath)}) + } + + if len(qids) == len(fc.Wname) && fc.Newfid != fc.Fid { + var lastQid plan9.Qid + if len(qids) > 0 { + lastQid = qids[len(qids)-1] + } else { + lastQid = f.qid + } + cs.mu.Lock() + cs.fids[fc.Newfid] = &fid{ + path: newPath, + qid: lastQid, + } + cs.mu.Unlock() + } else if len(qids) == len(fc.Wname) && len(qids) > 0 { + cs.mu.Lock() + f.path = newPath + f.qid = qids[len(qids)-1] + cs.mu.Unlock() + } + + return &plan9.Fcall{Type: plan9.Rwalk, Tag: fc.Tag, Wqid: qids} +} + +func handleOpen(cs *connState, fc *plan9.Fcall, tree *fsedsl.Tree) *plan9.Fcall { + cs.mu.Lock() + defer cs.mu.Unlock() + + f, ok := cs.fids[fc.Fid] + if !ok { + return errFcall(fc, "unknown fid") + } + + // Directories don't need an entry + if f.qid.Type&plan9.QTDIR != 0 { + f.opened = true + f.mode = fc.Mode + return &plan9.Fcall{Type: plan9.Ropen, Tag: fc.Tag, Qid: f.qid, Iounit: 8192} + } + + path := f.path + if path == "/" { + path = "" + } else { + path = path[1:] // strip leading / + } + + entry, err := tree.Open(path) + if err != nil { + return errFcall(fc, err.Error()) + } + + f.entry = entry + f.opened = true + f.mode = fc.Mode + f.readBuf = nil + f.readDone = false + + return &plan9.Fcall{Type: plan9.Ropen, Tag: fc.Tag, Qid: f.qid, Iounit: 8192} +} + +func handleRead(cs *connState, fc *plan9.Fcall, tree *fsedsl.Tree) *plan9.Fcall { + // Check if reading from auth fid + cs.mu.RLock() + af, isAuthFid := cs.authFids[fc.Fid] + f, isFid := cs.fids[fc.Fid] + cs.mu.RUnlock() + + if isAuthFid { + return handleAuthRead(cs, fc, af) + } + + if !isFid { + return errFcall(fc, "unknown fid") + } + if !f.opened { + return errFcall(fc, "not open") + } + + if f.qid.Type&plan9.QTDIR != 0 { + return handleDirRead(cs, f, fc, tree) + } + + if f.entry == nil { + return errFcall(fc, "no entry") + } + + // For rdwr files, check if we've already read + if f.entry.RequestRespMode() && f.readDone { + return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: nil} + } + + // Buffer read data if not already done + if f.readBuf == nil { + data, err := f.entry.Read() + if err != nil && err != io.EOF { + return errFcall(fc, err.Error()) + } + f.readBuf = data + } + + // Handle offset/count + offset := int(fc.Offset) + if offset >= len(f.readBuf) { + f.readDone = true + return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: nil} + } + + end := offset + int(fc.Count) + if end > len(f.readBuf) { + end = len(f.readBuf) + } + + return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: f.readBuf[offset:end]} +} + +func handleDirRead(cs *connState, f *fid, fc *plan9.Fcall, tree *fsedsl.Tree) *plan9.Fcall { + if fc.Offset == 0 || f.dirCache == nil { + path := f.path + if path == "/" { + path = "" + } else { + path = path[1:] + } + + entries, err := tree.Readdir(path) + if err != nil { + return errFcall(fc, err.Error()) + } + + var buf []byte + for _, e := range entries { + var qt uint8 = plan9.QTFILE + mode := plan9.Perm(0644) + if e.IsDir() { + qt = plan9.QTDIR + mode = plan9.Perm(plan9.DMDIR | 0755) + } + dir := plan9.Dir{ + Name: e.Name(), + Qid: plan9.Qid{Type: qt, Path: pathHash(f.path + "/" + e.Name())}, + Mode: mode, + Uid: "toolsrv", + Gid: "toolsrv", + } + data, _ := dir.Bytes() + buf = append(buf, data...) + } + f.dirCache = buf + } + + offset := int(fc.Offset) + if offset >= len(f.dirCache) { + return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: nil} + } + + end := offset + int(fc.Count) + if end > len(f.dirCache) { + end = len(f.dirCache) + } + + return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: f.dirCache[offset:end]} +} + +func handleWrite(cs *connState, fc *plan9.Fcall, tree *fsedsl.Tree) *plan9.Fcall { + // Check if writing to auth fid + cs.mu.Lock() + af, isAuthFid := cs.authFids[fc.Fid] + f, isFid := cs.fids[fc.Fid] + cs.mu.Unlock() + + if isAuthFid { + return handleAuthWrite(cs, fc, af) + } + + if !isFid { + return errFcall(fc, "unknown fid") + } + if !f.opened { + return errFcall(fc, "not open") + } + if f.entry == nil { + return errFcall(fc, "no entry") + } + + err := f.entry.Write(fc.Data) + if err != nil { + return errFcall(fc, err.Error()) + } + + return &plan9.Fcall{Type: plan9.Rwrite, Tag: fc.Tag, Count: uint32(len(fc.Data))} +} + +func handleClunk(cs *connState, fc *plan9.Fcall) *plan9.Fcall { + cs.mu.Lock() + delete(cs.fids, fc.Fid) + delete(cs.authFids, fc.Fid) + cs.mu.Unlock() + return &plan9.Fcall{Type: plan9.Rclunk, Tag: fc.Tag} +} + +func handleStat(cs *connState, fc *plan9.Fcall, tree *fsedsl.Tree) *plan9.Fcall { + cs.mu.RLock() + f, ok := cs.fids[fc.Fid] + cs.mu.RUnlock() + if !ok { + return errFcall(fc, "unknown fid") + } + + path := f.path + if path == "/" { + path = "" + } else { + path = path[1:] + } + + info, err := tree.Stat(path) + if err != nil { + return errFcall(fc, err.Error()) + } + + var qt uint8 = plan9.QTFILE + mode := plan9.Perm(0644) + if info.IsDir() { + qt = plan9.QTDIR + mode = plan9.Perm(plan9.DMDIR | 0755) + } + + name := info.Name() + if name == "" { + name = "/" + } + + dir := plan9.Dir{ + Name: name, + Qid: plan9.Qid{Type: qt, Path: pathHash(f.path)}, + Mode: mode, + Uid: "toolsrv", + Gid: "toolsrv", + } + data, _ := dir.Bytes() + + return &plan9.Fcall{Type: plan9.Rstat, Tag: fc.Tag, Stat: data} +} + +func pathHash(p string) uint64 { + var h uint64 = 5381 + for _, c := range p { + h = h*33 + uint64(c) + } + return h +} diff --git a/cmd/toolsrv/server9p.go b/cmd/toolsrv/server9p.go new file mode 100644 index 0000000..354addd --- /dev/null +++ b/cmd/toolsrv/server9p.go @@ -0,0 +1,101 @@ +// server9p.go - State management for 9P toolsrv +package main + +import ( + "crypto/rand" + "encoding/hex" + "fmt" + "runtime" + "sync" + + "ollie/fsedsl" +) + +// Server9P holds the state for the 9P tool server. +type Server9P struct { + mu sync.RWMutex + + secret string // set on first auth, verified on subsequent + token string // session token returned after auth + registry *Registry // tool registry + sessID string // session ID for registry scoping + yolo bool // skip sandbox + + // Process management + fs *FS9P +} + +// NewServer9P creates a new 9P tool server. +func NewServer9P() *Server9P { + return &Server9P{ + fs: NewFS9P(""), // cwd set per-agent + } +} + +// SetRegistry configures the tool registry and session ID. +func (s *Server9P) SetRegistry(r *Registry, sessID string) { + s.mu.Lock() + s.registry = r + s.sessID = sessID + s.fs.SetRegistry(r, sessID) + s.mu.Unlock() +} + +// SetYolo enables/disables sandbox bypass. +func (s *Server9P) SetYolo(yolo bool) { + s.mu.Lock() + s.yolo = yolo + s.fs.SetYolo(yolo) + s.mu.Unlock() +} + +// Authenticate handles secret verification. +// First call sets the secret; subsequent calls must match it. +// Returns (token, nil) on success, ("", error) on failure. +func (s *Server9P) Authenticate(clientSecret string) (string, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.secret == "" { + // First auth - set the secret + s.secret = clientSecret + s.token = randomToken() + return s.token, nil + } + + // Subsequent auth - verify secret + if clientSecret != s.secret { + return "", fmt.Errorf("authentication failed") + } + + return s.token, nil +} + +// Token returns the current session token (empty if not authenticated). +func (s *Server9P) Token() string { + s.mu.RLock() + defer s.mu.RUnlock() + return s.token +} + +// HostInfo returns platform info as key=value lines. +func (s *Server9P) HostInfo() string { + return fmt.Sprintf("platform=%s\narch=%s\n", runtime.GOOS, runtime.GOARCH) +} + +// FS returns the underlying filesystem state. +func (s *Server9P) FS() *FS9P { + return s.fs +} + +// BuildTree creates the fsedsl tree for this server. +func (s *Server9P) BuildTree() *fsedsl.Tree { + ctx := ToolsrvCtx{Server: s} + return fsedsl.BuildTree(ToolsrvSpec(), ctx) +} + +func randomToken() string { + b := make([]byte, 16) + rand.Read(b) + return hex.EncodeToString(b) +} diff --git a/cmd/toolsrv/server9p_test.go b/cmd/toolsrv/server9p_test.go new file mode 100644 index 0000000..10f541e --- /dev/null +++ b/cmd/toolsrv/server9p_test.go @@ -0,0 +1,134 @@ +package main + +import ( + "testing" +) + +func TestServer9P_AuthenticateFirst(t *testing.T) { + srv := NewServer9P() + + secret := "test-secret-0123456789abcdef" + + // First auth sets the secret and returns a token + token, err := srv.Authenticate(secret) + if err != nil { + t.Fatalf("Authenticate failed: %v", err) + } + if token == "" { + t.Fatal("Authenticate returned empty token") + } + + // Token() should return the same value + if srv.Token() != token { + t.Errorf("Token() = %q, want %q", srv.Token(), token) + } +} + +func TestServer9P_AuthenticateSameSecret(t *testing.T) { + srv := NewServer9P() + + secret := "test-secret-0123456789abcdef" + + // First auth + token1, err := srv.Authenticate(secret) + if err != nil { + t.Fatalf("First Authenticate failed: %v", err) + } + + // Second auth with same secret should succeed and return same token + token2, err := srv.Authenticate(secret) + if err != nil { + t.Fatalf("Second Authenticate failed: %v", err) + } + + if token1 != token2 { + t.Errorf("Token changed: %q -> %q", token1, token2) + } +} + +func TestServer9P_AuthenticateWrongSecret(t *testing.T) { + srv := NewServer9P() + + secret := "test-secret-0123456789abcdef" + wrongSecret := "wrong-secret-fedcba9876543210" + + // First auth establishes the secret + _, err := srv.Authenticate(secret) + if err != nil { + t.Fatalf("First Authenticate failed: %v", err) + } + + // Second auth with wrong secret should fail + _, err = srv.Authenticate(wrongSecret) + if err == nil { + t.Fatal("Authenticate should fail with wrong secret") + } +} + +func TestServer9P_TokenBeforeAuth(t *testing.T) { + srv := NewServer9P() + + // Token should be empty before authentication + if srv.Token() != "" { + t.Errorf("Token() = %q before auth, want empty", srv.Token()) + } +} + +func TestServer9P_BuildTree(t *testing.T) { + srv := NewServer9P() + + tree := srv.BuildTree() + if tree == nil { + t.Fatal("BuildTree returned nil") + } +} + +func TestServer9P_HostInfo(t *testing.T) { + srv := NewServer9P() + + info := srv.HostInfo() + if info == "" { + t.Fatal("HostInfo returned empty string") + } + + // Should contain platform info + if !contains(info, "platform=") { + t.Error("HostInfo missing platform=") + } + if !contains(info, "arch=") { + t.Error("HostInfo missing arch=") + } +} + +func TestGenerateSecret(t *testing.T) { + s1, err := GenerateSecret() + if err != nil { + t.Fatalf("GenerateSecret failed: %v", err) + } + if s1 == "" { + t.Fatal("GenerateSecret returned empty string") + } + // 32 bytes -> 64 hex chars + if len(s1) != 64 { + t.Errorf("GenerateSecret returned %d chars, want 64", len(s1)) + } + + // Two calls should return different secrets + s2, _ := GenerateSecret() + if s1 == s2 { + t.Error("GenerateSecret returned same secret twice") + } +} + +func contains(s, substr string) bool { + return len(s) >= len(substr) && (s == substr || len(s) > 0 && containsAt(s, substr)) +} + +func containsAt(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/cmd/toolsrv/spec9p.go b/cmd/toolsrv/spec9p.go new file mode 100644 index 0000000..b308754 --- /dev/null +++ b/cmd/toolsrv/spec9p.go @@ -0,0 +1,220 @@ +// spec9p.go - fsedsl-based namespace specification for toolsrv +package main + +import ( + "context" + "fmt" + "os" + "strconv" + "strings" + + "ollie/fsedsl" +) + +// ToolsrvCtx is the context passed through the fsedsl tree. +type ToolsrvCtx struct { + Server *Server9P + Proc *Proc9P // set when traversing into /proc/{pid} +} + +// Type aliases from fsedsl, specialized to ToolsrvCtx. +type ( + FsNodeDecl = fsedsl.FsNodeDecl[ToolsrvCtx] + Binding = fsedsl.Binding[ToolsrvCtx] +) + +// Specialized fsedsl constructors for ToolsrvCtx. +var ( + DirNode = fsedsl.DirNode[ToolsrvCtx] + FileNode = fsedsl.FileNode[ToolsrvCtx] + Each = fsedsl.Each[ToolsrvCtx] + Doc = fsedsl.Doc[ToolsrvCtx] + Read = fsedsl.Read[ToolsrvCtx] + Write = fsedsl.Write[ToolsrvCtx] + BlockOnce = fsedsl.BlockOnce[ToolsrvCtx] + Request = fsedsl.Request[ToolsrvCtx] +) + +// ToolsrvSpec returns the fsedsl specification for the toolsrv namespace. +// Authentication is handled via 9P Tauth, not a /register file. +func ToolsrvSpec() FsNodeDecl { + return DirNode("/", + FileNode("ctl", 0222, + Doc("Control: write 'load ' or 'unload '"), + Write(handleCtlWrite), + ), + FileNode("tools", 0666, + Doc("Tools: read to list, write tool name to load"), + Read(handleToolsRead), + Write(handleToolsWrite), + ), + FileNode("info", 0444, + Doc("Host info: platform, arch"), + Read(handleInfoRead), + ), + DirNode("proc", + FileNode("new", 0666, + Doc("Execute tool: write token + tool + args, read result (blocking)"), + Request(handleProcNew), + ), + FileNode("new.bg", 0222, + Doc("Execute tool in background: write token + tool + args, returns pid"), + Write(handleProcNewBg), + ), + Each("{pid}", procBindings, + FileNode("out", 0444, + Doc("Process output"), + Read(handleProcOutRead), + ), + FileNode("wait", 0444, + Doc("Block until process exits, returns exit code"), + BlockOnce(handleProcWait), + ), + FileNode("stat", 0444, + Doc("Process status: running/exited, runtime, tool"), + Read(handleProcStatRead), + ), + FileNode("ctl", 0222, + Doc("Process control: write 'signal ' or 'dismiss'"), + Write(handleProcCtlWrite), + ), + ), + ), + ) +} + +// --- Handlers --- + +func handleCtlWrite(ctx ToolsrvCtx, data []byte) error { + return ctx.Server.fs.HandleCtl(strings.TrimSpace(string(data))) +} + +func handleToolsRead(ctx ToolsrvCtx) ([]byte, error) { + return []byte(ctx.Server.fs.HandleToolsRead()), nil +} + +func handleToolsWrite(ctx ToolsrvCtx, data []byte) error { + return ctx.Server.fs.HandleToolsWrite(strings.TrimSpace(string(data))) +} + +func handleInfoRead(ctx ToolsrvCtx) ([]byte, error) { + return []byte(ctx.Server.HostInfo()), nil +} + +func handleProcNew(ctx ToolsrvCtx, data []byte) ([]byte, error) { + // Pass raw payload through - fs9p.go's parsePayload will handle escaping. + // Token verification is done at connection level via Tauth. + result, err := ctx.Server.fs.HandleProcNew(context.Background(), string(data)) + if err != nil { + return nil, err + } + return []byte(result), nil +} + +func handleProcNewBg(ctx ToolsrvCtx, data []byte) error { + // Pass raw payload through - fs9p.go's parsePayload will handle escaping. + // Token verification is done at connection level via Tauth. + _, err := ctx.Server.fs.HandleProcNewBg(context.Background(), string(data)) + return err +} + +// procBindings returns bindings for each running process. +func procBindings(ctx ToolsrvCtx) ([]Binding, error) { + pids := ctx.Server.fs.ListProcs() + var bindings []Binding + + for _, pid := range pids { + proc := ctx.Server.fs.GetProc(pid) + if proc == nil { + continue + } + p := proc // capture for closure + bindings = append(bindings, Binding{ + Name: strconv.Itoa(pid), + Applier: func(c ToolsrvCtx) ToolsrvCtx { + c.Proc = p + return c + }, + }) + } + + return bindings, nil +} + +func handleProcOutRead(ctx ToolsrvCtx) ([]byte, error) { + if ctx.Proc == nil { + return nil, fmt.Errorf("no process context") + } + return []byte(ctx.Proc.Output()), nil +} + +func handleProcWait(ctx ToolsrvCtx, cctx context.Context, base string) ([]byte, string, error) { + if ctx.Proc == nil { + return nil, "", fmt.Errorf("no process context") + } + exitCode := ctx.Proc.Wait() + return []byte(fmt.Sprintf("%d\n", exitCode)), "", nil +} + +func handleProcStatRead(ctx ToolsrvCtx) ([]byte, error) { + if ctx.Proc == nil { + return nil, fmt.Errorf("no process context") + } + return []byte(ctx.Proc.Stat()), nil +} + +func handleProcCtlWrite(ctx ToolsrvCtx, data []byte) error { + if ctx.Proc == nil { + return fmt.Errorf("no process context") + } + return ctx.Server.fs.HandleProcCtl(ctx.Proc.PID, strings.TrimSpace(string(data))) +} + +// Mode constants for convenience +const ( + ModeReadOnly os.FileMode = 0444 + ModeWriteOnly os.FileMode = 0222 + ModeReadWrite os.FileMode = 0666 +) + +// parseKV parses key=value lines into a map. +func parseKV(input string) map[string]string { + m := make(map[string]string) + for _, line := range strings.Split(input, "\n") { + line = strings.TrimSpace(line) + if line == "" { + continue + } + idx := strings.Index(line, "=") + if idx < 0 { + continue + } + k := line[:idx] + v := line[idx+1:] + // Unescape \n and \\ in one pass + v = unescapeKV(v) + m[k] = v + } + return m +} + +// unescapeKV reverses escape sequences: \n -> newline, \\ -> backslash +func unescapeKV(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) +} diff --git a/cmd/ollie-remote/test/.gitignore b/cmd/toolsrv/test/.gitignore similarity index 100% rename from cmd/ollie-remote/test/.gitignore rename to cmd/toolsrv/test/.gitignore diff --git a/cmd/ollie-remote/test/.ollie-sandbox.yaml b/cmd/toolsrv/test/.ollie-sandbox.yaml similarity index 100% rename from cmd/ollie-remote/test/.ollie-sandbox.yaml rename to cmd/toolsrv/test/.ollie-sandbox.yaml diff --git a/cmd/ollie-remote/test/Containerfile b/cmd/toolsrv/test/Containerfile similarity index 100% rename from cmd/ollie-remote/test/Containerfile rename to cmd/toolsrv/test/Containerfile diff --git a/cmd/ollie-remote/test/hello.go b/cmd/toolsrv/test/hello.go similarity index 100% rename from cmd/ollie-remote/test/hello.go rename to cmd/toolsrv/test/hello.go diff --git a/cmd/ollie-remote/test/run.sh b/cmd/toolsrv/test/run.sh similarity index 100% rename from cmd/ollie-remote/test/run.sh rename to cmd/toolsrv/test/run.sh diff --git a/data/tools/tool_load b/data/tools/tool_load index 8fa3ccd..5ad5c1a 100644 --- a/data/tools/tool_load +++ b/data/tools/tool_load @@ -3,8 +3,7 @@ # ollie:prompt # ## tool_load # -# Load a tool into the current session/agent by writing its name to -# `session/{sid}/agent/{aid}/tools` via ollie-9p. +# Load a tool into the current session/agent. # # **Args**: `name` (required) — the tool name to load. # @@ -15,6 +14,7 @@ set -euo pipefail +# This tool just outputs the name - the actual loading is handled by +# the agent framework when it receives the tool result. name=$(grep -oP '"name"\s*:\s*"\K[^"]+') -echo "$name" | ollie-9p write "session/$OLLIE_SESSION_ID/agent/$OLLIE_UNAME/tools" -echo "loaded: $name" \ No newline at end of file +echo "STATUS=ok TOOL_LOAD=$name" diff --git a/justfile b/justfile index 5830d35..5fb5fb9 100644 --- a/justfile +++ b/justfile @@ -22,13 +22,13 @@ kf5: build-kf5 install-kf5 # === Build targets === -build: core ninep lib9p ollie-remote lsp-tools kde +build: core ninep lib9p toolsrv lsp-tools kde -build-kf5: core ninep lib9p ollie-remote kde-kf5 +build-kf5: core ninep lib9p toolsrv kde-kf5 core: @echo === Building ollie-core === - go build ./agent/... ./backend/... ./bypass/... ./env/... ./fs/... ./log/... ./paths/... ./prompts/... ./sandbox/... ./session/... ./tools/lsp/... ./toolsrv/... + go build ./... ninep: @echo === Building ollie-9p === @@ -61,9 +61,9 @@ lsp-tools: go build -o {{cfg}}/tools/lsp_diagnostics ./tools/lsp/cmd/lsp_diagnostics/ install -m644 ./tools/lsp/cmd/lsp_diagnostics/lsp_diagnostics.meta {{cfg}}/tools/lsp_diagnostics.meta -ollie-remote: - @echo === Building ollie-remote === - go build -buildvcs=false -o {{bin}}/ollie-remote ./cmd/ollie-remote/ +toolsrv: + @echo === Building toolsrv === + go build -buildvcs=false -o {{bin}}/toolsrv ./cmd/toolsrv/ kde: @echo === Building ollie-kde === @@ -92,7 +92,7 @@ install-el: migrate-bin: #!/usr/bin/env sh old="$HOME/bin" - for f in olliesrv ollie-9p ollie-remote ollie-gui ollie-remount logseq-cli o; do + for f in olliesrv ollie-9p ollie-remote toolsrv ollie-gui ollie-remount logseq-cli o; do if [ -f "$old/$f" ]; then echo " removing stale $old/$f" rm -f "$old/$f" @@ -127,8 +127,8 @@ install-data: install -m644 data/prompts/user-preferences.md {{cfg}}/prompts/user-preferences.md # Sandbox mkdir -p {{cfg}}/sandbox - install -m644 sandbox/default.yaml {{cfg}}/sandbox/default.yaml - install -m644 sandbox/restricted.yaml {{cfg}}/sandbox/restricted.yaml + install -m644 cmd/toolsrv/internal/sandbox/default.yaml {{cfg}}/sandbox/default.yaml + install -m644 cmd/toolsrv/internal/sandbox/restricted.yaml {{cfg}}/sandbox/restricted.yaml # Skills mkdir -p {{cfg}}/skills/acme-api {{cfg}}/skills/agent-kb {{cfg}}/skills/duckduckgo-search {{cfg}}/skills/github-cli {{cfg}}/skills/rc-shell \ {{cfg}}/skills/web-dev-browser-screencapture {{cfg}}/skills/write-skills @@ -175,6 +175,7 @@ install-data: install -m755 data/tools/logseq {{cfg}}/tools/logseq install -m755 data/tools/skill_list {{cfg}}/tools/skill_list install -m755 data/tools/skill_load {{cfg}}/tools/skill_load + install -m755 data/tools/tool_load {{cfg}}/tools/tool_load install -m755 data/tools/system_logs_dmesg {{cfg}}/tools/system_logs_dmesg install -m755 data/tools/system_logs_journald {{cfg}}/tools/system_logs_journald # Tool metadata (.meta sidecar files) @@ -252,20 +253,20 @@ test: test-core test-9p test-core: @echo === Testing ollie-core === - go test ./agent/... ./backend/... ./bypass/... ./env/... ./log/... ./paths/... ./sandbox/... ./session/... ./toolsrv/... ./prompts/... + go test ./cmd/olliesrv/... ./cmd/toolsrv/... ./env/... ./log/... ./paths/... ./toolsrv/... test-9p: @echo === Testing ollie-9p === - go test ./fs/... + go test ./lib9p/... test-remote: - @echo === Testing ollie-remote === - go build ./cmd/ollie-remote/... + @echo === Testing toolsrv === + go build ./cmd/toolsrv/... # === Clean === clean: - go clean ./agent/... ./backend/... 2>/dev/null || true + go clean ./cmd/... 2>/dev/null || true rm -f {{bin}}/Ollie {{bin}}/Kmpl {{bin}}/Oi {{bin}}/Theo rm -f {{bin}}/ollie-remote just -f kde/justfile clean diff --git a/toolsrv/client9p.go b/toolsrv/client9p.go new file mode 100644 index 0000000..f373073 --- /dev/null +++ b/toolsrv/client9p.go @@ -0,0 +1,468 @@ +// client9p.go - 9P client for toolsrv. +package toolsrv + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "syscall" + + "9fans.net/go/plan9" + p9client "9fans.net/go/plan9/client" +) + +// Conn is a 9P connection to a toolsrv. +type Conn struct { + fsys *p9client.Fsys + conn *p9client.Conn + token string // session token from auth + secret string // secret used for auth (for reconnect) + onToolsChanged func() // callback for tool changes (client must poll) +} + +// Dial connects to a toolsrv at the given Unix socket path and authenticates. +// If secret is empty, a random one is generated (first connection). +// Returns the connection and the secret used (caller should save for reconnect). +func Dial(socketPath string, secret string) (*Conn, error) { + conn, err := p9client.Dial("unix", socketPath) + if err != nil { + return nil, err + } + + // Generate secret if not provided + if secret == "" { + secret = randomSecret() + } + + // Authenticate via Tauth + afid, err := conn.Auth("agent", "") + if err != nil { + conn.Close() + return nil, fmt.Errorf("auth: %w", err) + } + + // Write secret to auth fid + if _, err := afid.Write([]byte(secret)); err != nil { + afid.Close() + conn.Close() + return nil, fmt.Errorf("write secret: %w", err) + } + + // Read token from auth fid (use ReadAt to read from offset 0, + // since Write advances the file offset) + tokenBuf := make([]byte, 1024) + n, err := afid.ReadAt(tokenBuf, 0) + if err != nil && err != io.EOF { + afid.Close() + conn.Close() + return nil, fmt.Errorf("read token: %w", err) + } + token := string(tokenBuf[:n]) + // Trim newline + if len(token) > 0 && token[len(token)-1] == '\n' { + token = token[:len(token)-1] + } + + // Attach using the authenticated afid + fsys, err := conn.Attach(afid, "agent", "") + if err != nil { + afid.Close() + conn.Close() + return nil, fmt.Errorf("attach: %w", err) + } + + afid.Close() + + return &Conn{ + fsys: fsys, + conn: conn, + token: token, + secret: secret, + }, nil +} + +// Secret returns the secret used for authentication. +// Save this for reconnecting to the same toolsrv instance. +func (c *Conn) Secret() string { + return c.secret +} + +// Token returns the session token. +func (c *Conn) Token() string { + return c.token +} + +// ListTools returns the list of loaded tools. +func (c *Conn) ListTools() ([]ToolInfo, error) { + fid, err := c.fsys.Open("tools", plan9.OREAD) + if err != nil { + return nil, err + } + defer fid.Close() + + data, err := io.ReadAll(fid) + if err != nil { + return nil, err + } + + var tools []ToolInfo + if err := json.Unmarshal(data, &tools); err != nil { + return nil, fmt.Errorf("parse tools: %w", err) + } + return tools, nil +} + +// CallTool executes a tool and returns the result. +func (c *Conn) CallTool(ctx context.Context, name string, args json.RawMessage) (json.RawMessage, error) { + fid, err := c.fsys.Open("proc/new", plan9.ORDWR) + if err != nil { + return nil, fmt.Errorf("open proc/new: %w", err) + } + defer fid.Close() + + // Convert JSON args to key=value format + var argMap map[string]interface{} + if err := json.Unmarshal(args, &argMap); err != nil { + argMap = make(map[string]interface{}) + } + + payload := fmt.Sprintf("token=%s\ntool=%s\n", c.token, name) + for k, v := range argMap { + escaped := escapeValue(fmt.Sprintf("%v", v)) + payload += fmt.Sprintf("%s=%s\n", k, escaped) + } + + if _, err := fid.Write([]byte(payload)); err != nil { + return nil, fmt.Errorf("write: %w", err) + } + + // Read result from offset 0 (Write advances the file offset, but + // for request-response files we need to read the result from the start) + var result []byte + buf := make([]byte, 8192) + for offset := int64(0); ; { + n, err := fid.ReadAt(buf, offset) + if n > 0 { + result = append(result, buf[:n]...) + offset += int64(n) + } + if err == io.EOF || n == 0 { + break + } + if err != nil { + return nil, fmt.Errorf("read: %w", err) + } + } + + return json.RawMessage(result), nil +} + +// LoadTool loads a tool by name. +func (c *Conn) LoadTool(name string) error { + fid, err := c.fsys.Open("ctl", plan9.OWRITE) + if err != nil { + return err + } + defer fid.Close() + _, err = fid.Write([]byte("load " + name + "\n")) + return err +} + +// UnloadTool unloads a tool by name. +func (c *Conn) UnloadTool(name string) error { + fid, err := c.fsys.Open("ctl", plan9.OWRITE) + if err != nil { + return err + } + defer fid.Close() + _, err = fid.Write([]byte("unload " + name + "\n")) + return err +} + +// Ping checks if the connection is alive. +func (c *Conn) Ping() error { + _, err := c.fsys.Stat("info") + return err +} + +// Close closes the connection. +func (c *Conn) Close() { + if c.conn != nil { + c.conn.Close() + } +} + +// HostInfo returns host information from the server. +func (c *Conn) HostInfo() (HostInfo, error) { + fid, err := c.fsys.Open("info", plan9.OREAD) + if err != nil { + return HostInfo{}, err + } + defer fid.Close() + + data, err := io.ReadAll(fid) + if err != nil { + return HostInfo{}, err + } + + info := HostInfo{} + for _, line := range splitLines(string(data)) { + if idx := indexOf(line, '='); idx >= 0 { + key := line[:idx] + val := line[idx+1:] + switch key { + case "platform": + info.Platform = val + case "git": + info.IsGitRepo = val == "true" + } + } + } + return info, nil +} + +// --- Environment and Configuration --- + +// SetEnv sets an environment variable via ctl command. +func (c *Conn) SetEnv(key, value string) error { + fid, err := c.fsys.Open("ctl", plan9.OWRITE) + if err != nil { + return err + } + defer fid.Close() + _, err = fid.Write([]byte(fmt.Sprintf("env %s=%s\n", key, value))) + return err +} + +// SetCWD sets the working directory via ctl command. +func (c *Conn) SetCWD(dir string) error { + fid, err := c.fsys.Open("ctl", plan9.OWRITE) + if err != nil { + return err + } + defer fid.Close() + _, err = fid.Write([]byte(fmt.Sprintf("cwd %s\n", dir))) + return err +} + +// SetOnToolsChanged stores a callback that is invoked when tools change. +// In 9P model, this is handled by polling ToolRegistryRevision(). +// The callback is stored but you must call CheckToolsChanged() periodically. +func (c *Conn) SetOnToolsChanged(fn func()) { + c.onToolsChanged = fn +} + +// ToolRegistryRevision returns the current tool registry revision. +// Returns 0 if not available (always refresh). +func (c *Conn) ToolRegistryRevision() uint64 { + fid, err := c.fsys.Open("info", plan9.OREAD) + if err != nil { + return 0 + } + defer fid.Close() + + data, err := io.ReadAll(fid) + if err != nil { + return 0 + } + + for _, line := range splitLines(string(data)) { + if idx := indexOf(line, '='); idx >= 0 { + if line[:idx] == "tools_rev" { + var rev uint64 + fmt.Sscanf(line[idx+1:], "%d", &rev) + return rev + } + } + } + return 0 +} + +// --- Process management --- + +// Detach detaches the current foreground process to background. +// Returns true if successful. +func (c *Conn) Detach() bool { + fid, err := c.fsys.Open("ctl", plan9.OWRITE) + if err != nil { + return false + } + defer fid.Close() + _, err = fid.Write([]byte("detach\n")) + return err == nil +} + +// ListDetachedRaw returns raw info about detached processes. +func (c *Conn) ListDetachedRaw() []any { + fid, err := c.fsys.Open("proc", plan9.OREAD) + if err != nil { + return nil + } + defer fid.Close() + + var result []any + // Read directory entries + for { + dirs, err := fid.Dirread() + if err != nil || len(dirs) == 0 { + break + } + for _, d := range dirs { + if d.Name == "new" || d.Name == "new.bg" { + continue + } + info := c.readProcInfo(d.Name) + if info != nil { + result = append(result, info) + } + } + } + return result +} + +// readProcInfo reads info about a single proc. +func (c *Conn) readProcInfo(name string) map[string]any { + path := fmt.Sprintf("proc/%s/stat", name) + fid, err := c.fsys.Open(path, plan9.OREAD) + if err != nil { + return nil + } + defer fid.Close() + + data, err := io.ReadAll(fid) + if err != nil { + return nil + } + + info := make(map[string]any) + for _, line := range splitLines(string(data)) { + if idx := indexOf(line, '='); idx >= 0 { + key := line[:idx] + val := line[idx+1:] + switch key { + case "pid": + var pid int + fmt.Sscanf(val, "%d", &pid) + info["pid"] = pid + case "command": + info["command"] = val + case "started": + var ts int64 + fmt.Sscanf(val, "%d", &ts) + info["started"] = ts + case "exited": + info["exited"] = val == "true" + case "exit_code": + var code int + fmt.Sscanf(val, "%d", &code) + info["exit_code"] = code + } + } + } + return info +} + +// SignalDetached sends a signal to a detached process. +func (c *Conn) SignalDetached(pid int, sig syscall.Signal) error { + path := fmt.Sprintf("proc/%d/ctl", pid) + fid, err := c.fsys.Open(path, plan9.OWRITE) + if err != nil { + return err + } + defer fid.Close() + _, err = fid.Write([]byte(fmt.Sprintf("signal %d\n", sig))) + return err +} + +// DismissDetached dismisses a detached process. +func (c *Conn) DismissDetached(pid int) bool { + path := fmt.Sprintf("proc/%d/ctl", pid) + fid, err := c.fsys.Open(path, plan9.OWRITE) + if err != nil { + return false + } + defer fid.Close() + _, err = fid.Write([]byte("dismiss\n")) + return err == nil +} + +// GetDetachedOutput gets output from a detached process. +func (c *Conn) GetDetachedOutput(pid int) (string, error) { + path := fmt.Sprintf("proc/%d/out", pid) + fid, err := c.fsys.Open(path, plan9.OREAD) + if err != nil { + return "", err + } + defer fid.Close() + data, err := io.ReadAll(fid) + return string(data), err +} + +// --- Helper types --- + +// HostInfo contains information about the host running toolsrv. +type HostInfo struct { + Platform string + IsGitRepo bool +} + +// RateLimitedError indicates a rate limit was hit. +type RateLimitedError struct { + Err error + Remaining int +} + +func (e RateLimitedError) Error() string { + return e.Err.Error() +} + +// --- Helpers --- + +func randomSecret() string { + b := make([]byte, 16) + rand.Read(b) + return hex.EncodeToString(b) +} + +func splitLines(s string) []string { + var lines []string + start := 0 + for i := 0; i < len(s); i++ { + if s[i] == '\n' { + lines = append(lines, s[start:i]) + start = i + 1 + } + } + if start < len(s) { + lines = append(lines, s[start:]) + } + return lines +} + +func indexOf(s string, c byte) int { + for i := 0; i < len(s); i++ { + if s[i] == c { + return i + } + } + return -1 +} + +func escapeValue(s string) string { + // Escape newlines and backslashes + result := make([]byte, 0, len(s)) + for i := 0; i < len(s); i++ { + switch s[i] { + case '\n': + result = append(result, '\\', 'n') + case '\\': + result = append(result, '\\', '\\') + default: + result = append(result, s[i]) + } + } + return string(result) +} diff --git a/toolsrv/conn.go b/toolsrv/conn.go deleted file mode 100644 index fa6d18a..0000000 --- a/toolsrv/conn.go +++ /dev/null @@ -1,304 +0,0 @@ -package toolsrv - -import ( - "bufio" - "context" - "encoding/json" - "fmt" - "io" - "sync" - "sync/atomic" - "syscall" -) - -// Conn is a unified JSON-RPC client connected to a tool server subprocess -// (ollie-remote). It satisfies Runner and all extended interfaces that the -// agent layer uses via type assertions. Both local and remote connections -// return the same *Conn — the only difference is the underlying transport. -type Conn struct { - writeMu sync.Mutex - pendingMu sync.Mutex - pending map[int64]*pendingCall - closed chan struct{} - closeOnce sync.Once - readErr error - - enc *json.Encoder - dec *json.Decoder - nextID atomic.Int64 - Info HostInfo - - // cleanup is called on Close (kills subprocess, closes SSH, etc.) - cleanup func() - - // onToolsChanged is a callback for tool directory change notifications. - onToolsChanged func() -} - -// NewConn wraps an io.ReadWriteCloser as a JSON-RPC connection. -// The cleanup function is called when Close() is invoked. -func NewConn(rwc io.ReadWriteCloser, cleanup func()) *Conn { - c := &Conn{ - enc: json.NewEncoder(rwc), - dec: json.NewDecoder(bufio.NewReader(rwc)), - cleanup: cleanup, - pending: make(map[int64]*pendingCall), - closed: make(chan struct{}), - } - go c.readLoop() - return c -} - -// NewConnSplit creates a Conn from separate reader/writer (e.g. stdin/stdout pipes). -func NewConnSplit(r io.Reader, w io.WriteCloser, cleanup func()) *Conn { - c := &Conn{ - enc: json.NewEncoder(w), - dec: json.NewDecoder(bufio.NewReader(r)), - cleanup: cleanup, - pending: make(map[int64]*pendingCall), - closed: make(chan struct{}), - } - go c.readLoop() - return c -} - -// --- Runner interface --- - -func (c *Conn) ListTools() ([]ToolInfo, error) { - resp, err := c.call("list_tools", nil) - if err != nil { - return nil, err - } - var infos []ToolInfo - if err := json.Unmarshal(resp, &infos); err != nil { - return nil, fmt.Errorf("list_tools unmarshal: %w", err) - } - return infos, nil -} - -func (c *Conn) CallTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) { - return c.callStreaming(ctx, tool, args) -} - -// --- Extended interfaces (used by agent via type assertion) --- - -func (c *Conn) SetEnv(key, value string) { - params, _ := json.Marshal(map[string]string{"key": key, "value": value}) - c.call("set_env", params) -} - -func (c *Conn) SetCWD(dir string) { - params, _ := json.Marshal(map[string]string{"dir": dir}) - c.call("set_cwd", params) -} - -func (c *Conn) SetOnToolsChanged(fn func()) { - c.pendingMu.Lock() - c.onToolsChanged = fn - c.pendingMu.Unlock() -} - -// ToolRegistryRevision returns 0 for remote connections (revision tracking -// is not supported over RPC; use SetOnToolsChanged for push-based refresh). -func (c *Conn) ToolRegistryRevision() uint64 { return 0 } - -func (c *Conn) Detach() bool { - resp, err := c.call("detach", nil) - if err != nil { - return false - } - var ok bool - json.Unmarshal(resp, &ok) - return ok -} - -func (c *Conn) ListDetachedRaw() []any { - resp, err := c.call("list_detached", nil) - if err != nil { - return nil - } - var out []any - json.Unmarshal(resp, &out) - return out -} - -func (c *Conn) SignalDetached(pid int, sig syscall.Signal) error { - params, _ := json.Marshal(map[string]any{"pid": pid, "signal": int(sig)}) - _, err := c.call("signal_detached", params) - return err -} - -func (c *Conn) GetDetachedOutput(pid int) (string, error) { - params, _ := json.Marshal(map[string]any{"pid": pid}) - resp, err := c.call("get_detached_output", params) - if err != nil { - return "", err - } - var out string - json.Unmarshal(resp, &out) - return out, nil -} - -func (c *Conn) DismissDetached(pid int) bool { - params, _ := json.Marshal(map[string]any{"pid": pid}) - resp, err := c.call("dismiss_detached", params) - if err != nil { - return false - } - var ok bool - json.Unmarshal(resp, &ok) - return ok -} - -// Close shuts down the connection and subprocess. -func (c *Conn) Close() { - c.fail(fmt.Errorf("connection closed")) - if c.cleanup != nil { - c.cleanup() - } -} - -// Ping verifies the remote server is responsive. -func (c *Conn) Ping() error { - _, err := c.call("ping", nil) - return err -} - -// FetchHostInfo retrieves environment details from the server. -func (c *Conn) FetchHostInfo() (HostInfo, error) { - resp, err := c.call("host_info", nil) - if err != nil { - return HostInfo{}, err - } - var info HostInfo - if err := json.Unmarshal(resp, &info); err != nil { - return HostInfo{}, err - } - return info, nil -} - -type pendingCall struct { - response chan rpcResponse - ctx context.Context -} - -func (c *Conn) readLoop() { - for { - var resp rpcResponse - if err := c.dec.Decode(&resp); err != nil { - c.fail(fmt.Errorf("rpc read: %w", err)) - return - } - - c.pendingMu.Lock() - call := c.pending[resp.ID] - c.pendingMu.Unlock() - if call == nil { - continue - } - - // Tool output is sent as an intermediate response with the request ID. - var notif outputNotification - if resp.Stream && resp.Error == nil && json.Unmarshal(resp.Result, ¬if) == nil { - if notif.Data != "" { - StreamOutput(call.ctx, notif.Data) - } - continue - } - - c.pendingMu.Lock() - delete(c.pending, resp.ID) - c.pendingMu.Unlock() - call.response <- resp - } -} - -func (c *Conn) fail(err error) { - c.closeOnce.Do(func() { - c.pendingMu.Lock() - c.readErr = err - close(c.closed) - for id, call := range c.pending { - delete(c.pending, id) - call.response <- rpcResponse{Error: &rpcError{Code: -32000, Message: err.Error()}} - } - c.pendingMu.Unlock() - }) -} - -// --- JSON-RPC internals --- - -func (c *Conn) call(method string, params json.RawMessage) (json.RawMessage, error) { - return c.callStreaming(context.Background(), method, params) -} - -func (c *Conn) callStreaming(ctx context.Context, method string, params json.RawMessage) (json.RawMessage, error) { - id := c.nextID.Add(1) - call := &pendingCall{response: make(chan rpcResponse, 1), ctx: ctx} - c.pendingMu.Lock() - select { - case <-c.closed: - c.pendingMu.Unlock() - return nil, fmt.Errorf("rpc %s: connection closed", method) - default: - } - c.pending[id] = call - c.pendingMu.Unlock() - - c.writeMu.Lock() - err := c.enc.Encode(rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params}) - c.writeMu.Unlock() - if err != nil { - c.pendingMu.Lock() - delete(c.pending, id) - c.pendingMu.Unlock() - return nil, fmt.Errorf("rpc write %s: %w", method, err) - } - - select { - case resp := <-call.response: - if resp.Error != nil { - return nil, fmt.Errorf("rpc %s: %s", method, resp.Error.Message) - } - return resp.Result, nil - case <-ctx.Done(): - c.pendingMu.Lock() - delete(c.pending, id) - c.pendingMu.Unlock() - return nil, ctx.Err() - case <-c.closed: - c.pendingMu.Lock() - err := c.readErr - c.pendingMu.Unlock() - if err == nil { - err = fmt.Errorf("connection closed") - } - return nil, fmt.Errorf("rpc %s: %w", method, err) - } -} - -// JSON-RPC 2.0 types shared between Conn (client) and transport layer. - -type rpcRequest struct { - JSONRPC string `json:"jsonrpc"` - ID int64 `json:"id"` - Method string `json:"method"` - Params json.RawMessage `json:"params,omitempty"` -} - -type rpcResponse struct { - JSONRPC string `json:"jsonrpc"` - ID int64 `json:"id"` - Result json.RawMessage `json:"result,omitempty"` - Error *rpcError `json:"error,omitempty"` - Stream bool `json:"stream,omitempty"` -} - -type rpcError struct { - Code int `json:"code"` - Message string `json:"message"` -} - -type outputNotification struct { - Data string `json:"data"` -} diff --git a/toolsrv/detach.go b/toolsrv/detach.go deleted file mode 100644 index 8464b2c..0000000 --- a/toolsrv/detach.go +++ /dev/null @@ -1,99 +0,0 @@ -package toolsrv - -import ( - "sync" - "time" -) - -const RingBufSize = 64 * 1024 // 64KB ring buffer per detached process - -// DetachedProcess represents a process that the agent has detached from -// but which continues running. The user can view its output and signal it. -type DetachedProcess struct { - PID int - Command string - Started time.Time - Exited bool - ExitCode int - - Ring *RingBuffer - Done chan struct{} - Mu sync.Mutex -} - -// Info returns a plain-data snapshot of this process for external consumers. -type InfoData struct { - PID int - Command string - Started int64 // unix timestamp - Exited bool - ExitCode int -} - -func (p *DetachedProcess) Info() InfoData { - p.Mu.Lock() - defer p.Mu.Unlock() - return InfoData{ - PID: p.PID, - Command: p.Command, - Started: p.Started.Unix(), - Exited: p.Exited, - ExitCode: p.ExitCode, - } -} - -// Output returns the current contents of the ring buffer. -func (p *DetachedProcess) Output() string { - p.Mu.Lock() - defer p.Mu.Unlock() - return p.Ring.String() -} - -// ringBuffer is a fixed-size circular byte buffer. -type RingBuffer struct { - buf []byte - size int - pos int - full bool -} - -func NewRingBuffer(size int) *RingBuffer { - return &RingBuffer{buf: make([]byte, size), size: size} -} - -// Write implements io.Writer. -func (r *RingBuffer) Write(p []byte) (int, error) { - n := len(p) - if n >= r.size { - // Data larger than buffer: just keep the tail - copy(r.buf, p[n-r.size:]) - r.pos = 0 - r.full = true - return n, nil - } - if r.pos+n <= r.size { - copy(r.buf[r.pos:], p) - } else { - first := r.size - r.pos - copy(r.buf[r.pos:], p[:first]) - copy(r.buf, p[first:]) - r.full = true - } - r.pos = (r.pos + n) % r.size - if r.pos == 0 && n > 0 { - r.full = true - } - return n, nil -} - -// String returns the buffer contents in order. -func (r *RingBuffer) String() string { - if !r.full { - return string(r.buf[:r.pos]) - } - // Buffer has wrapped: data from pos..end + 0..pos - out := make([]byte, r.size) - copy(out, r.buf[r.pos:]) - copy(out[r.size-r.pos:], r.buf[:r.pos]) - return string(out) -} diff --git a/toolsrv/dial.go b/toolsrv/dial.go deleted file mode 100644 index 0c582b9..0000000 --- a/toolsrv/dial.go +++ /dev/null @@ -1,25 +0,0 @@ -package toolsrv - -import ( - "context" - "fmt" - "net" -) - -// Dial connects to a running tool server at the given Unix socket path. -// The server must already be running (started via Spawn or SpawnRemote). -func Dial(ctx context.Context, socket string) (*Conn, error) { - if socket == "" { - return nil, fmt.Errorf("toolsrv.Dial: socket path required") - } - conn, err := net.Dial("unix", socket) - if err != nil { - return nil, fmt.Errorf("toolsrv.Dial: %w", err) - } - c := NewConn(conn, func() { conn.Close() }) - if err := c.Ping(); err != nil { - c.Close() - return nil, fmt.Errorf("toolsrv.Dial: ping failed: %w", err) - } - return c, nil -} diff --git a/toolsrv/exec.go b/toolsrv/exec.go deleted file mode 100644 index 9d1af6f..0000000 --- a/toolsrv/exec.go +++ /dev/null @@ -1,594 +0,0 @@ -package toolsrv - -import ( - "bytes" - "context" - "encoding/binary" - "encoding/json" - "fmt" - "io" - "net" - "os" - "os/exec" - "path/filepath" - "strings" - "sync" - "syscall" - "time" - - "ollie/paths" - "ollie/sandbox" -) - -func loadSandboxConfig(name string) (*sandbox.Config, error) { - if name == "" { - name = "default" - } - path := filepath.Join(paths.CfgDir(), "sandbox", name+".yaml") - f, err := os.Open(path) - if err != nil { - return nil, fmt.Errorf("sandbox %q not found: %w", name, err) - } - defer f.Close() - return sandbox.LoadSandbox(f) -} - -type limitedWriter struct { - mu sync.Mutex - w io.Writer - written int - limit int - truncated bool - stream func(string) // if non-nil, called with each chunk of output -} - -func (lw *limitedWriter) Write(p []byte) (n int, err error) { - lw.mu.Lock() - defer lw.mu.Unlock() - if lw.written >= lw.limit { - lw.truncated = true - return len(p), nil - } - - remaining := lw.limit - lw.written - toWrite := p - if len(p) > remaining { - toWrite = p[:remaining] - lw.truncated = true - } - - written, err := lw.w.Write(toWrite) - lw.written += written - if lw.stream != nil { - lw.stream(string(toWrite[:written])) - } - if err != nil { - return written, err - } - return len(p), nil -} - -// Connects to the broker socket, sends the request with the current env, -// and streams the framed response back. -func (s *Server) executeBypass(ctx context.Context, cmd, dir string, timeout int, doDetach ...bool) (string, error) { - return s.executeBypassOpts(ctx, cmd, dir, timeout, false, doDetach...) -} - -func (s *Server) executeBypassSudo(ctx context.Context, cmd, dir string, timeout int) (string, error) { - return s.executeBypassOpts(ctx, cmd, dir, timeout, true) -} - -func (s *Server) executeBypassOpts(ctx context.Context, cmd, dir string, timeout int, sudo bool, doDetach ...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") - - wantDetach := len(doDetach) > 0 && doDetach[0] - - var cancel context.CancelFunc - if wantDetach || timeout <= 0 { - ctx, cancel = context.WithCancel(ctx) - } else { - ctx, cancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Second) - } - - // Connect to broker - conn, err := net.DialTimeout("unix", sockPath, 5*time.Second) - if err != nil { - cancel() - return "", fmt.Errorf("bypass not available: %w", err) - } - - // Send request — merge process env with session-scoped envExtra. - envMap := make(map[string]string, len(os.Environ())) - for _, kv := range os.Environ() { - if k, v, ok := strings.Cut(kv, "="); ok { - envMap[k] = v - } - } - s.envMu.RLock() - for k, v := range s.envExtra { - envMap[k] = v - } - s.envMu.RUnlock() - reqJSON, _ := json.Marshal(struct { - Cmd string `json:"cmd"` - Cwd string `json:"cwd"` - Env map[string]string `json:"env"` - Session string `json:"session,omitempty"` - Sudo bool `json:"sudo,omitempty"` - }{Cmd: cmd, Cwd: dir, Env: envMap, Session: s.sessionID, Sudo: sudo}) - reqJSON = append(reqJSON, '\n') - if _, err := conn.Write(reqJSON); err != nil { - conn.Close() - cancel() - return "", fmt.Errorf("bypass execution failed: write: %w", err) - } - - // Set up detach channel - detachCh := make(chan struct{}, 1) - s.detachMu.Lock() - s.detachCh = detachCh - s.detachMu.Unlock() - defer func() { - s.detachMu.Lock() - s.detachCh = nil - s.detachMu.Unlock() - }() - - if wantDetach { - close(detachCh) - } - - // readFrames reads from the connection, writing to w and streaming. - // Returns exit code when 'x' frame arrives or -1 on error. - readFrames := func(w io.Writer, stream func(string)) int { - header := make([]byte, 5) - for { - conn.SetReadDeadline(time.Now().Add(1 * time.Second)) //nolint:errcheck - _, err := io.ReadFull(conn, header) - if err != nil { - if ne, ok := err.(net.Error); ok && ne.Timeout() { - // Check for context cancellation - 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)) //nolint:errcheck - if _, err := io.ReadFull(conn, payload); err != nil { - return -1 - } - } - - switch frameType { - case 'd': - w.Write(payload) //nolint:errcheck - if stream != nil { - stream(string(payload)) - } - case 'x': - var code int - fmt.Sscanf(string(payload), "%d", &code) - return code - } - } - } - - // Check if detach was requested - select { - case <-detachCh: - // Detach: background the socket read into a goroutine - cmdStr := "bypass: " + cmd - if len(cmdStr) > 80 { - cmdStr = cmdStr[:77] + "..." - } - ring := NewRingBuffer(RingBufSize) - pid := int(time.Now().UnixNano() & 0x7FFFFFFF) // synthetic PID - - proc := &DetachedProcess{ - PID: pid, - Command: cmdStr, - Started: time.Now(), - Ring: ring, - Done: make(chan struct{}), - } - s.detachMu.Lock() - s.detached = append(s.detached, proc) - s.detachMu.Unlock() - - go func() { - defer conn.Close() - defer cancel() - exitCode := readFrames(ring, nil) - proc.Mu.Lock() - proc.Exited = true - proc.ExitCode = exitCode - proc.Mu.Unlock() - close(proc.Done) - if s.OnExit != nil { - s.OnExit(proc.PID, proc.ExitCode) - } - }() - - if s.OnDetach != nil { - s.OnDetach(proc.PID, cmdStr) - } - return fmt.Sprintf("[detached: pid %d]", pid), nil - - default: - // Normal (foreground) execution — must also handle manual detach. - var outputBuf bytes.Buffer - streamFn := StreamFunc(ctx) - lw := &limitedWriter{ - w: &outputBuf, - limit: 10 * 1024 * 1024, - stream: streamFn, - } - - // Run readFrames in a goroutine so we can select on detach. - type frameResult struct{ exitCode int } - frameCh := make(chan frameResult, 1) - go func() { - code := readFrames(lw, nil) - frameCh <- frameResult{code} - }() - - select { - case fr := <-frameCh: - // Normal completion - conn.Close() - cancel() - combined := outputBuf.String() - if fr.exitCode != 0 { - if combined == "" { - return "", fmt.Errorf("bypass execution failed (exit %d)", fr.exitCode) - } - errOutput := combined - const maxErrOutput = 8192 - if len(errOutput) > maxErrOutput { - errOutput = errOutput[:maxErrOutput] + fmt.Sprintf("\n[...truncated %d bytes in error]", len(combined)-maxErrOutput) - } - return combined, fmt.Errorf("bypass execution failed (exit %d)\nOutput: %s", fr.exitCode, errOutput) - } - return combined, nil - - case <-ctx.Done(): - // Timeout or external cancellation - conn.Close() - cancel() - combined := outputBuf.String() - if ctx.Err() == context.DeadlineExceeded { - return combined, fmt.Errorf("bypass execution timeout after %d seconds", timeout) - } - return combined, fmt.Errorf("bypass execution interrupted") - - case <-detachCh: - // Manual detach: background the connection read - cmdStr := "bypass: " + cmd - if len(cmdStr) > 80 { - cmdStr = cmdStr[:77] + "..." - } - ring := NewRingBuffer(RingBufSize) - pid := int(time.Now().UnixNano() & 0x7FFFFFFF) - - // Splice: future output goes to ring buffer, stop streaming - lw.mu.Lock() - lw.w = ring - lw.stream = nil - lw.mu.Unlock() - - proc := &DetachedProcess{ - PID: pid, - Command: cmdStr, - Started: time.Now(), - Ring: ring, - Done: make(chan struct{}), - } - s.detachMu.Lock() - s.detached = append(s.detached, proc) - s.detachMu.Unlock() - - go func() { - fr := <-frameCh - conn.Close() - cancel() - proc.Mu.Lock() - proc.Exited = true - proc.ExitCode = fr.exitCode - proc.Mu.Unlock() - close(proc.Done) - if s.OnExit != nil { - s.OnExit(proc.PID, proc.ExitCode) - } - }() - - if s.OnDetach != nil { - s.OnDetach(proc.PID, cmdStr) - } - - partial := outputBuf.String() - return partial + fmt.Sprintf("\n[detached: pid %d]", pid), nil - } - } -} - -// executeWithStdin runs code in a sandbox and returns combined stdout+stderr. -// For languages where code is itself passed via stdin (ed, expect, bc), stdinData is ignored. -func (s *Server) executeWithStdin(ctx context.Context, code, language string, timeout int, sandboxName string, trusted bool, stdinData string, doDetach ...bool) (string, error) { - if timeout < 0 { - timeout = 30 - } - - if !trusted { - if err := s.ValidateCode(code, language); err != nil { - return "", err - } - } - - var cfg *sandbox.Config - var err error - if !s.Yolo { - cfg, err = loadSandboxConfig(sandboxName) - if err != nil { - return "", err - } - } - - s.wdMu.RLock() - workDir := s.cwd - s.wdMu.RUnlock() - if workDir == "" { - workDir, _ = os.Getwd() - } - - var cancel context.CancelFunc - if timeout > 0 { - ctx, cancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Second) - } else { - // timeout=0 means no timeout (run indefinitely, e.g. for daemons) - ctx, cancel = context.WithCancel(ctx) - } - // NOTE: do NOT defer cancel() here — the detach path must avoid cancelling - // the context (which would kill the detached process via cmd.Cancel). - // Each exit path calls cancel() explicitly except detach. - - var cmd *exec.Cmd - var interpreter []string - // codeStdin: non-empty means the code itself is fed via stdin (ed, expect, bc). - // In these cases stdinData cannot be used simultaneously. - var codeStdin string - switch language { - case "bash", "": - interpreter = []string{"bash", "-c", code} - default: - cancel() - return "", fmt.Errorf("unsupported language: %s (only bash is supported)", language) - } - s.envMu.RLock() - envMap := make(map[string]string, len(s.envExtra)) - for k, v := range s.envExtra { - envMap[k] = v - } - s.envMu.RUnlock() - for _, ev := range os.Environ() { - k, v, _ := strings.Cut(ev, "=") - if _, ok := envMap[k]; !ok { - envMap[k] = v - } - } - if _, ok := envMap["NAMESPACE"]; !ok { - if ns := plan9Namespace(envMap); ns != "" { - envMap["NAMESPACE"] = ns - } - } - getenv := func(key string) string { return envMap[key] } - - if s.Yolo { - cmd = exec.CommandContext(ctx, interpreter[0], interpreter[1:]...) - } else { - wrapped, wrapErr := sandbox.WrapCommand(cfg, interpreter, workDir, getenv) - if wrapErr != nil { - cancel() - return "", wrapErr - } - cmd = exec.CommandContext(ctx, wrapped[0], wrapped[1:]...) - } - cmd.Dir = workDir - switch { - case codeStdin != "": - cmd.Stdin = strings.NewReader(codeStdin) - case stdinData != "": - cmd.Stdin = strings.NewReader(stdinData) - } - - s.envMu.RLock() - baseEnv := os.Environ() - // Remove keys that envExtra overrides so duplicates don't leak through. - filtered := baseEnv[:0] - for _, ev := range baseEnv { - k, _, _ := strings.Cut(ev, "=") - if _, overridden := s.envExtra[k]; !overridden { - filtered = append(filtered, ev) - } - } - cmd.Env = filtered - for k, v := range s.envExtra { - cmd.Env = append(cmd.Env, k+"="+v) - } - s.envMu.RUnlock() - // Export computed NAMESPACE so tools like ollie-9p can find the server. - if ns := envMap["NAMESPACE"]; ns != "" { - cmd.Env = append(cmd.Env, "NAMESPACE="+ns) - } - - 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 - - // Set up detach channel for this execution - detachCh := make(chan struct{}, 1) - s.detachMu.Lock() - s.detachCh = detachCh - s.detachMu.Unlock() - defer func() { - s.detachMu.Lock() - s.detachCh = nil - s.detachMu.Unlock() - }() - - if err = cmd.Start(); err != nil { - cancel() - return "", fmt.Errorf("execution failed: %w", err) - } - - // If detach requested, immediately signal the detach channel - if len(doDetach) > 0 && doDetach[0] { - close(detachCh) - } - - // Wait for completion, context cancellation, or detach signal - waitCh := make(chan error, 1) - go func() { waitCh <- cmd.Wait() }() - - select { - case err = <-waitCh: - // Normal completion - cancel() - output := outputBuf.Bytes() - if lw.truncated { - output = append(output, []byte("\n[output truncated at 10MB]")...) - } - if ctx.Err() == context.DeadlineExceeded { - return "", fmt.Errorf("execution timeout after %d seconds", timeout) - } - if err != nil { - // Cap output in error message to avoid flooding the agent context. - // The full output is still returned as the first return value. - 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 - - case <-ctx.Done(): - // Context cancelled (interrupt or timeout) - cancel() - syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) - <-waitCh // reap - output := outputBuf.Bytes() - if ctx.Err() == context.DeadlineExceeded { - return string(output), fmt.Errorf("execution timeout after %d seconds", timeout) - } - return string(output), fmt.Errorf("execution interrupted") - - case <-detachCh: - // Detach: move process to registry, continue running - // Neutralize cmd.Cancel and WaitDelay so context cancellation - // won't kill the process after the timeout fires. - cmd.Cancel = nil - cmd.WaitDelay = 0 - cancel() - - cmdStr := strings.Join(interpreter, " ") - if len(cmdStr) > 80 { - cmdStr = cmdStr[:77] + "..." - } - ring := NewRingBuffer(RingBufSize) - // Splice: future output goes to ring buffer instead of outputBuf. - lw.mu.Lock() - lw.w = ring - lw.stream = nil - lw.mu.Unlock() - - proc := &DetachedProcess{ - PID: cmd.Process.Pid, - Command: cmdStr, - Started: time.Now(), - Ring: ring, - Done: make(chan struct{}), - } - s.detachMu.Lock() - s.detached = append(s.detached, proc) - s.detachMu.Unlock() - - // Monitor for exit in background - go func() { - waitErr := <-waitCh - proc.Mu.Lock() - proc.Exited = true - if waitErr != nil { - if exitErr, ok := waitErr.(*exec.ExitError); ok { - proc.ExitCode = exitErr.ExitCode() - } else { - proc.ExitCode = -1 - } - } - proc.Mu.Unlock() - close(proc.Done) - if s.OnExit != nil { - s.OnExit(proc.PID, proc.ExitCode) - } - }() - - if s.OnDetach != nil { - s.OnDetach(proc.PID, cmdStr) - } - - partial := outputBuf.String() - return partial + fmt.Sprintf("\n[detached: pid %d]", proc.PID), nil - } -} - -// plan9Namespace computes the plan9port namespace directory from the environment. -// Convention: $NAMESPACE if set, else /tmp/ns.. where display is -// $DISPLAY with trailing .0 stripped and / replaced by _. -func plan9Namespace(env map[string]string) string { - if ns := env["NAMESPACE"]; ns != "" { - return ns - } - user := env["USER"] - if user == "" { - return "" - } - disp := env["DISPLAY"] - if disp == "" { - return "" - } - // Canonicalize: strip trailing .0 - if strings.HasSuffix(disp, ".0") { - disp = disp[:len(disp)-2] - } - // Replace / with _ - disp = strings.ReplaceAll(disp, "/", "_") - return "/tmp/ns." + user + "." + disp -} diff --git a/toolsrv/processkeeper.go b/toolsrv/processkeeper.go deleted file mode 100644 index 90076fc..0000000 --- a/toolsrv/processkeeper.go +++ /dev/null @@ -1,112 +0,0 @@ -package toolsrv - -import ( - "context" - "fmt" - "sync" - "time" -) - -// ProcessKeeper maintains a running ollie-remote process and respawns it -// on failure. Multiple agents can Dial() the same keeper to get independent -// connections. If the process dies, the next Dial() triggers a respawn. -type ProcessKeeper struct { - mu sync.Mutex - ctx context.Context - proc *Process - spawn func(context.Context) (*Process, error) -} - -// NewProcessKeeper creates a keeper with an initial process and a spawn factory. -func NewProcessKeeper(ctx context.Context, proc *Process, spawn func(context.Context) (*Process, error)) *ProcessKeeper { - return &ProcessKeeper{ - ctx: ctx, - proc: proc, - spawn: spawn, - } -} - -// Dial returns a new Conn to the managed process. If the process is dead -// (socket unreachable), it respawns before dialing. -func (pk *ProcessKeeper) Dial() (*Conn, error) { - pk.mu.Lock() - sock := pk.proc.Socket - pk.mu.Unlock() - - // Try dialing the existing socket. - conn, err := Dial(pk.ctx, sock) - if err == nil { - return conn, nil - } - - // Dial failed — process likely dead. Respawn. - return pk.respawnAndDial() -} - -// Socket returns the current process socket path. -func (pk *ProcessKeeper) Socket() string { - pk.mu.Lock() - defer pk.mu.Unlock() - return pk.proc.Socket -} - -// Process returns the current underlying Process. -func (pk *ProcessKeeper) Process() *Process { - pk.mu.Lock() - defer pk.mu.Unlock() - return pk.proc -} - -// Close shuts down the managed process. -func (pk *ProcessKeeper) Close() { - pk.mu.Lock() - defer pk.mu.Unlock() - if pk.proc != nil { - pk.proc.Close() - } -} - -// SetContext updates the keeper's context. Used on session resume to -// provide a fresh context after the previous one was cancelled on pause. -func (pk *ProcessKeeper) SetContext(ctx context.Context) { - pk.mu.Lock() - defer pk.mu.Unlock() - pk.ctx = ctx -} - -func (pk *ProcessKeeper) respawnAndDial() (*Conn, error) { - pk.mu.Lock() - defer pk.mu.Unlock() - - // Another goroutine may have respawned already — try current socket first. - if conn, err := Dial(pk.ctx, pk.proc.Socket); err == nil { - return conn, nil - } - - // Clean up old process. - pk.proc.Close() - - var lastErr error - for attempt := 0; attempt < 3; attempt++ { - if pk.ctx.Err() != nil { - return nil, pk.ctx.Err() - } - proc, err := pk.spawn(pk.ctx) - if err != nil { - lastErr = err - time.Sleep(time.Duration(attempt+1) * 2 * time.Second) - continue - } - pk.proc = proc - - conn, err := Dial(pk.ctx, proc.Socket) - if err != nil { - lastErr = err - proc.Close() - time.Sleep(time.Duration(attempt+1) * 2 * time.Second) - continue - } - return conn, nil - } - return nil, fmt.Errorf("respawn failed after 3 attempts: %w", lastErr) -} diff --git a/toolsrv/registry.go b/toolsrv/registry.go index f77766c..0ca8313 100644 --- a/toolsrv/registry.go +++ b/toolsrv/registry.go @@ -1,121 +1,15 @@ +// registry.go - Tool discovery and shared types. package toolsrv import ( "encoding/json" - "fmt" "os" "path/filepath" - "sort" "strings" - "sync" "ollie/paths" ) -type Registry struct { - mu sync.RWMutex - sessions map[string]map[string]ToolInfo - revisions map[string]uint64 -} - -func NewRegistry() (*Registry, error) { - r := &Registry{ - sessions: make(map[string]map[string]ToolInfo), - revisions: make(map[string]uint64), - } - // Validate that the tools directory is readable. - dir := ToolsPath() - if _, err := os.ReadDir(dir); err != nil { - return nil, fmt.Errorf("read tools dir %s: %w", dir, err) - } - return r, nil -} - -// discover scans the tools directory and returns all valid tool infos. -func (r *Registry) discover() map[string]ToolInfo { - infos := DiscoverTools() - tools := make(map[string]ToolInfo, len(infos)) - for _, info := range infos { - tools[info.Name] = info - } - return tools -} - -func (r *Registry) Load(sessionID, name string) error { - tools := r.discover() - tool, exists := tools[name] - if !exists { - return fmt.Errorf("tool not found: %s", name) - } - - r.mu.Lock() - defer r.mu.Unlock() - - if r.sessions[sessionID] == nil { - r.sessions[sessionID] = make(map[string]ToolInfo) - } - - r.sessions[sessionID][name] = tool - r.revisions[sessionID]++ - return nil -} - -func (r *Registry) Unload(sessionID, name string) error { - r.mu.Lock() - defer r.mu.Unlock() - - sessionTools, ok := r.sessions[sessionID] - if !ok { - return nil - } - - if _, exists := sessionTools[name]; !exists { - return nil - } - - delete(sessionTools, name) - r.revisions[sessionID]++ - return nil -} - -func (r *Registry) Loaded(sessionID string) []ToolInfo { - r.mu.RLock() - defer r.mu.RUnlock() - - sessionTools, ok := r.sessions[sessionID] - if !ok { - return nil - } - - var loaded []ToolInfo - for _, info := range sessionTools { - loaded = append(loaded, info) - } - sort.Slice(loaded, func(i, j int) bool { return loaded[i].Name < loaded[j].Name }) - return loaded -} - -func (r *Registry) Lookup(sessionID, name string) (ToolInfo, bool) { - r.mu.RLock() - defer r.mu.RUnlock() - - sessionTools, ok := r.sessions[sessionID] - if !ok { - return ToolInfo{}, false - } - - tool, exists := sessionTools[name] - return tool, exists -} - -func (r *Registry) Revision(sessionID string) uint64 { - r.mu.RLock() - defer r.mu.RUnlock() - return r.revisions[sessionID] -} - - - // ToolInfo describes a tool provided by a server. type ToolInfo struct { Server string diff --git a/toolsrv/remote.go b/toolsrv/remote.go deleted file mode 100644 index 3ef6a90..0000000 --- a/toolsrv/remote.go +++ /dev/null @@ -1,52 +0,0 @@ -// Package toolsrv remote types. -// -// RemoteConfig and HostInfo are used by the transport layer and session -// creation. The actual client implementation lives in conn.go; the -// transport (SSH bootstrap) lives in transport.go. -package toolsrv - -import ( - "context" -) - -// HostInfo holds environment details from the remote host. -type HostInfo struct { - Platform string `json:"platform"` - Arch string `json:"arch"` - IsGitRepo bool `json:"is_git_repo"` -} - -// RemoteConfig holds the parameters for connecting to a remote host. -type RemoteConfig struct { - // SSHTarget is the SSH destination (e.g., "user@host" or an SSH config alias). - SSHTarget string - // CWD is the working directory on the remote host. - CWD string - // ToolsPath overrides the remote tool scripts path. - ToolsPath string - // Yolo disables sandbox enforcement on the remote. - Yolo bool -} - -// RemoteDial opens an SSH connection to the remote host and returns a *Conn. -// Convenience wrapper: spawns the remote process and dials a single connection. -func RemoteDial(ctx context.Context, cfg RemoteConfig) (*Conn, error) { - proc, err := SpawnRemote(ctx, cfg) - if err != nil { - return nil, err - } - conn, err := Dial(ctx, proc.Socket) - if err != nil { - proc.Close() - return nil, err - } - origClose := conn.cleanup - conn.cleanup = func() { - if origClose != nil { - origClose() - } - proc.Close() - } - conn.Info = proc.Info - return conn, nil -} diff --git a/toolsrv/rpc_integration_test.go b/toolsrv/rpc_integration_test.go deleted file mode 100644 index d05e45f..0000000 --- a/toolsrv/rpc_integration_test.go +++ /dev/null @@ -1,322 +0,0 @@ -package toolsrv_test - -import ( - "context" - "encoding/json" - "net" - "os" - "path/filepath" - "runtime" - "sync" - "syscall" - "testing" - - "ollie/toolsrv" -) - -// --- Test infrastructure --- - -// rpcRequest mirrors the JSON-RPC 2.0 request format (test-local). -type rpcRequest struct { - JSONRPC string `json:"jsonrpc"` - ID int64 `json:"id"` - Method string `json:"method"` - Params json.RawMessage `json:"params,omitempty"` -} - -// rpcResponse mirrors the JSON-RPC 2.0 response format (test-local). -type rpcResponse struct { - JSONRPC string `json:"jsonrpc"` - ID int64 `json:"id"` - Result json.RawMessage `json:"result,omitempty"` - Error *rpcError `json:"error,omitempty"` - Stream bool `json:"stream,omitempty"` -} - -type rpcError struct { - Code int `json:"code"` - Message string `json:"message"` -} - -// testServer wraps a toolsrv.Server with a pipe-based RPC connection. -type testServer struct { - srv *toolsrv.Server - conn net.Conn // client side - enc *json.Encoder - dec *json.Decoder - wg sync.WaitGroup -} - -func newTestServer(t *testing.T) *testServer { - t.Helper() - cwd := t.TempDir() - os.MkdirAll(filepath.Join(cwd, ".git"), 0755) - - reg, err := toolsrv.NewRegistry() - if err != nil { - t.Fatalf("NewRegistry: %v", err) - } - - sessionID := "test-session-" + t.Name() - srv := toolsrv.New(cwd) - srv.SetToolRegistry(reg, sessionID) - - clientConn, serverConn := net.Pipe() - ts := &testServer{ - srv: srv, - conn: clientConn, - enc: json.NewEncoder(clientConn), - dec: json.NewDecoder(clientConn), - } - - ctx, cancel := context.WithCancel(context.Background()) - ts.wg.Add(1) - go func() { - defer ts.wg.Done() - srv.ServeRPC(ctx, serverConn, serverConn) - }() - - t.Cleanup(func() { - cancel() - clientConn.Close() - serverConn.Close() - ts.wg.Wait() - }) - - return ts -} - -// call sends an RPC request and returns the response (skipping stream frames). -func (ts *testServer) call(method string, params any) rpcResponse { - var raw json.RawMessage - if params != nil { - raw, _ = json.Marshal(params) - } - ts.enc.Encode(rpcRequest{JSONRPC: "2.0", ID: 1, Method: method, Params: raw}) - for { - var resp rpcResponse - ts.dec.Decode(&resp) - if !resp.Stream { - return resp - } - } -} - -// --- Integration tests --- - -func TestRPC_Ping(t *testing.T) { - ts := newTestServer(t) - resp := ts.call("ping", nil) - - if resp.Error != nil { - t.Fatalf("unexpected error: %v", resp.Error.Message) - } - var result string - json.Unmarshal(resp.Result, &result) - if result != "pong" { - t.Errorf("expected 'pong', got %q", result) - } -} - -func TestRPC_HostInfo(t *testing.T) { - ts := newTestServer(t) - resp := ts.call("host_info", nil) - - if resp.Error != nil { - t.Fatalf("unexpected error: %v", resp.Error.Message) - } - - var info toolsrv.HostInfo - json.Unmarshal(resp.Result, &info) - if info.Platform != runtime.GOOS { - t.Errorf("platform: expected %q, got %q", runtime.GOOS, info.Platform) - } - if info.Arch != runtime.GOARCH { - t.Errorf("arch: expected %q, got %q", runtime.GOARCH, info.Arch) - } - if !info.IsGitRepo { - t.Error("expected is_git_repo=true (we created .git dir)") - } -} - -func TestRPC_ListTools_Empty(t *testing.T) { - ts := newTestServer(t) - resp := ts.call("list_tools", nil) - - if resp.Error != nil { - t.Fatalf("unexpected error: %v", resp.Error.Message) - } - - var tools []toolsrv.ToolInfo - json.Unmarshal(resp.Result, &tools) - if len(tools) != 0 { - t.Errorf("expected 0 tools, got %d", len(tools)) - } -} - -func TestRPC_SetEnv(t *testing.T) { - ts := newTestServer(t) - - var envKey, envValue string - ts.srv.OnEnvSet = func(k, v string) { - envKey, envValue = k, v - } - - resp := ts.call("set_env", map[string]string{"key": "TEST_VAR", "value": "test_value"}) - - if resp.Error != nil { - t.Fatalf("unexpected error: %v", resp.Error.Message) - } - if envKey != "TEST_VAR" || envValue != "test_value" { - t.Errorf("env not set: got key=%q value=%q", envKey, envValue) - } -} - -func TestRPC_SetCWD(t *testing.T) { - ts := newTestServer(t) - resp := ts.call("set_cwd", map[string]string{"dir": "/tmp/test-cwd"}) - - if resp.Error != nil { - t.Fatalf("unexpected error: %v", resp.Error.Message) - } - var result bool - json.Unmarshal(resp.Result, &result) - if !result { - t.Error("expected result=true") - } -} - -func TestRPC_Detach_NoRunningProcess(t *testing.T) { - ts := newTestServer(t) - resp := ts.call("detach", nil) - - if resp.Error != nil { - t.Fatalf("unexpected error: %v", resp.Error.Message) - } - var result bool - json.Unmarshal(resp.Result, &result) - if result { - t.Error("expected false when no process running") - } -} - -func TestRPC_ListDetached_Empty(t *testing.T) { - ts := newTestServer(t) - resp := ts.call("list_detached", nil) - - if resp.Error != nil { - t.Fatalf("unexpected error: %v", resp.Error.Message) - } - var list []any - json.Unmarshal(resp.Result, &list) - if len(list) != 0 { - t.Errorf("expected empty list, got %d items", len(list)) - } -} - -func TestRPC_SignalDetached_NotFound(t *testing.T) { - ts := newTestServer(t) - resp := ts.call("signal_detached", map[string]any{"pid": 99999, "signal": int(syscall.SIGTERM)}) - - if resp.Error == nil { - t.Fatal("expected error for nonexistent PID") - } - if resp.Error.Code != -32000 { - t.Errorf("expected error code -32000, got %d", resp.Error.Code) - } -} - -func TestRPC_GetDetachedOutput_NotFound(t *testing.T) { - ts := newTestServer(t) - resp := ts.call("get_detached_output", map[string]any{"pid": 99999}) - - if resp.Error == nil { - t.Fatal("expected error for nonexistent PID") - } -} - -func TestRPC_DismissDetached_NotFound(t *testing.T) { - ts := newTestServer(t) - resp := ts.call("dismiss_detached", map[string]any{"pid": 99999}) - - if resp.Error != nil { - t.Fatalf("unexpected error: %v", resp.Error.Message) - } - var result bool - json.Unmarshal(resp.Result, &result) - if result { - t.Error("expected false for nonexistent PID") - } -} - -func TestRPC_CallTool_NotLoaded(t *testing.T) { - ts := newTestServer(t) - resp := ts.call("nonexistent_tool", map[string]any{"arg": "value"}) - - if resp.Error == nil { - t.Fatal("expected error for unloaded tool") - } - if resp.Error.Code != -32000 { - t.Errorf("expected error code -32000, got %d", resp.Error.Code) - } -} - -func TestRPC_InvalidParams(t *testing.T) { - ts := newTestServer(t) - // Send an array where an object is expected — json.Unmarshal will fail. - ts.enc.Encode(rpcRequest{JSONRPC: "2.0", ID: 1, Method: "set_env", Params: json.RawMessage(`[1,2,3]`)}) - for { - var resp rpcResponse - if err := ts.dec.Decode(&resp); err != nil { - t.Fatalf("decode: %v", err) - } - if resp.Stream { - continue - } - if resp.Error == nil { - t.Fatal("expected error for invalid params") - } - if resp.Error.Code != -32602 { - t.Errorf("expected error code -32602 (invalid params), got %d", resp.Error.Code) - } - break - } -} - -// --- Full pipe-based client/server tests --- - -func TestRPC_PipeConnection(t *testing.T) { - cwd := t.TempDir() - os.MkdirAll(filepath.Join(cwd, ".git"), 0755) - - reg, _ := toolsrv.NewRegistry() - srv := toolsrv.New(cwd) - srv.SetToolRegistry(reg, "pipe-test") - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - clientConn, serverConn := net.Pipe() - defer clientConn.Close() - defer serverConn.Close() - - go srv.ServeRPC(ctx, serverConn, serverConn) - - // Use toolsrv.Conn as the client - conn := toolsrv.NewConn(clientConn, nil) - - if err := conn.Ping(); err != nil { - t.Errorf("Ping failed: %v", err) - } - - tools, err := conn.ListTools() - if err != nil { - t.Errorf("ListTools failed: %v", err) - } - if len(tools) != 0 { - t.Errorf("expected 0 tools, got %d", len(tools)) - } - - conn.SetEnv("TEST_KEY", "TEST_VALUE") - conn.SetCWD("/tmp/test") -} diff --git a/toolsrv/rpcserver.go b/toolsrv/rpcserver.go deleted file mode 100644 index cb599a5..0000000 --- a/toolsrv/rpcserver.go +++ /dev/null @@ -1,263 +0,0 @@ -package toolsrv - -import ( - "context" - "encoding/json" - "fmt" - "io" - "net" - "os" - "path/filepath" - "runtime" - "sync" - "syscall" - - "ollie/paths" -) - -// RPCResponse is a JSON-RPC 2.0 response. -type RPCResponse struct { - JSONRPC string `json:"jsonrpc"` - ID json.RawMessage `json:"id"` - Result any `json:"result,omitempty"` - Error *RPCError `json:"error,omitempty"` - Stream bool `json:"stream,omitempty"` -} - -// lockedEncoder serializes JSON writes to a single writer. -type lockedEncoder struct { - mu sync.Mutex - enc *json.Encoder -} - -func (e *lockedEncoder) Encode(v any) error { - e.mu.Lock() - defer e.mu.Unlock() - return e.enc.Encode(v) -} - -// ServeSocket accepts JSON-RPC connections on a Unix socket. Blocks until -// ctx is cancelled. -func (s *Server) ServeSocket(ctx context.Context, sockPath string) error { - os.Remove(sockPath) - os.MkdirAll(filepath.Dir(sockPath), 0700) - - ln, err := net.Listen("unix", sockPath) - if err != nil { - return fmt.Errorf("listen %s: %w", sockPath, err) - } - defer ln.Close() - defer os.Remove(sockPath) - - fmt.Fprintln(os.Stderr, "ListenReady") - - var mu sync.Mutex - var activeConns []net.Conn - var wg sync.WaitGroup - - go func() { - <-ctx.Done() - ln.Close() - mu.Lock() - for _, c := range activeConns { - c.Close() - } - mu.Unlock() - }() - - for { - conn, err := ln.Accept() - if err != nil { - if ctx.Err() != nil { - break - } - continue - } - mu.Lock() - activeConns = append(activeConns, conn) - mu.Unlock() - - wg.Add(1) - go func() { - defer wg.Done() - defer conn.Close() - s.ServeRPC(ctx, conn, conn) - mu.Lock() - for i, c := range activeConns { - if c == conn { - activeConns = append(activeConns[:i], activeConns[i+1:]...) - break - } - } - mu.Unlock() - }() - } - wg.Wait() - return nil -} - -// ServeRPC reads JSON-RPC 2.0 requests from in and writes responses to out. -// Blocks until in is closed or ctx is cancelled. -func (s *Server) ServeRPC(ctx context.Context, in io.Reader, out io.Writer) { - dec := json.NewDecoder(in) - enc := &lockedEncoder{enc: json.NewEncoder(out)} - - for { - var req RPCRequest - if err := dec.Decode(&req); err != nil { - if err == io.EOF || ctx.Err() != nil { - return - } - enc.Encode(RPCResponse{ - JSONRPC: "2.0", - ID: req.ID, - Error: &RPCError{Code: -32700, Message: "parse error: " + err.Error()}, - }) - continue - } - - go s.handleRPC(ctx, req, enc) - } -} - -// CWD returns the server's current working directory. -func (s *Server) CWD() string { - s.wdMu.RLock() - defer s.wdMu.RUnlock() - return s.cwd -} - -// handleRPC dispatches a single JSON-RPC request. -// -// Reserved method names (control plane): -// - list_tools, host_info, ping, set_env, set_cwd -// - detach, list_detached, signal_detached, get_detached_output, dismiss_detached -// - tool_load -// -// Any other method name is treated as a tool invocation (data plane). -// A tool with the same name as a reserved method is unreachable. -func (s *Server) handleRPC(ctx context.Context, req RPCRequest, enc *lockedEncoder) { - switch req.Method { - case "list_tools": - tools, err := s.ListTools() - if err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}}) - } else { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: tools}) - } - - case "host_info": - info := HostInfo{ - Platform: runtime.GOOS, - Arch: runtime.GOARCH, - IsGitRepo: paths.IsGitRepo(s.CWD()), - } - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: info}) - - case "ping": - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: "pong"}) - - case "set_env": - var params struct { - Key string `json:"key"` - Value string `json:"value"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) - return - } - s.SetEnv(params.Key, params.Value) - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) - - case "set_cwd": - var params struct { - Dir string `json:"dir"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) - return - } - s.SetCWD(params.Dir) - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) - - case "detach": - ok := s.Detach() - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: ok}) - - case "list_detached": - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: s.ListDetachedRaw()}) - - case "signal_detached": - var params struct { - PID int `json:"pid"` - Signal int `json:"signal"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) - return - } - if err := s.SignalDetached(params.PID, syscall.Signal(params.Signal)); err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}}) - } else { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) - } - - case "tool_load": - var params struct { - Name string `json:"name"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) - return - } - if err := s.LoadTool(params.Name); err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}}) - } else { - res, _ := json.Marshal(map[string]string{"result": "loaded: " + params.Name}) - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(res)}) - } - - case "get_detached_output": - var params struct { - PID int `json:"pid"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) - return - } - out, err := s.GetDetachedOutput(params.PID) - if err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}}) - } else { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: out}) - } - - case "dismiss_detached": - var params struct { - PID int `json:"pid"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) - return - } - ok := s.DismissDetached(params.PID) - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: ok}) - - default: - // Data plane: method name is the tool name. - streamCtx := WithOutputStream(ctx, func(data string) { - enc.Encode(RPCResponse{ - JSONRPC: "2.0", - ID: req.ID, - Result: OutputNotification{Data: data}, - Stream: true, - }) - }) - result, err := s.CallTool(streamCtx, req.Method, req.Params) - if err != nil { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}}) - } else { - enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(result)}) - } - } -} diff --git a/toolsrv/rpcwire.go b/toolsrv/rpcwire.go deleted file mode 100644 index ea07b46..0000000 --- a/toolsrv/rpcwire.go +++ /dev/null @@ -1,51 +0,0 @@ -package toolsrv - -import "encoding/json" - -// RPCRequest is a JSON-RPC 2.0 request. -type RPCRequest struct { - JSONRPC string `json:"jsonrpc"` - ID json.RawMessage `json:"id"` - Method string `json:"method"` - Params json.RawMessage `json:"params,omitempty"` -} - -// RPCError is a JSON-RPC 2.0 error object. -type RPCError struct { - Code int `json:"code"` - Message string `json:"message"` -} - -// OutputNotification carries streamed tool output data. -type OutputNotification struct { - Data string `json:"data"` -} - -// ToolResult is the JSON envelope returned by tool execution. -// Tool scripts write this to stdout; the agent loop parses it via -// extractToolResult in agent/runtime.go. -// -// Schema: -// -// { -// "isError": false, -// "content": [ -// {"type": "text", "text": "..."}, -// {"type": "image", "media_type": "image/png", "data": ""} -// ] -// } -// -// If a tool emits raw text (not JSON), the agent treats it as a single -// text content block with isError=false. -type ToolResult struct { - IsError bool `json:"isError,omitempty"` - Content []ToolResultContent `json:"content"` -} - -// ToolResultContent is a single content block within a ToolResult. -type ToolResultContent struct { - Type string `json:"type"` - Text string `json:"text,omitempty"` - MediaType string `json:"media_type,omitempty"` - Data string `json:"data,omitempty"` -} diff --git a/toolsrv/server.go b/toolsrv/server.go deleted file mode 100644 index 2d2eeb3..0000000 --- a/toolsrv/server.go +++ /dev/null @@ -1,369 +0,0 @@ -package toolsrv - -import ( - "context" - "encoding/json" - "fmt" - "strings" - "sync" - "syscall" - "time" - - "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 - - 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 the set of loaded tools changes. - // The receiver should re-fetch the tool list to update its state. - OnToolsChanged func() - - // 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 []*DetachedProcess - 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 } } - - - -// 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. -// If OnToolsChanged is set, it is called to signal that the available tools -// have changed. The caller is responsible for re-fetching the current list. -func (s *Server) LoadTool(name string) error { - if s.toolRegistry == nil { - return fmt.Errorf("no tool registry configured") - } - if s.sessionID == "" { - return fmt.Errorf("no session ID configured") - } - if err := s.toolRegistry.Load(s.sessionID, name); err != nil { - return err - } - if s.OnToolsChanged != nil { - s.OnToolsChanged() - } - return nil -} - - - -// ListTools implements Server, returning tools loaded in the session's -// tool registry. With zero built-in tools, only dynamically loaded tools -// appear here. -func (s *Server) ListTools() ([]ToolInfo, error) { - var all []ToolInfo - - if s.toolRegistry != nil && s.sessionID != "" { - all = append(all, s.toolRegistry.Loaded(s.sessionID)...) - } - return all, nil -} - -// CallTool implements Runner. -// With zero built-in tools, all tools must be loaded in the registry first. -func (s *Server) CallTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) { - if s.toolRegistry != nil && s.sessionID != "" { - if _, promoted := s.toolRegistry.Lookup(s.sessionID, tool); promoted { - return s.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 (s *Server) SetCWD(dir string) { - s.wdMu.Lock() - s.cwd = paths.ExpandHome(dir) - s.wdMu.Unlock() -} - -// SetEnv adds a session-scoped environment variable injected into all -// subsequent subprocess invocations for this session. -func (s *Server) SetEnv(key, value string) { - s.envMu.Lock() - if s.envExtra == nil { - s.envExtra = make(map[string]string) - } - s.envExtra[key] = value - s.envMu.Unlock() - if s.OnEnvSet != nil { - s.OnEnvSet(key, value) - } -} - -// SetToolRegistry attaches a session-local tool registry. -func (s *Server) SetToolRegistry(r *Registry, sessionID string) { - s.toolRegistry = r - s.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 (s *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 dispatch-level flags (not passed to tool script). - bypassed := false - 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 (from .meta). - needsSudo := false - if m, err := LoadMetaFile(tool); err == nil && m != nil { - resolved := m.Resolve() - if resolved != nil && resolved.Sudo { - needsSudo = true - bypassed = true // sudo implies bypass - } - } - - // The tool script receives JSON args on stdin. - code := path - stdinData := string(args) - - var result string - if needsSudo { - s.wdMu.RLock() - workDir := s.cwd - s.wdMu.RUnlock() - sudoCode := fmt.Sprintf("cat <<'OLLIE_EOF' | %s\n%s\nOLLIE_EOF", code, stdinData) - result, err = s.executeBypassSudo(ctx, sudoCode, workDir, timeout) - } else if bypassed { - s.wdMu.RLock() - workDir := s.cwd - s.wdMu.RUnlock() - // Broker protocol has no stdin support; pipe JSON via heredoc. - bypassCode := fmt.Sprintf("cat <<'OLLIE_EOF' | %s\n%s\nOLLIE_EOF", code, stdinData) - result, err = s.executeBypass(ctx, bypassCode, workDir, timeout) - } else { - result, err = s.executeWithStdin(ctx, code, "bash", timeout, sandboxName, false, stdinData) - } - 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}}, - }) -} - -// Close is called when the session ends. Calls OnClose hook if registered. -func (s *Server) Close() { - s.cleanupDetached() - if s.OnClose != nil { - s.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 (s *Server) Detach() bool { - s.detachMu.Lock() - ch := s.detachCh - s.detachMu.Unlock() - if ch == nil { - return false - } - select { - case ch <- struct{}{}: - return true - default: - return false - } -} - -// 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 (s *Server) ListDetachedRaw() []any { - s.detachMu.Lock() - defer s.detachMu.Unlock() - out := make([]any, len(s.detached)) - for i, p := range s.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 (s *Server) SignalDetached(pid int, sig syscall.Signal) error { - s.detachMu.Lock() - defer s.detachMu.Unlock() - for _, p := range s.detached { - if p.PID == pid { - p.Mu.Lock() - defer p.Mu.Unlock() - if p.Exited { - return fmt.Errorf("process %d already exited", pid) - } - return syscall.Kill(-pid, sig) - } - } - return fmt.Errorf("no detached process with pid %d", pid) -} - -// GetDetachedOutput returns the ring buffer contents for a detached process. -func (s *Server) GetDetachedOutput(pid int) (string, error) { - s.detachMu.Lock() - defer s.detachMu.Unlock() - for _, p := range s.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 (s *Server) DismissDetached(pid int) bool { - s.detachMu.Lock() - defer s.detachMu.Unlock() - for i, p := range s.detached { - if p.PID == pid && p.Exited { - s.detached = append(s.detached[:i], s.detached[i+1:]...) - return true - } - } - return false -} - -// cleanupDetached sends SIGTERM to all running detached processes. -func (s *Server) cleanupDetached() { - for _, p := range s.detached { - p.Mu.Lock() - if !p.Exited { - syscall.Kill(-p.PID, syscall.SIGTERM) - } - p.Mu.Unlock() - } -} - -// SetOnToolsChanged sets the callback for tool directory changes. -func (s *Server) SetOnToolsChanged(fn func()) { - s.OnToolsChanged = fn -} - -// ToolRegistryRevision returns the current registry revision for this server's session. -// Returns 0 if no registry is configured. -func (s *Server) ToolRegistryRevision() uint64 { - if s.toolRegistry == nil || s.sessionID == "" { - return 0 - } - return s.toolRegistry.Revision(s.sessionID) -} - - diff --git a/toolsrv/shell_validate.go b/toolsrv/shell_validate.go deleted file mode 100644 index b8957d4..0000000 --- a/toolsrv/shell_validate.go +++ /dev/null @@ -1,86 +0,0 @@ -package toolsrv - -import ( - "fmt" - "regexp" - "strings" - "time" -) - -// universalPatterns apply to all code. -var universalPatterns = []*regexp.Regexp{ - regexp.MustCompile(`\bmkfs\b`), - regexp.MustCompile(`\bdd\b.*\bif=/dev/`), - regexp.MustCompile(`\b(sudo|su)\s`), - regexp.MustCompile(`/etc/(shadow|sudoers)`), -} - -// bashPatterns apply to bash (flag syntax, redirects, shell-specific constructs). -var bashPatterns = []*regexp.Regexp{ - regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s*/(home|var|usr|etc|boot|root|bin|sbin|lib|opt|srv)?`), - regexp.MustCompile(`rm\s+(-[a-z]*f[a-z]*\s+)*-[a-z]*r[a-z]*\s*/(home|var|usr|etc|boot|root|bin|sbin|lib|opt|srv)?`), - regexp.MustCompile(`rm\s+.*--recursive.*--force`), - regexp.MustCompile(`rm\s+.*--force.*--recursive`), - regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+\.\.(/|\s|$)`), - regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+~`), - regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+\*`), - regexp.MustCompile(`:\s*\(\s*\)\s*\{\s*:\s*\|\s*:\s*&`), // fork bomb - regexp.MustCompile(`>\s*/dev/sd`), - regexp.MustCompile(`\beval\s+".*\$`), -} - -var languagePatterns = map[string][]*regexp.Regexp{ - "bash": bashPatterns, - "": bashPatterns, -} - -var whitespacePattern = regexp.MustCompile(`\s+`) - -// ValidateCode checks code against dangerous patterns. -func (s *Server) ValidateCode(code, language string) error { - if err := s.checkRateLimit(); err != nil { - return err - } - - normalized := strings.ToLower(code) - normalized = whitespacePattern.ReplaceAllString(normalized, " ") - - patterns := append(universalPatterns, languagePatterns[language]...) - for _, pattern := range patterns { - if pattern.MatchString(normalized) { - s.recordValidationFailure() - return fmt.Errorf("dangerous pattern detected") - } - } - return nil -} - -func (s *Server) checkRateLimit() error { - s.rateLimitMu.Lock() - defer s.rateLimitMu.Unlock() - - now := time.Now() - if now.Before(s.blockedUntil) { - remaining := s.blockedUntil.Sub(now).Round(time.Second) - return &RateLimitedError{Remaining: remaining} - } - return nil -} - -func (s *Server) recordValidationFailure() { - s.rateLimitMu.Lock() - defer s.rateLimitMu.Unlock() - - now := time.Now() - if now.Sub(s.lastFailure) > failureWindow { - s.validationFailures = 0 - } - - s.validationFailures++ - s.lastFailure = now - - if s.validationFailures >= maxFailures { - s.blockedUntil = now.Add(blockDuration) - s.validationFailures = 0 - } -} diff --git a/toolsrv/socket_test.go b/toolsrv/socket_test.go deleted file mode 100644 index ace0fb6..0000000 --- a/toolsrv/socket_test.go +++ /dev/null @@ -1,164 +0,0 @@ -package toolsrv_test - -import ( - "context" - "net" - "os" - "path/filepath" - "testing" - "time" - - "ollie/toolsrv" -) - -// TestSocketSpawnAndDial verifies the full lifecycle: -// spawn ollie-remote → dial the socket → call RPCs → close. -func TestSocketSpawnAndDial(t *testing.T) { - // Skip if ollie-remote binary not available. - if _, err := findBinary(); err != nil { - t.Skipf("ollie-remote not found: %v", err) - } - - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - cwd, err := os.Getwd() - if err != nil { - t.Fatal(err) - } - - // Spawn the toolsrv subprocess. - proc, err := toolsrv.Spawn(ctx, cwd, toolsrv.WithYolo()) - if err != nil { - t.Fatalf("Spawn: %v", err) - } - defer proc.Close() - - if proc.Socket == "" { - t.Fatal("expected non-empty socket path") - } - - // Verify socket file exists. - if _, err := os.Stat(proc.Socket); err != nil { - t.Fatalf("socket file not found: %v", err) - } - - // Dial a connection. - conn, err := toolsrv.Dial(ctx, proc.Socket) - if err != nil { - t.Fatalf("Dial: %v", err) - } - defer conn.Close() - - // Ping. - if err := conn.Ping(); err != nil { - t.Fatalf("Ping: %v", err) - } - - // FetchHostInfo. - info, err := conn.FetchHostInfo() - if err != nil { - t.Fatalf("FetchHostInfo: %v", err) - } - if info.Platform == "" { - t.Error("expected non-empty platform") - } - - // SetEnv (fire-and-forget, just verify no crash). - conn.SetEnv("TEST_KEY", "test_value") - - // SetCWD. - conn.SetCWD("/tmp") - - // ListTools. The remote server starts with no promoted tools; tools are - // loaded dynamically through the session tool registry. - if _, err := conn.ListTools(); err != nil { - t.Fatalf("ListTools: %v", err) - } -} - -// TestSocketMultipleClients verifies multiple clients can connect simultaneously. -func TestSocketMultipleClients(t *testing.T) { - if _, err := findBinary(); err != nil { - t.Skipf("ollie-remote not found: %v", err) - } - - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - cwd, err := os.Getwd() - if err != nil { - t.Fatal(err) - } - - proc, err := toolsrv.Spawn(ctx, cwd, toolsrv.WithYolo()) - if err != nil { - t.Fatalf("Spawn: %v", err) - } - defer proc.Close() - - // Dial three clients concurrently. - conns := make([]*toolsrv.Conn, 3) - for i := range conns { - c, err := toolsrv.Dial(ctx, proc.Socket) - if err != nil { - t.Fatalf("Dial[%d]: %v", i, err) - } - conns[i] = c - } - - // Each client can ping independently. - for i, c := range conns { - if err := c.Ping(); err != nil { - t.Errorf("Ping[%d]: %v", i, err) - } - } - - // Close all. - for _, c := range conns { - c.Close() - } -} - -// TestSocketCleanup verifies the socket is removed after Close. -func TestSocketCleanup(t *testing.T) { - if _, err := findBinary(); err != nil { - t.Skipf("ollie-remote not found: %v", err) - } - - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - cwd, err := os.Getwd() - if err != nil { - t.Fatal(err) - } - - proc, err := toolsrv.Spawn(ctx, cwd, toolsrv.WithYolo()) - if err != nil { - t.Fatalf("Spawn: %v", err) - } - sockPath := proc.Socket - proc.Close() - - // Socket should be cleaned up. - time.Sleep(100 * time.Millisecond) - if _, err := net.DialTimeout("unix", sockPath, time.Second); err == nil { - t.Error("expected socket to be removed after Close") - } -} - -// findBinary checks if ollie-remote is available. -func findBinary() (string, error) { - home, _ := os.UserHomeDir() - candidates := []string{ - filepath.Join(home, "bin", "ollie-remote"), - filepath.Join(home, ".local", "bin", "ollie-remote"), - } - for _, p := range candidates { - if _, err := os.Stat(p); err == nil { - return p, nil - } - } - return "", os.ErrNotExist -} diff --git a/toolsrv/spawn.go b/toolsrv/spawn.go deleted file mode 100644 index ff2aa2a..0000000 --- a/toolsrv/spawn.go +++ /dev/null @@ -1,76 +0,0 @@ -package toolsrv - -import ( - "context" - "fmt" -) - -// Process represents a running toolsrv subprocess. -// The session owns this and calls Close() when the session ends. -type Process struct { - // Socket is the Unix socket path clients should Dial(). - Socket string - // Info holds host environment details (platform, git repo status). - Info HostInfo - // cleanup shuts down the subprocess and removes the socket. - cleanup func() -} - -// Close shuts down the toolsrv subprocess and removes the socket file. -func (p *Process) Close() { - if p.cleanup != nil { - p.cleanup() - } -} - -// Spawn starts a local ollie-remote subprocess listening on a Unix socket. -// The returned Process owns the subprocess lifecycle. Callers connect -// via Dial(LocalAddr(proc.Socket)). -func Spawn(ctx context.Context, cwd string, opts ...Option) (*Process, error) { - StreamOutput(ctx, "[toolsrv] spawning local toolsrv...\n") - - t, err := localDial(ctx, cwd, opts...) - if err != nil { - return nil, err - } - - // Verify the server is responsive via the initial connection. - conn := NewConnSplit(t.r, t.w, func() {}) - var info HostInfo - if hi, err := conn.FetchHostInfo(); err == nil { - info = hi - } - - StreamOutput(ctx, fmt.Sprintf("[toolsrv] local ready: %s (socket %s)\n", cwd, t.sockPath)) - - return &Process{ - Socket: t.sockPath, - Info: info, - cleanup: t.cleanup, - }, nil -} - -// SpawnRemote starts ollie-remote on a remote host via SSH, sets up Unix -// socket forwarding, and returns a Process with the local forwarded socket. -func SpawnRemote(ctx context.Context, cfg RemoteConfig) (*Process, error) { - StreamOutput(ctx, fmt.Sprintf("[toolsrv] dialing %s (cwd: %s)...\n", cfg.SSHTarget, cfg.CWD)) - - t, err := sshDial(ctx, cfg) - if err != nil { - return nil, fmt.Errorf("spawn remote: %w", err) - } - - conn := NewConnSplit(t.r, t.w, func() {}) - var info HostInfo - if hi, err := conn.FetchHostInfo(); err == nil { - info = hi - } - - StreamOutput(ctx, fmt.Sprintf("[toolsrv] remote ready: %s → %s (%s/%s)\n", cfg.SSHTarget, t.sockPath, info.Platform, info.Arch)) - - return &Process{ - Socket: t.sockPath, - Info: info, - cleanup: t.cleanup, - }, nil -} diff --git a/toolsrv/stream.go b/toolsrv/stream.go deleted file mode 100644 index 145fd37..0000000 --- a/toolsrv/stream.go +++ /dev/null @@ -1,30 +0,0 @@ -package toolsrv - -import "context" - -// OutputFunc is a callback for streaming partial tool output. -type OutputFunc func(data string) - -type streamKey struct{} - -// WithOutputStream returns a context carrying an output streaming callback. -// Tool servers can call StreamOutput(ctx, data) to emit partial results. -func WithOutputStream(ctx context.Context, fn OutputFunc) context.Context { - return context.WithValue(ctx, streamKey{}, fn) -} - -// StreamOutput sends partial output to the streaming callback in ctx, if any. -func StreamOutput(ctx context.Context, data string) { - if fn, ok := ctx.Value(streamKey{}).(OutputFunc); ok && fn != nil { - fn(data) - } -} - -// StreamFunc returns the streaming callback from ctx, or nil if none set. -// Useful for passing to writers that call the function directly. -func StreamFunc(ctx context.Context) func(string) { - if fn, ok := ctx.Value(streamKey{}).(OutputFunc); ok { - return func(data string) { fn(data) } - } - return nil -} diff --git a/toolsrv/transport.go b/toolsrv/transport.go deleted file mode 100644 index 2dcfd58..0000000 --- a/toolsrv/transport.go +++ /dev/null @@ -1,287 +0,0 @@ -package toolsrv - -import ( - "bufio" - "context" - "fmt" - "io" - "net" - "os" - "os/exec" - "path/filepath" - "strings" - "sync/atomic" - "syscall" - "time" - - "ollie/paths" -) - -// socketSeq is a monotonic counter to ensure unique socket paths. -var socketSeq atomic.Int64 - -// transport encapsulates the plumbing for a JSON-RPC connection to a -// toolsrv subprocess. Both local and SSH transports produce the same -// result: a reader, writer, socket path, and cleanup function. -type transport struct { - r io.Reader - w io.WriteCloser - sockPath string - cleanup func() -} - -// localDial spawns ollie-remote as a local subprocess listening on a Unix -// socket, waits for it to be ready, and returns the socket path. -func localDial(ctx context.Context, cwd string, opts ...Option) (*transport, error) { - // Bootstrap: verify ollie-remote binary is present. - binPath, err := findLocalBinary() - if err != nil { - return nil, fmt.Errorf("bootstrap: ollie-remote not found") - } - - // Verify tools directory is accessible via ToolsPath(). - toolsDir := ToolsPath() - if _, err := os.Stat(toolsDir); err != nil { - return nil, fmt.Errorf("bootstrap: tools directory not found at %s", toolsDir) - } - - // Generate a unique socket path. - sockDir := filepath.Join(paths.RuntimeDir(), "ollie") - os.MkdirAll(sockDir, 0700) - sockPath := filepath.Join(sockDir, fmt.Sprintf("toolsrv-%d-%d.sock", os.Getpid(), socketSeq.Add(1))) - - args := []string{"serve", "--cwd", cwd, "--listen", sockPath} - - probe := &Server{} - for _, o := range opts { - o(probe) - } - if probe.Yolo { - args = append(args, "--yolo") - } - - cmd := exec.CommandContext(ctx, binPath, args...) - cmd.SysProcAttr = &syscall.SysProcAttr{ - Setpgid: true, - Pdeathsig: syscall.SIGTERM, - } - - // Inherit environment. - cmd.Env = os.Environ() - - // Watch stderr for readiness signal. - stderrPipe, err := cmd.StderrPipe() - if err != nil { - return nil, fmt.Errorf("local stderr pipe: %w", err) - } - - if err := cmd.Start(); err != nil { - return nil, fmt.Errorf("local start: %w", err) - } - - // Wait for ListenReady on stderr. - readyCh := make(chan struct{}) - go func() { - scanner := bufio.NewScanner(stderrPipe) - for scanner.Scan() { - if strings.TrimSpace(scanner.Text()) == "ListenReady" { - close(readyCh) - return - } - } - }() - - select { - case <-readyCh: - case <-time.After(10 * time.Second): - syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM) - cmd.Wait() - return nil, fmt.Errorf("local toolsrv startup timeout") - } - - // Connect to the socket. - conn, err := net.Dial("unix", sockPath) - if err != nil { - syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM) - cmd.Wait() - return nil, fmt.Errorf("local socket connect: %w", err) - } - - cleanup := func() { - conn.Close() - syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM) - done := make(chan struct{}) - go func() { cmd.Wait(); close(done) }() - select { - case <-done: - case <-time.After(3 * time.Second): - syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) - <-done - } - os.Remove(sockPath) - } - - return &transport{r: conn, w: conn, sockPath: sockPath, cleanup: cleanup}, nil -} - -// sshDial connects to a remote host via SSH, bootstraps ollie-remote -// with --listen mode, sets up SSH Unix socket forwarding, and returns -// a transport connected to the forwarded local socket. -// -// Assumes pubkey auth (BatchMode=yes). No password handling. -func sshDial(ctx context.Context, cfg RemoteConfig) (*transport, error) { - // Bootstrap: verify local environment has everything needed. - cacheDir := paths.CfgDir() - if _, err := os.Stat(cacheDir); err != nil { - return nil, fmt.Errorf("bootstrap: local environment not found at %s", cacheDir) - } - - // Generate socket paths. - sockDir := filepath.Join(paths.RuntimeDir(), "ollie") - os.MkdirAll(sockDir, 0700) - localSock := filepath.Join(sockDir, fmt.Sprintf("remote-%s-%d-%d.sock", cfg.SSHTarget, os.Getpid(), socketSeq.Add(1))) - remoteSock := fmt.Sprintf("/tmp/ollie-toolsrv-%d.sock", os.Getpid()) - - // Remote bootstrap: extract tarball, set env, exec serve. - bootstrap := fmt.Sprintf(`#!/bin/sh -set -e -CACHE_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/ollie" -mkdir -p "$CACHE_DIR" -tar xzf - -C "$CACHE_DIR" -# CACHE_DIR is the extracted ollie config dir; set XDG_CONFIG_HOME so -# $XDG_CONFIG_HOME/ollie resolves to it. -export XDG_CONFIG_HOME="$(dirname "$CACHE_DIR")" -export PATH="$CACHE_DIR/bin:$PATH" -exec "$CACHE_DIR/bin/ollie-remote" serve --cwd %s --listen %s -`, shellEscape(cfg.CWD), shellEscape(remoteSock)) - - sshArgs := []string{ - "-o", "BatchMode=yes", - "-o", "StrictHostKeyChecking=accept-new", - "-o", "StreamLocalBindUnlink=yes", - "-L", localSock + ":" + remoteSock, - } - sshTarget := cfg.SSHTarget - if host, port, ok := strings.Cut(sshTarget, ":"); ok && port != "" { - sshTarget = host - sshArgs = append(sshArgs, "-p", port) - } - sshArgs = append(sshArgs, sshTarget, "bash -s") - - cmd := exec.CommandContext(ctx, "ssh", sshArgs...) - cmd.SysProcAttr = &syscall.SysProcAttr{ - Setpgid: true, - Pdeathsig: syscall.SIGTERM, - } - - stdin, err := cmd.StdinPipe() - if err != nil { - return nil, fmt.Errorf("ssh stdin pipe: %w", err) - } - stderrPipe, err := cmd.StderrPipe() - if err != nil { - return nil, fmt.Errorf("ssh stderr pipe: %w", err) - } - - if err := cmd.Start(); err != nil { - return nil, fmt.Errorf("ssh start: %w", err) - } - - // Send tarball over stdin, then the bootstrap script. - if err := writeTarball(stdin, cacheDir); err != nil { - cmd.Process.Kill() - return nil, fmt.Errorf("send tarball: %w", err) - } - if _, err := stdin.Write([]byte(bootstrap + "\n")); err != nil { - cmd.Process.Kill() - return nil, fmt.Errorf("send bootstrap: %w", err) - } - - // Monitor stderr for ListenReady. - readyCh := make(chan struct{}, 1) - go func() { - scanner := bufio.NewScanner(stderrPipe) - for scanner.Scan() { - if strings.TrimSpace(scanner.Text()) == "ListenReady" { - close(readyCh) - return - } - } - }() - - select { - case <-readyCh: - case <-time.After(30 * time.Second): - cmd.Process.Kill() - return nil, fmt.Errorf("bootstrap timeout waiting for listen ready") - } - - // Poll for local socket forwarding. - var conn net.Conn - for i := 0; i < 20; i++ { - conn, err = net.Dial("unix", localSock) - if err == nil { - break - } - time.Sleep(50 * time.Millisecond) - } - if err != nil { - cmd.Process.Kill() - return nil, fmt.Errorf("ssh socket connect: %w", err) - } - - cleanup := func() { - conn.Close() - stdin.Close() - syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM) - done := make(chan struct{}) - go func() { cmd.Wait(); close(done) }() - select { - case <-done: - case <-time.After(3 * time.Second): - syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) - <-done - } - os.Remove(localSock) - } - - return &transport{r: conn, w: conn, sockPath: localSock, cleanup: cleanup}, nil -} - -// writeTarball creates a gzipped tarball of the directory at dir and writes -// it to w. Used to transfer the ollie environment to remote hosts. -func writeTarball(w io.Writer, dir string) error { - cmd := exec.Command("tar", "czf", "-", "-C", dir, ".") - cmd.Stdout = w - cmd.Stderr = os.Stderr - return cmd.Run() -} - -// findLocalBinary locates the ollie-remote binary on the local system. -func findLocalBinary() (string, error) { - home, _ := os.UserHomeDir() - candidates := []string{ - filepath.Join(home, "bin", "ollie-remote"), - filepath.Join(home, ".local", "bin", "ollie-remote"), - } - for _, p := range candidates { - if _, err := os.Stat(p); err == nil { - return p, nil - } - } - if p, err := exec.LookPath("ollie-remote"); err == nil { - return p, nil - } - return "", fmt.Errorf("ollie-remote not found (checked %s)", strings.Join(candidates, ", ")) -} - -// shellEscape escapes a string for safe use in a shell command. -func shellEscape(s string) string { - if s == "" { - return "''" - } - if !strings.ContainsAny(s, " \t\n\r'\"\\$`!#&|;(){}") { - return s - } - return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'" -} diff --git a/toolsrv/types.go b/toolsrv/types.go new file mode 100644 index 0000000..fddf7fa --- /dev/null +++ b/toolsrv/types.go @@ -0,0 +1,16 @@ +// types.go - Shared types for toolsrv +package toolsrv + +// ToolResult is the structured result from a tool execution. +type ToolResult struct { + Content []ToolResultContent `json:"content"` + IsError bool `json:"isError,omitempty"` +} + +// ToolResultContent is a single content item in a tool result. +type ToolResultContent struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + MediaType string `json:"mediaType,omitempty"` + Data string `json:"data,omitempty"` +}