toolsrv: unified socket-based execution
Execution is now identical for local and remote: - Session calls Spawn()/SpawnRemote() → owns *Process with socket - Clients call Dial(LocalAddr(socket)) → *Conn over Unix socket - ollie-remote serves --listen mode with multi-client accept loop - Builtins (shell, tool_load, skill_load) run in subprocess with registries Lifecycle: - Pdeathsig ensures children die with parent (no orphans) - Cleanup: SIGTERM → 3s grace → SIGKILL (guaranteed) - Socket paths use atomic counter for uniqueness - Session Kill() closes Process Protocol: - Extended RPC: set_env, set_cwd, detach, list_detached, etc. - Shell goes through builtin dispatch (parses 'cmd' field correctly) - Streaming output via JSON-RPC notifications
This commit is contained in:
Levi Neely
parent
eb7bacc243
commit
59b672966f
|
|
@ -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 <path>")
|
||||
fmt.Fprintln(os.Stderr, "usage: ollie-remote serve --cwd <path> [--listen <socket>]")
|
||||
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) {
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ filesystem:
|
|||
- "/lib64"
|
||||
- "/bin"
|
||||
- "/sbin"
|
||||
- "{HOME}/bin"
|
||||
- "{HOME}/go/bin"
|
||||
- "{HOME}/.cargo/bin"
|
||||
- "{HOME}/.local/bin"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
6
justfile
6
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'
|
||||
|
|
|
|||
Binary file not shown.
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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, "'", "'\\''") + "'"
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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"`
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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()))
|
||||
}
|
||||
Loading…
Reference in New Issue