diff --git a/cmd/ollie-remote/main.go b/cmd/ollie-remote/main.go index 304e57b..a92d3c5 100644 --- a/cmd/ollie-remote/main.go +++ b/cmd/ollie-remote/main.go @@ -1,12 +1,16 @@ -// ollie-remote — minimal remote execution server for split-brain mode. +// ollie-remote — tool execution server. // -// Accepts JSON-RPC 2.0 over stdin/stdout. Executes tool calls in a sandboxed -// environment on the remote host. Deployed automatically by olliesrv/ollied -// via SSH bootstrap. +// 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] // ollie-remote serve --cwd /path/to/project [--tools /path/to/tools] [--yolo] +// +// When --listen is provided, accepts connections on a Unix socket. +// Otherwise, serves a single session over stdin/stdout (bootstrap mode). package main import ( @@ -17,13 +21,17 @@ import ( "fmt" "io" "io/fs" + "net" "os" "os/signal" "path/filepath" "runtime" + "sync" "syscall" + "ollie/skills" "ollie/toolsrv" + "ollie/tools/builtin" ) //go:embed sandbox-remote.yaml @@ -33,15 +41,16 @@ var defaultSandboxConfig []byte var embeddedTools embed.FS var ( - cwd = flag.String("cwd", ".", "working directory for execution") - toolsPath = flag.String("tools", "", "path to tool scripts (default: ~/.cache/ollie/tools)") - yolo = flag.Bool("yolo", false, "skip sandbox enforcement") + cwd = flag.String("cwd", ".", "working directory for execution") + toolsPath = flag.String("tools", "", "path to tool scripts (default: ~/.cache/ollie/tools)") + listenPath = flag.String("listen", "", "Unix socket path to listen on (default: stdio mode)") + yolo = flag.Bool("yolo", false, "skip sandbox enforcement") ) func main() { // "serve" subcommand must come first; flags follow it. if len(os.Args) < 2 || os.Args[1] != "serve" { - fmt.Fprintln(os.Stderr, "usage: ollie-remote serve --cwd ") + fmt.Fprintln(os.Stderr, "usage: ollie-remote serve --cwd [--listen ]") os.Exit(1) } // Parse flags after the subcommand @@ -71,14 +80,25 @@ func main() { // otherwise use the embedded default. setupSandboxConfig(*cwd) - // Build the tool server + // Build the tool server with builtins and registries. + toolReg, _ := toolsrv.NewRegistry() + skillsReg := skills.NewRegistry() + agentID := fmt.Sprintf("agent-%d", os.Getpid()) + var opts []toolsrv.Option if *yolo { opts = append(opts, toolsrv.WithYolo()) } opts = append(opts, toolsrv.WithStrict()) + opts = append(opts, toolsrv.WithBuiltins(builtin.Builtins())) + if toolReg != nil { + opts = append(opts, toolsrv.WithToolRegistry(toolReg, agentID)) + } + opts = append(opts, toolsrv.WithSkillsRegistry(skillsReg)) server := toolsrv.New(*cwd) - for _, o := range opts { o(server) } + for _, o := range opts { + o(server) + } // Signal handling ctx, cancel := context.WithCancel(context.Background()) @@ -90,8 +110,77 @@ func main() { cancel() }() - // Serve JSON-RPC over stdin/stdout - serveRPC(ctx, server, os.Stdin, os.Stdout) + if *listenPath != "" { + // Socket mode: accept multiple connections. + serveSocket(ctx, server, *listenPath) + } else { + // Stdio mode: single session over stdin/stdout (bootstrap/legacy). + serveRPC(ctx, server, os.Stdin, os.Stdout) + } +} + +// serveSocket listens on a Unix socket and handles multiple concurrent clients. +func serveSocket(ctx context.Context, srv *toolsrv.Server, sockPath string) { + // Remove stale socket. + os.Remove(sockPath) + os.MkdirAll(filepath.Dir(sockPath), 0700) + + ln, err := net.Listen("unix", sockPath) + if err != nil { + fmt.Fprintf(os.Stderr, "listen %s: %v\n", sockPath, err) + os.Exit(1) + } + defer ln.Close() + defer os.Remove(sockPath) + + // Signal readiness on stderr (transport layer watches for this). + fmt.Fprintln(os.Stderr, "OLLIE_LISTEN_READY") + + // Track active connections so we can close them on shutdown. + var mu sync.Mutex + var activeConns []net.Conn + var wg sync.WaitGroup + + go func() { + <-ctx.Done() + ln.Close() + // Close all active connections to unblock their reads. + 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() + serveRPC(ctx, srv, conn, conn) + // Remove from active list. + mu.Lock() + for i, c := range activeConns { + if c == conn { + activeConns = append(activeConns[:i], activeConns[i+1:]...) + break + } + } + mu.Unlock() + }() + } + wg.Wait() } // --- JSON-RPC 2.0 --- @@ -115,13 +204,6 @@ type rpcError struct { Message string `json:"message"` } -type shellParams struct { - Command string `json:"command"` - Language string `json:"language,omitempty"` - Timeout int `json:"timeout,omitempty"` - Stdin string `json:"stdin,omitempty"` -} - type outputNotif struct { Data string `json:"data"` } @@ -145,43 +227,6 @@ func serveRPC(ctx context.Context, srv *toolsrv.Server, in io.Reader, out io.Wri } switch req.Method { - case "shell": - var params shellParams - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Error: &rpcError{Code: -32602, Message: "invalid params: " + err.Error()}, - }) - continue - } - // Wire streaming output to JSON-RPC notifications. - streamCtx := toolsrv.WithOutputStream(ctx, func(data string) { - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - Result: outputNotif{Data: data}, - }) - }) - result, err := srv.ExecuteInSandbox(streamCtx, params.Command, params.Language, params.Timeout, "default", false, params.Stdin, false) - if err != nil { - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Result: map[string]any{ - "isError": true, - "content": []map[string]string{{"type": "text", "text": err.Error()}}, - }, - }) - } else { - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Result: map[string]any{ - "content": []map[string]string{{"type": "text", "text": result}}, - }, - }) - } - case "list_tools": tools, err := srv.ListTools() if err != nil { @@ -217,6 +262,122 @@ func serveRPC(ctx context.Context, srv *toolsrv.Server, in io.Reader, out io.Wri 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()}}) + continue + } + srv.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()}}) + continue + } + srv.SetCWD(params.Dir) + enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) + + case "set_allow_tools": + var params struct { + Names []string `json:"names"` + } + 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()}}) + continue + } + srv.SetAllowTools(params.Names) + enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) + + case "detach": + ok := srv.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: srv.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()}}) + continue + } + if err := srv.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 "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()}}) + continue + } + out, err := srv.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()}}) + continue + } + ok := srv.DismissDetached(params.PID) + enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: ok}) + + case "is_parallel_read": + 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()}}) + continue + } + ok := srv.IsParallelRead(params.Name) + enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: ok}) + + case "result_tier": + 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()}}) + continue + } + tier := srv.ResultTier(params.Name) + enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: tier}) + + case "result_tier_args": + var params struct { + Name string `json:"name"` + Args json.RawMessage `json:"args"` + } + 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()}}) + continue + } + tier := srv.ResultTierArgs(params.Name, params.Args) + enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: tier}) + default: // Dispatch to named tool via the tool server's CallTool. streamCtx := toolsrv.WithOutputStream(ctx, func(data string) { diff --git a/cmd/ollie-remote/sandbox-remote.yaml b/cmd/ollie-remote/sandbox-remote.yaml index fcc16a1..49b2f51 100644 --- a/cmd/ollie-remote/sandbox-remote.yaml +++ b/cmd/ollie-remote/sandbox-remote.yaml @@ -20,6 +20,7 @@ filesystem: - "/lib64" - "/bin" - "/sbin" + - "{HOME}/bin" - "{HOME}/go/bin" - "{HOME}/.cargo/bin" - "{HOME}/.local/bin" diff --git a/fs/session/create.go b/fs/session/create.go index 606cc16..8d1e56d 100644 --- a/fs/session/create.go +++ b/fs/session/create.go @@ -10,7 +10,6 @@ import ( "ollie/backend" coresession "ollie/session" "ollie/toolsrv" - "ollie/tools/builtin" "ollie/paths" "ollie/prompts" ) @@ -85,6 +84,7 @@ func Create(rs *rootState, args []string) (string, error) { var core *coresession.Session var sessPtr *Session + var proc *toolsrv.Process uname := rs.nextUname() if rs.cfg.NewCore != nil { var err error @@ -119,44 +119,57 @@ func Create(rs *rootState, args []string) (string, error) { return "", fmt.Errorf("sessions dir: %w", err) } - var execOpts []toolsrv.Option + // Build dial options — only strict/yolo apply to the subprocess. + var dialOpts []toolsrv.Option if rs.cfg.Strict { - execOpts = append(execOpts, toolsrv.WithStrict()) + dialOpts = append(dialOpts, toolsrv.WithStrict()) } if rs.cfg.Yolo { - execOpts = append(execOpts, toolsrv.WithYolo()) + dialOpts = append(dialOpts, toolsrv.WithYolo()) } - if cfg != nil && len(cfg.AllowTools) > 0 { - execOpts = append(execOpts, toolsrv.WithAllowTools(cfg.AllowTools)) - } - if rs.cfg.ToolRegistry != nil { - execOpts = append(execOpts, toolsrv.WithToolRegistry(rs.cfg.ToolRegistry, sessID)) - } - if rs.cfg.SkillsRegistry != nil { - execOpts = append(execOpts, toolsrv.WithSkillsRegistry(rs.cfg.SkillsRegistry)) - execOpts = append(execOpts, toolsrv.WithBuiltins(builtin.Builtins())) - } var newToolServer func() toolsrv.Runner var remoteEnv []string if remoteTarget != "" { - rsrv, dialErr := toolsrv.RemoteDial(rs.cfg.Ctx, toolsrv.RemoteConfig{ + var spawnErr error + proc, spawnErr = toolsrv.SpawnRemote(rs.cfg.Ctx, toolsrv.RemoteConfig{ SSHTarget: remoteTarget, CWD: cwd, }) - if dialErr != nil { - return "", fmt.Errorf("remote dial: %w", dialErr) + if spawnErr != nil { + return "", fmt.Errorf("remote spawn: %w", spawnErr) + } + if cfg != nil && len(cfg.AllowTools) > 0 { + // Apply allow-tools on first connection. + if c, err := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket)); err == nil { + c.SetAllowTools(cfg.AllowTools) + c.Close() + } + } + newToolServer = func() toolsrv.Runner { + conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket)) + return conn } - newToolServer = func() toolsrv.Runner { return rsrv } remoteEnv = []string{ "PRIME_CWD=" + cwd, - "PRIME_PLATFORM=" + rsrv.Info.Platform, - "PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", rsrv.Info.IsGitRepo), + "PRIME_PLATFORM=" + proc.Info.Platform, + "PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", proc.Info.IsGitRepo), } } else { + var spawnErr error + proc, spawnErr = toolsrv.Spawn(rs.cfg.Ctx, cwd, dialOpts...) + if spawnErr != nil { + return "", fmt.Errorf("local spawn: %w", spawnErr) + } + if cfg != nil && len(cfg.AllowTools) > 0 { + if c, err := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket)); err == nil { + c.SetAllowTools(cfg.AllowTools) + c.Close() + } + } newToolServer = func() toolsrv.Runner { - srv, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(cwd), execOpts...) - return srv + conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket)) + return conn } } @@ -231,6 +244,7 @@ func Create(rs *rootState, args []string) (string, error) { sess := NewSession(sessID, core, sessionCtx, sessionCancel) sessPtr = sess sess.remote = remoteTarget + sess.proc = proc rs.mu.Lock() sess.uname = uname diff --git a/fs/session/persist.go b/fs/session/persist.go index a3397f8..f7575b2 100644 --- a/fs/session/persist.go +++ b/fs/session/persist.go @@ -13,7 +13,6 @@ import ( "ollie/backend" coresession "ollie/session" "ollie/toolsrv" - "ollie/tools/builtin" "ollie/prompts" ) @@ -26,6 +25,9 @@ func Kill(rs *rootState, id string) { if sess != nil { sess.Cancel() sess.Core.Close() + if sess.proc != nil { + sess.proc.Close() + } removePersistedAgent(rs, id) rs.cfg.Log.Info("killed session %s", id) if rs.cfg.OnSessionKilled != nil { @@ -239,40 +241,43 @@ func restoreSession(rs *rootState, ps *agent.PersistedAgent) error { uname := rs.nextUname() var newToolServer func() toolsrv.Runner var promptEnv []string + var proc *toolsrv.Process remoteTarget := ps.Remote if remoteTarget != "" { - rsrv, dialErr := toolsrv.RemoteDial(rs.cfg.Ctx, toolsrv.RemoteConfig{ + var spawnErr error + proc, spawnErr = toolsrv.SpawnRemote(rs.cfg.Ctx, toolsrv.RemoteConfig{ SSHTarget: remoteTarget, CWD: cwd, }) - if dialErr != nil { - return fmt.Errorf("remote dial: %w", dialErr) - } - newToolServer = func() toolsrv.Runner { return rsrv } - promptEnv = []string{ - "PRIME_CWD=" + cwd, - "PRIME_PLATFORM=" + rsrv.Info.Platform, - "PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", rsrv.Info.IsGitRepo), - } - } else { - var execOpts []toolsrv.Option - if rs.cfg.Strict { - execOpts = append(execOpts, toolsrv.WithStrict()) - } - if rs.cfg.Yolo { - execOpts = append(execOpts, toolsrv.WithYolo()) - } - if rs.cfg.ToolRegistry != nil { - execOpts = append(execOpts, toolsrv.WithToolRegistry(rs.cfg.ToolRegistry, sessID)) - } - if rs.cfg.SkillsRegistry != nil { - execOpts = append(execOpts, toolsrv.WithSkillsRegistry(rs.cfg.SkillsRegistry)) - execOpts = append(execOpts, toolsrv.WithBuiltins(builtin.Builtins())) + if spawnErr != nil { + return fmt.Errorf("remote spawn: %w", spawnErr) } newToolServer = func() toolsrv.Runner { - srv, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(cwd), execOpts...) - return srv + conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket)) + return conn + } + promptEnv = []string{ + "PRIME_CWD=" + cwd, + "PRIME_PLATFORM=" + proc.Info.Platform, + "PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", proc.Info.IsGitRepo), + } + } else { + var dialOpts []toolsrv.Option + if rs.cfg.Strict { + dialOpts = append(dialOpts, toolsrv.WithStrict()) + } + if rs.cfg.Yolo { + dialOpts = append(dialOpts, toolsrv.WithYolo()) + } + var spawnErr error + proc, spawnErr = toolsrv.Spawn(rs.cfg.Ctx, cwd, dialOpts...) + if spawnErr != nil { + return fmt.Errorf("local spawn: %w", spawnErr) + } + newToolServer = func() toolsrv.Runner { + conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket)) + return conn } promptEnv = agent.PromptEnv(cwd) } @@ -344,6 +349,7 @@ func restoreSession(rs *rootState, ps *agent.PersistedAgent) error { sessPtr = sess sess.uname = uname sess.remote = remoteTarget + sess.proc = proc replayMessagesToLog(sess, ps.Messages) diff --git a/fs/session/types.go b/fs/session/types.go index 6c18fc2..94039a4 100644 --- a/fs/session/types.go +++ b/fs/session/types.go @@ -8,6 +8,7 @@ import ( "ollie/agent" "ollie/session" + "ollie/toolsrv" ) // Session holds all state for one agent session exposed via 9P. @@ -24,6 +25,7 @@ type Session struct { plan []byte prevPrompt []byte // last submitted prompt; overwritten on each new submission remote string // SSH target for remote execution (empty = local) + proc *toolsrv.Process // toolsrv subprocess (session owns lifecycle) // Cached model list (expensive API call; refreshed every 24h). modelsMu sync.Mutex diff --git a/justfile b/justfile index f0715db..ce363f3 100644 --- a/justfile +++ b/justfile @@ -20,9 +20,9 @@ kf5: build-kf5 install-kf5 # === Build targets === -build: core ninep acme lsp-tools kde +build: core ninep ollie-remote acme lsp-tools kde -build-kf5: core ninep acme kde-kf5 +build-kf5: core ninep ollie-remote acme kde-kf5 core: @echo === Building ollie-core === @@ -262,4 +262,4 @@ help: @echo ' just uninstall-kf5 Remove all installed files (KF5)' @echo ' just clean Clean build artifacts' @echo '' - @echo 'Individual: just core | ninep | acme | kde | kde-kf5 | ollie-remote' + @echo 'Individual: just core | ninep | ollie-remote | acme | kde | kde-kf5 | lsp-tools' diff --git a/ollie-remote b/ollie-remote new file mode 100755 index 0000000..e6b4538 Binary files /dev/null and b/ollie-remote differ diff --git a/toolsrv/conn.go b/toolsrv/conn.go new file mode 100644 index 0000000..5879c9f --- /dev/null +++ b/toolsrv/conn.go @@ -0,0 +1,247 @@ +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 { + mu sync.Mutex + 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(string) +} + +// 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 { + return &Conn{ + enc: json.NewEncoder(rwc), + dec: json.NewDecoder(bufio.NewReader(rwc)), + cleanup: cleanup, + } +} + +// NewConnSplit creates a Conn from separate reader/writer (e.g. stdin/stdout pipes). +func NewConnSplit(r io.Reader, w io.WriteCloser, cleanup func()) *Conn { + return &Conn{ + enc: json.NewEncoder(w), + dec: json.NewDecoder(bufio.NewReader(r)), + cleanup: cleanup, + } +} + +// --- 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) SetAllowTools(names []string) { + params, _ := json.Marshal(map[string]any{"names": names}) + c.call("set_allow_tools", params) +} + +func (c *Conn) SetOnToolsChanged(fn func(string)) { + c.mu.Lock() + c.onToolsChanged = fn + c.mu.Unlock() +} + +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 +} + +// IsParallelRead implements ParallelClassifier. +func (c *Conn) IsParallelRead(name string) bool { + params, _ := json.Marshal(map[string]string{"name": name}) + resp, err := c.call("is_parallel_read", params) + if err != nil { + return false + } + var ok bool + json.Unmarshal(resp, &ok) + return ok +} + +// ResultTier implements TierClassifier. +func (c *Conn) ResultTier(name string) string { + params, _ := json.Marshal(map[string]string{"name": name}) + resp, err := c.call("result_tier", params) + if err != nil { + return "hot" + } + var tier string + json.Unmarshal(resp, &tier) + return tier +} + +// ResultTierArgs implements TierClassifier. +func (c *Conn) ResultTierArgs(name string, args json.RawMessage) string { + params, _ := json.Marshal(map[string]any{"name": name, "args": args}) + resp, err := c.call("result_tier_args", params) + if err != nil { + return "hot" + } + var tier string + json.Unmarshal(resp, &tier) + return tier +} + +// Close shuts down the connection and subprocess. +func (c *Conn) Close() { + 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 +} + +// --- 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) + req := rpcRequest{ + JSONRPC: "2.0", + ID: id, + Method: method, + Params: params, + } + + c.mu.Lock() + if err := c.enc.Encode(req); err != nil { + c.mu.Unlock() + return nil, fmt.Errorf("rpc write %s: %w", method, err) + } + + // Read responses, handling interleaved streaming notifications. + for { + var resp rpcResponse + if err := c.dec.Decode(&resp); err != nil { + c.mu.Unlock() + return nil, fmt.Errorf("rpc read %s: %w", method, err) + } + + // Notification: id == 0 means streaming output or event. + if resp.ID == 0 && resp.Result != nil { + var notif outputNotification + if json.Unmarshal(resp.Result, ¬if) == nil && notif.Data != "" { + StreamOutput(ctx, notif.Data) + } + continue + } + + // Actual response for our request. + c.mu.Unlock() + if resp.Error != nil { + return nil, fmt.Errorf("rpc %s: %s", method, resp.Error.Message) + } + return resp.Result, nil + } +} diff --git a/toolsrv/dial.go b/toolsrv/dial.go index 8174e4c..f66afc4 100644 --- a/toolsrv/dial.go +++ b/toolsrv/dial.go @@ -3,21 +3,25 @@ package toolsrv import ( "context" "fmt" + "net" ) // Addr describes how to reach a tool server. type Addr struct { - // Network is the transport: "local", "ssh", or (future) "unix". + // Network is the transport: "unix" (pre-existing socket) or "ssh". Network string // Host is the SSH target (user@host or user@host:port). Empty for local. Host string // CWD is the working directory on the target. CWD string + // Socket is the Unix socket path to connect to. + // For local: set by Spawn(). For SSH: set by SpawnRemote(). + Socket string } -// LocalAddr returns an Addr for an in-process tool server. -func LocalAddr(cwd string) Addr { - return Addr{Network: "local", CWD: cwd} +// LocalAddr returns an Addr for connecting to a local socket. +func LocalAddr(socket string) Addr { + return Addr{Network: "unix", Socket: socket} } // SSHAddr returns an Addr for a remote tool server over SSH. @@ -25,30 +29,54 @@ func SSHAddr(host, cwd string) Addr { return Addr{Network: "ssh", Host: host, CWD: cwd} } -// Dial connects to a tool server at the given address. -// Options (WithStrict, WithBuiltins, etc.) are applied to local servers. -// -// TODO: Unify local and remote behind a single protocol. Local execution -// should eventually go through a Unix socket (same JSON-RPC as remote), -// making the agent always a client. The elevation broker already proves -// this pattern. Benefits: process isolation, shared tool server across -// sessions, independent lifecycle. Transport becomes the only difference: -// - Network "local" → unix socket to co-located toolsrv process -// - Network "ssh" → SSH tunnel to remote toolsrv process -// - Network "unix" → pre-existing unix socket -func Dial(ctx context.Context, addr Addr, opts ...Option) (Runner, error) { +// Dial connects to a running tool server at the given address. +// The server must already be running (started via Spawn or SpawnRemote). +func Dial(ctx context.Context, addr Addr) (*Conn, error) { switch addr.Network { - case "", "local": - srv := New(addr.CWD) - for _, o := range opts { - o(srv) + case "", "unix": + if addr.Socket == "" { + return nil, fmt.Errorf("toolsrv.Dial: socket path required for unix network") } - return srv, nil + conn, err := net.Dial("unix", addr.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 + case "ssh": - return RemoteDial(ctx, RemoteConfig{ + // For SSH, the caller should use SpawnRemote first, then Dial the + // returned local socket. Direct SSH dial is a convenience that does both. + proc, err := SpawnRemote(ctx, RemoteConfig{ SSHTarget: addr.Host, CWD: addr.CWD, }) + if err != nil { + return nil, err + } + conn, err := net.Dial("unix", proc.Socket) + if err != nil { + proc.Close() + return nil, fmt.Errorf("toolsrv.Dial ssh: %w", err) + } + // Wrap cleanup: closing the conn + the process. + c := NewConn(conn, func() { + conn.Close() + proc.Close() + }) + if err := c.Ping(); err != nil { + c.Close() + return nil, fmt.Errorf("toolsrv.Dial ssh: ping failed: %w", err) + } + if info, err := c.FetchHostInfo(); err == nil { + c.Info = info + } + return c, nil + default: return nil, fmt.Errorf("toolsrv.Dial: unsupported network %q", addr.Network) } diff --git a/toolsrv/remote.go b/toolsrv/remote.go index 49bac0f..d79595c 100644 --- a/toolsrv/remote.go +++ b/toolsrv/remote.go @@ -1,35 +1,14 @@ -// Package remote provides the SSH bootstrap and JSON-RPC client for -// split-brain remote execution. It connects to a remote host over SSH, -// ensures ollie-remote is deployed, and returns a Server that -// forwards execution calls over the RPC channel. +// 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 ( - "bufio" - "bytes" - "compress/gzip" "context" - "crypto/sha256" - _ "embed" - "encoding/base64" - "encoding/hex" - "encoding/json" - "fmt" - "io" - "os" - "os/exec" - "path/filepath" - "strings" - "sync" - "sync/atomic" - "syscall" - "time" - ) -//go:embed bootstrap.sh -var bootstrapTemplate string - // HostInfo holds environment details from the remote host. type HostInfo struct { Platform string `json:"platform"` @@ -37,25 +16,7 @@ type HostInfo struct { IsGitRepo bool `json:"is_git_repo"` } -// Server implements Server by forwarding calls to a remote -// ollie-remote process over SSH. -type RemoteServer struct { - mu sync.Mutex - stdin io.WriteCloser - stdout io.ReadCloser - cmd *exec.Cmd - enc *json.Encoder - dec *json.Decoder - nextID atomic.Int64 - Info HostInfo - - // For auto-upgrade on hash mismatch. - cfg RemoteConfig - ctx context.Context - deployedHash string -} - -// Config holds the parameters for connecting to a remote host. +// 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 @@ -67,439 +28,33 @@ type RemoteConfig struct { Yolo bool } -// Dial opens an SSH connection to the remote host, bootstraps ollie-remote -// (deploying the binary if needed), and returns a Server ready for tool calls. -func RemoteDial(ctx context.Context, cfg RemoteConfig) (*RemoteServer, error) { - // Find the local ollie-remote binary to compute hash and transfer if needed. - localBin, err := findLocalBinary() +// RemoteDial opens an SSH connection to the remote host and returns a *Conn. +// Convenience wrapper: spawns the remote process and dials a single connection. +// For session-level use, prefer SpawnRemote() + Dial() separately. +func RemoteDial(ctx context.Context, cfg RemoteConfig) (*Conn, error) { + proc, err := SpawnRemote(ctx, cfg) if err != nil { - return nil, fmt.Errorf("find ollie-remote binary: %w", err) + return nil, err } - binData, err := os.ReadFile(localBin) + conn, err := Dial(ctx, LocalAddr(proc.Socket)) if err != nil { - return nil, fmt.Errorf("read ollie-remote binary: %w", err) + proc.Close() + return nil, err } - hash := sha256.Sum256(binData) - hashStr := hex.EncodeToString(hash[:]) - - // Find landrun binary to deploy alongside ollie-remote. - landrunPath, err := exec.LookPath("landrun") - if err != nil { - return nil, fmt.Errorf("find landrun binary: %w", err) - } - landrunData, err := os.ReadFile(landrunPath) - if err != nil { - return nil, fmt.Errorf("read landrun binary: %w", err) - } - landrunHash := sha256.Sum256(landrunData) - landrunHashStr := hex.EncodeToString(landrunHash[:]) - - // Generate bootstrap script with hashes and CWD baked in. - bootstrap := bootstrapTemplate - bootstrap = strings.ReplaceAll(bootstrap, "@@HASH@@", hashStr) - bootstrap = strings.ReplaceAll(bootstrap, "@@LANDRUN_HASH@@", landrunHashStr) - cwdEscaped := shellEscape(cfg.CWD) - extraArgs := "" - if cfg.ToolsPath != "" { - extraArgs += " --tools " + shellEscape(cfg.ToolsPath) - } - if cfg.Yolo { - extraArgs += " --yolo" - } - bootstrap = strings.ReplaceAll(bootstrap, "@@CWD@@", cwdEscaped+extraArgs) - - // Parse optional port from SSHTarget (user@host:port or host:port) - sshArgs := []string{ - "-o", "BatchMode=yes", - "-o", "StrictHostKeyChecking=accept-new", - } - sshTarget := cfg.SSHTarget - if host, port, ok := strings.Cut(sshTarget, ":"); ok && port != "" { - sshTarget = host - sshArgs = append(sshArgs, "-p", port) - } - // Run bash on the remote; we'll send the bootstrap script over stdin. - sshArgs = append(sshArgs, sshTarget, "bash -s") - - cmd := exec.CommandContext(ctx, "ssh", sshArgs...) - cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} - - stdin, err := cmd.StdinPipe() - if err != nil { - return nil, fmt.Errorf("ssh stdin pipe: %w", err) - } - stdout, err := cmd.StdoutPipe() - if err != nil { - return nil, fmt.Errorf("ssh stdout 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 bootstrap script over stdin. - if _, err := stdin.Write([]byte(bootstrap + "\n")); err != nil { - cmd.Process.Kill() - return nil, fmt.Errorf("send bootstrap: %w", err) - } - - // Monitor stderr for bootstrap messages. - type loaderMsg struct { - NeedDownload bool `json:"need_download"` - NeedLandrun bool `json:"need_landrun"` - } - - stderrScanner := bufio.NewScanner(stderrPipe) - loaderCh := make(chan loaderMsg, 1) - readyCh := make(chan struct{}, 1) - go func() { - for stderrScanner.Scan() { - line := strings.TrimSpace(stderrScanner.Text()) - if after, ok := strings.CutPrefix(line, "OLLIE_LOADER_START "); ok { - var msg loaderMsg - json.Unmarshal([]byte(after), &msg) - loaderCh <- msg - if !msg.NeedDownload && !msg.NeedLandrun { - close(readyCh) - } - } else if line == "OLLIE_LOADER_READY" { - close(readyCh) - } + // Wrap: closing this conn also closes the process. + origClose := conn.cleanup + conn.cleanup = func() { + if origClose != nil { + origClose() } - }() - - // Wait for loader start message. - select { - case msg := <-loaderCh: - if msg.NeedDownload { - if err := transferBinary(stdin, binData); err != nil { - cmd.Process.Kill() - return nil, fmt.Errorf("transfer ollie-remote: %w", err) - } - } - if msg.NeedLandrun { - if err := transferBinary(stdin, landrunData); err != nil { - cmd.Process.Kill() - return nil, fmt.Errorf("transfer landrun: %w", err) - } - } - if msg.NeedDownload || msg.NeedLandrun { - // Wait for ready signal. - select { - case <-readyCh: - case <-time.After(30 * time.Second): - cmd.Process.Kill() - return nil, fmt.Errorf("bootstrap timeout waiting for ready") - } - } - case <-time.After(30 * time.Second): - cmd.Process.Kill() - return nil, fmt.Errorf("bootstrap timeout waiting for loader start") + proc.Close() } - - // Bootstrap complete — stdin/stdout are now JSON-RPC. - s := &RemoteServer{ - stdin: stdin, - stdout: stdout, - cmd: cmd, - enc: json.NewEncoder(stdin), - dec: json.NewDecoder(bufio.NewReader(stdout)), - cfg: cfg, - ctx: ctx, - deployedHash: hashStr, - } - - // Verify connectivity with a ping. - if err := s.ping(ctx); err != nil { - s.Close() - return nil, fmt.Errorf("remote ping failed: %w", err) - } - - // Fetch remote host info for prompt resolution. - // FIXME: This duplicates information that prime scripts also try to detect - // locally (platform, is_git_repo). Eventually unify so prompt resolution - // always uses these env vars rather than running local checks against a - // non-existent remote CWD. - if info, err := s.fetchHostInfo(ctx); err == nil { - s.Info = info - } - - return s, nil + conn.Info = proc.Info + return conn, nil } -// findLocalBinary locates the ollie-remote binary on the local system. -func findLocalBinary() (string, error) { - // Check common locations in priority order. - 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 - } - } - // Try PATH as last resort. - if p, err := exec.LookPath("ollie-remote"); err == nil { - return p, nil - } - return "", fmt.Errorf("ollie-remote not found (checked %s)", strings.Join(candidates, ", ")) +// DialSSH connects to a remote tool server via SSH with full config control. +// Alias for RemoteDial. +func DialSSH(ctx context.Context, cfg RemoteConfig) (*Conn, error) { + return RemoteDial(ctx, cfg) } - -// transferBinary sends the ollie-remote binary gzipped+base64 over stdin. -func transferBinary(stdin io.Writer, binData []byte) error { - var compressed bytes.Buffer - gzw := gzip.NewWriter(&compressed) - if _, err := gzw.Write(binData); err != nil { - return err - } - if err := gzw.Close(); err != nil { - return err - } - encoded := base64.StdEncoding.EncodeToString(compressed.Bytes()) - if _, err := fmt.Fprintf(stdin, "OLLIE_DOWNLOAD %d\n", len(encoded)); err != nil { - return err - } - if _, err := stdin.Write([]byte(encoded)); err != nil { - return err - } - return nil -} - -// Close shuts down the SSH connection. -func (s *RemoteServer) Close() error { - s.mu.Lock() - defer s.mu.Unlock() - if s.stdin != nil { - s.stdin.Close() - } - if s.cmd != nil && s.cmd.Process != nil { - syscall.Kill(-s.cmd.Process.Pid, syscall.SIGTERM) - s.cmd.Wait() - } - return nil -} - -// checkUpgrade compares the local ollie-remote binary hash to the deployed -// hash. If they differ, it closes the current connection and re-dials, -// which triggers the bootstrap to upload the new binary. -func (s *RemoteServer) checkUpgrade() error { - localBin, err := findLocalBinary() - if err != nil { - return nil // can't find binary, skip check - } - binData, err := os.ReadFile(localBin) - if err != nil { - return nil - } - hash := sha256.Sum256(binData) - hashStr := hex.EncodeToString(hash[:]) - - if hashStr == s.deployedHash { - return nil - } - - // Hash changed — reconnect. - StreamOutput(s.ctx, "[info] ollie-remote binary changed, upgrading remote...\n") - - s.Close() - - fresh, err := RemoteDial(s.ctx, s.cfg) - if err != nil { - return err - } - - // Swap internals from the fresh connection. - s.mu.Lock() - s.stdin = fresh.stdin - s.stdout = fresh.stdout - s.cmd = fresh.cmd - s.enc = fresh.enc - s.dec = fresh.dec - s.deployedHash = fresh.deployedHash - s.Info = fresh.Info - s.mu.Unlock() - - StreamOutput(s.ctx, "[info] remote binary upgraded successfully\n") - return nil -} - -// --- Server interface --- - -func (s *RemoteServer) ListTools() ([]ToolInfo, error) { - id := s.nextID.Add(1) - req := rpcRequest{ - JSONRPC: "2.0", - ID: id, - Method: "list_tools", - } - - s.mu.Lock() - if err := s.enc.Encode(req); err != nil { - s.mu.Unlock() - return nil, fmt.Errorf("remote list_tools write: %w", err) - } - var resp rpcResponse - if err := s.dec.Decode(&resp); err != nil { - s.mu.Unlock() - return nil, fmt.Errorf("remote list_tools read: %w", err) - } - s.mu.Unlock() - - if resp.Error != nil { - return nil, fmt.Errorf("remote list_tools: %s", resp.Error.Message) - } - - var infos []ToolInfo - if err := json.Unmarshal(resp.Result, &infos); err != nil { - return nil, fmt.Errorf("remote list_tools unmarshal: %w", err) - } - return infos, nil -} - -func (s *RemoteServer) CallTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) { - // Check if local binary has been updated; reconnect if so. - if err := s.checkUpgrade(); err != nil { - StreamOutput(ctx, "[info] remote upgrade failed: "+err.Error()+"\n") - } - - // Forward the call with the tool name as the RPC method. - id := s.nextID.Add(1) - - req := rpcRequest{ - JSONRPC: "2.0", - ID: id, - Method: tool, - Params: args, - } - - s.mu.Lock() - if err := s.enc.Encode(req); err != nil { - s.mu.Unlock() - return nil, fmt.Errorf("remote rpc write: %w", err) - } - - // Read responses, handling interleaved streaming notifications. - for { - var resp rpcResponse - if err := s.dec.Decode(&resp); err != nil { - s.mu.Unlock() - return nil, fmt.Errorf("remote rpc read: %w", err) - } - - // Notification: id == 0 and method is "output". - // (JSON-RPC notifications have no id; our struct decodes missing id as 0.) - if resp.ID == 0 && resp.Result != nil { - // Streaming output notification — emit to context callback. - var notif outputNotification - if json.Unmarshal(resp.Result, ¬if) == nil && notif.Data != "" { - StreamOutput(ctx, notif.Data) - } - continue - } - - // Actual response for our request. - s.mu.Unlock() - if resp.Error != nil { - return nil, fmt.Errorf("remote error: %s", resp.Error.Message) - } - return resp.Result, nil - } -} - -type outputNotification struct { - Data string `json:"data"` -} - -// fetchHostInfo retrieves environment details from the remote host. -func (s *RemoteServer) fetchHostInfo(ctx context.Context) (HostInfo, error) { - id := s.nextID.Add(1) - req := rpcRequest{ - JSONRPC: "2.0", - ID: id, - Method: "host_info", - } - - s.mu.Lock() - defer s.mu.Unlock() - - if err := s.enc.Encode(req); err != nil { - return HostInfo{}, err - } - var resp rpcResponse - if err := s.dec.Decode(&resp); err != nil { - return HostInfo{}, err - } - if resp.Error != nil { - return HostInfo{}, fmt.Errorf("%s", resp.Error.Message) - } - var info HostInfo - if err := json.Unmarshal(resp.Result, &info); err != nil { - return HostInfo{}, err - } - return info, nil -} - -// ping verifies the remote server is responsive. -func (s *RemoteServer) ping(ctx context.Context) error { - id := s.nextID.Add(1) - req := rpcRequest{ - JSONRPC: "2.0", - ID: id, - Method: "ping", - } - - s.mu.Lock() - defer s.mu.Unlock() - - if err := s.enc.Encode(req); err != nil { - return err - } - var resp rpcResponse - if err := s.dec.Decode(&resp); err != nil { - return err - } - if resp.Error != nil { - return fmt.Errorf("%s", resp.Error.Message) - } - return nil -} - -// --- JSON-RPC types --- - -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"` -} - -type rpcError struct { - Code int `json:"code"` - Message string `json:"message"` -} - -// --- helpers --- - -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/rpc.go b/toolsrv/rpc.go new file mode 100644 index 0000000..552a504 --- /dev/null +++ b/toolsrv/rpc.go @@ -0,0 +1,28 @@ +package toolsrv + +import "encoding/json" + +// 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"` +} + +type rpcError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +type outputNotification struct { + Data string `json:"data"` +} diff --git a/toolsrv/socket_test.go b/toolsrv/socket_test.go new file mode 100644 index 0000000..2ac17a7 --- /dev/null +++ b/toolsrv/socket_test.go @@ -0,0 +1,177 @@ +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, toolsrv.LocalAddr(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. + tools, err := conn.ListTools() + if err != nil { + t.Fatalf("ListTools: %v", err) + } + if len(tools) == 0 { + t.Error("expected at least one tool (shell)") + } + found := false + for _, ti := range tools { + if ti.Name == "shell" { + found = true + break + } + } + if !found { + t.Error("expected 'shell' in tool list") + } +} + +// 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, toolsrv.LocalAddr(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 new file mode 100644 index 0000000..ff2aa2a --- /dev/null +++ b/toolsrv/spawn.go @@ -0,0 +1,76 @@ +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/transport.go b/toolsrv/transport.go new file mode 100644 index 0000000..4c8f787 --- /dev/null +++ b/toolsrv/transport.go @@ -0,0 +1,381 @@ +package toolsrv + +import ( + "bufio" + "bytes" + "compress/gzip" + "context" + "crypto/sha256" + _ "embed" + "encoding/base64" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "net" + "os" + "os/exec" + "path/filepath" + "strings" + "sync/atomic" + "syscall" + "time" +) + +//go:embed bootstrap.sh +var bootstrapTemplate string + +// 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) { + binPath, err := findLocalBinary() + if err != nil { + return nil, fmt.Errorf("local dial: %w", err) + } + + // Generate a unique socket path. + sockDir := filepath.Join(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} + + // Apply options: extract yolo flag. + probe := &Server{} + for _, o := range opts { + o(probe) + } + if probe.Yolo { + args = append(args, "--yolo") + } + + // Pass tools path so the subprocess uses installed tools, not embedded. + if tp := os.Getenv("OLLIE_TOOLS_PATH"); tp != "" { + args = append(args, "--tools", tp) + } + + cmd := exec.CommandContext(ctx, binPath, args...) + cmd.SysProcAttr = &syscall.SysProcAttr{ + Setpgid: true, + Pdeathsig: syscall.SIGTERM, + } + + // Set OLLIE_TOOLS_PATH if the host has it configured. + cmd.Env = os.Environ() + if tp := os.Getenv("OLLIE_TOOLS_PATH"); tp != "" { + cmd.Env = appendEnvVar(cmd.Env, "OLLIE_TOOLS_PATH", tp) + } + + // 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 OLLIE_LISTEN_READY on stderr. + readyCh := make(chan struct{}) + go func() { + scanner := bufio.NewScanner(stderrPipe) + for scanner.Scan() { + if strings.TrimSpace(scanner.Text()) == "OLLIE_LISTEN_READY" { + 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) + // Wait up to 3 seconds for graceful exit, then force kill. + 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. +func sshDial(ctx context.Context, cfg RemoteConfig) (*transport, error) { + // Find the local ollie-remote binary to compute hash and transfer if needed. + localBin, err := findLocalBinary() + if err != nil { + return nil, fmt.Errorf("find ollie-remote binary: %w", err) + } + binData, err := os.ReadFile(localBin) + if err != nil { + return nil, fmt.Errorf("read ollie-remote binary: %w", err) + } + hash := sha256.Sum256(binData) + hashStr := hex.EncodeToString(hash[:]) + + // Find landrun binary to deploy alongside ollie-remote. + landrunPath, err := exec.LookPath("landrun") + if err != nil { + return nil, fmt.Errorf("find landrun binary: %w", err) + } + landrunData, err := os.ReadFile(landrunPath) + if err != nil { + return nil, fmt.Errorf("read landrun binary: %w", err) + } + landrunHash := sha256.Sum256(landrunData) + landrunHashStr := hex.EncodeToString(landrunHash[:]) + + // Generate socket paths. + localSockDir := filepath.Join(runtimeDir(), "ollie") + os.MkdirAll(localSockDir, 0700) + localSock := filepath.Join(localSockDir, fmt.Sprintf("remote-%s-%d-%d.sock", cfg.SSHTarget, os.Getpid(), socketSeq.Add(1))) + remoteSock := fmt.Sprintf("/tmp/ollie-toolsrv-%s.sock", hashStr[:12]) + + // Generate bootstrap script with hashes, CWD, and listen path baked in. + bootstrap := bootstrapTemplate + bootstrap = strings.ReplaceAll(bootstrap, "@@HASH@@", hashStr) + bootstrap = strings.ReplaceAll(bootstrap, "@@LANDRUN_HASH@@", landrunHashStr) + cwdEscaped := shellEscape(cfg.CWD) + extraArgs := " --listen " + shellEscape(remoteSock) + if cfg.ToolsPath != "" { + extraArgs += " --tools " + shellEscape(cfg.ToolsPath) + } + if cfg.Yolo { + extraArgs += " --yolo" + } + bootstrap = strings.ReplaceAll(bootstrap, "@@CWD@@", cwdEscaped+extraArgs) + + // Parse optional port from SSHTarget (user@host:port or host:port) + 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 bootstrap script over stdin. + if _, err := stdin.Write([]byte(bootstrap + "\n")); err != nil { + cmd.Process.Kill() + return nil, fmt.Errorf("send bootstrap: %w", err) + } + + // Monitor stderr for bootstrap messages. + type loaderMsg struct { + NeedDownload bool `json:"need_download"` + NeedLandrun bool `json:"need_landrun"` + } + + stderrScanner := bufio.NewScanner(stderrPipe) + loaderCh := make(chan loaderMsg, 1) + readyCh := make(chan struct{}, 1) + go func() { + for stderrScanner.Scan() { + line := strings.TrimSpace(stderrScanner.Text()) + if after, ok := strings.CutPrefix(line, "OLLIE_LOADER_START "); ok { + var msg loaderMsg + json.Unmarshal([]byte(after), &msg) + loaderCh <- msg + if !msg.NeedDownload && !msg.NeedLandrun { + // Binaries are present; still need LISTEN_READY. + } + } else if line == "OLLIE_LOADER_READY" { + // Bootstrap done, binary launched — still need LISTEN_READY. + } else if line == "OLLIE_LISTEN_READY" { + close(readyCh) + } + } + }() + + // Wait for loader start message and handle binary transfer. + select { + case msg := <-loaderCh: + if msg.NeedDownload { + if err := transferBinary(stdin, binData); err != nil { + cmd.Process.Kill() + return nil, fmt.Errorf("transfer ollie-remote: %w", err) + } + } + if msg.NeedLandrun { + if err := transferBinary(stdin, landrunData); err != nil { + cmd.Process.Kill() + return nil, fmt.Errorf("transfer landrun: %w", err) + } + } + case <-time.After(30 * time.Second): + cmd.Process.Kill() + return nil, fmt.Errorf("bootstrap timeout waiting for loader start") + } + + // Wait for the socket to be ready (OLLIE_LISTEN_READY). + select { + case <-readyCh: + case <-time.After(30 * time.Second): + cmd.Process.Kill() + return nil, fmt.Errorf("bootstrap timeout waiting for listen ready") + } + + // Give SSH a moment to set up the local socket forwarding. + // The forwarding is established when the remote socket exists, + // but the local socket may take a beat to appear. + 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 +} + +// 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, ", ")) +} + +// transferBinary sends a binary gzipped+base64 over stdin. +func transferBinary(stdin io.Writer, binData []byte) error { + var compressed bytes.Buffer + gzw := gzip.NewWriter(&compressed) + if _, err := gzw.Write(binData); err != nil { + return err + } + if err := gzw.Close(); err != nil { + return err + } + encoded := base64.StdEncoding.EncodeToString(compressed.Bytes()) + if _, err := fmt.Fprintf(stdin, "OLLIE_DOWNLOAD %d\n", len(encoded)); err != nil { + return err + } + if _, err := stdin.Write([]byte(encoded)); err != nil { + return err + } + return nil +} + +// 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, "'", "'\\''") + "'" +} + +// appendEnvVar replaces or appends a KEY=VALUE in an env slice. +func appendEnvVar(env []string, key, value string) []string { + prefix := key + "=" + for i, e := range env { + if strings.HasPrefix(e, prefix) { + env[i] = prefix + value + return env + } + } + return append(env, prefix+value) +} + +// runtimeDir returns XDG_RUNTIME_DIR or a fallback. +func runtimeDir() string { + if d := os.Getenv("XDG_RUNTIME_DIR"); d != "" { + return d + } + return filepath.Join(os.TempDir(), fmt.Sprintf("runtime-%d", os.Getuid())) +}