306 lines
7.4 KiB
Go
306 lines
7.4 KiB
Go
// ollie-remote — minimal remote execution server for split-brain mode.
|
|
//
|
|
// 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.
|
|
//
|
|
// Usage:
|
|
//
|
|
// ollie-remote serve --cwd /path/to/project [--tools /path/to/tools] [--yolo]
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"embed"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"io/fs"
|
|
"os"
|
|
"os/signal"
|
|
"path/filepath"
|
|
"runtime"
|
|
"syscall"
|
|
|
|
"ollie/toolsrv"
|
|
)
|
|
|
|
//go:embed sandbox-remote.yaml
|
|
var defaultSandboxConfig []byte
|
|
|
|
//go:embed all:tools
|
|
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")
|
|
)
|
|
|
|
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>")
|
|
os.Exit(1)
|
|
}
|
|
// Parse flags after the subcommand
|
|
flag.CommandLine.Parse(os.Args[2:])
|
|
|
|
if *toolsPath == "" {
|
|
// Extract embedded tool scripts to a temp dir.
|
|
tmpTools, err := extractEmbeddedTools()
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "extract tools: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
*toolsPath = tmpTools
|
|
// Clean up on exit.
|
|
defer os.RemoveAll(tmpTools)
|
|
}
|
|
|
|
// Set OLLIE_TOOLS_PATH for tool script discovery
|
|
os.Setenv("OLLIE_TOOLS_PATH", *toolsPath)
|
|
|
|
// Prepend our bin dir to PATH so landrun is found (deployed alongside us)
|
|
home, _ := os.UserHomeDir()
|
|
binDir := filepath.Join(home, ".cache", "ollie", "bin")
|
|
os.Setenv("PATH", binDir+":"+os.Getenv("PATH"))
|
|
|
|
// Set up sandbox config: .ollie-sandbox.yaml in CWD takes precedence,
|
|
// otherwise use the embedded default.
|
|
setupSandboxConfig(*cwd)
|
|
|
|
// Build the tool server
|
|
var opts []toolsrv.Option
|
|
if *yolo {
|
|
opts = append(opts, toolsrv.WithYolo())
|
|
}
|
|
opts = append(opts, toolsrv.WithStrict())
|
|
server := toolsrv.New(*cwd)
|
|
for _, o := range opts { o(server) }
|
|
|
|
// Signal handling
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
sigCh := make(chan os.Signal, 1)
|
|
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
|
go func() {
|
|
<-sigCh
|
|
cancel()
|
|
}()
|
|
|
|
// Serve JSON-RPC over stdin/stdout
|
|
serveRPC(ctx, server, os.Stdin, os.Stdout)
|
|
}
|
|
|
|
// --- JSON-RPC 2.0 ---
|
|
|
|
type rpcRequest struct {
|
|
JSONRPC string `json:"jsonrpc"`
|
|
ID json.RawMessage `json:"id"`
|
|
Method string `json:"method"`
|
|
Params json.RawMessage `json:"params"`
|
|
}
|
|
|
|
type rpcResponse struct {
|
|
JSONRPC string `json:"jsonrpc"`
|
|
ID json.RawMessage `json:"id"`
|
|
Result any `json:"result,omitempty"`
|
|
Error *rpcError `json:"error,omitempty"`
|
|
}
|
|
|
|
type rpcError struct {
|
|
Code int `json:"code"`
|
|
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"`
|
|
}
|
|
|
|
func serveRPC(ctx context.Context, srv *toolsrv.Server, in io.Reader, out io.Writer) {
|
|
dec := json.NewDecoder(in)
|
|
enc := json.NewEncoder(out)
|
|
|
|
for {
|
|
var req rpcRequest
|
|
if err := dec.Decode(&req); err != nil {
|
|
if err == io.EOF || ctx.Err() != nil {
|
|
return
|
|
}
|
|
enc.Encode(rpcResponse{
|
|
JSONRPC: "2.0",
|
|
ID: req.ID,
|
|
Error: &rpcError{Code: -32700, Message: "parse error: " + err.Error()},
|
|
})
|
|
continue
|
|
}
|
|
|
|
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 {
|
|
enc.Encode(rpcResponse{
|
|
JSONRPC: "2.0",
|
|
ID: req.ID,
|
|
Error: &rpcError{Code: -32000, Message: err.Error()},
|
|
})
|
|
} else {
|
|
enc.Encode(rpcResponse{
|
|
JSONRPC: "2.0",
|
|
ID: req.ID,
|
|
Result: tools,
|
|
})
|
|
}
|
|
|
|
case "host_info":
|
|
info := map[string]any{
|
|
"platform": runtime.GOOS,
|
|
"arch": runtime.GOARCH,
|
|
"is_git_repo": isGitRepo(*cwd),
|
|
}
|
|
enc.Encode(rpcResponse{
|
|
JSONRPC: "2.0",
|
|
ID: req.ID,
|
|
Result: info,
|
|
})
|
|
|
|
case "ping":
|
|
enc.Encode(rpcResponse{
|
|
JSONRPC: "2.0",
|
|
ID: req.ID,
|
|
Result: "pong",
|
|
})
|
|
|
|
default:
|
|
// Dispatch to named tool via the tool server's CallTool.
|
|
streamCtx := toolsrv.WithOutputStream(ctx, func(data string) {
|
|
enc.Encode(rpcResponse{
|
|
JSONRPC: "2.0",
|
|
Result: outputNotif{Data: data},
|
|
})
|
|
})
|
|
result, err := srv.CallTool(streamCtx, req.Method, req.Params)
|
|
if err != nil {
|
|
enc.Encode(rpcResponse{
|
|
JSONRPC: "2.0",
|
|
ID: req.ID,
|
|
Error: &rpcError{Code: -32000, Message: err.Error()},
|
|
})
|
|
} else {
|
|
enc.Encode(rpcResponse{
|
|
JSONRPC: "2.0",
|
|
ID: req.ID,
|
|
Result: json.RawMessage(result),
|
|
})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// setupSandboxConfig ensures a sandbox config is available where the execute
|
|
// server expects it (~/.config/ollie/sandbox/default.yaml).
|
|
func setupSandboxConfig(cwd string) {
|
|
home, _ := os.UserHomeDir()
|
|
sandboxDir := filepath.Join(home, ".config", "ollie", "sandbox")
|
|
target := filepath.Join(sandboxDir, "default.yaml")
|
|
|
|
// Check for per-project override
|
|
var configData []byte
|
|
projectConfig := filepath.Join(cwd, ".ollie-sandbox.yaml")
|
|
if data, err := os.ReadFile(projectConfig); err == nil {
|
|
configData = data
|
|
} else {
|
|
configData = defaultSandboxConfig
|
|
}
|
|
|
|
os.MkdirAll(sandboxDir, 0700)
|
|
os.WriteFile(target, configData, 0644)
|
|
}
|
|
|
|
// isGitRepo checks whether the given directory is inside a git repository.
|
|
func isGitRepo(dir string) bool {
|
|
_, err := os.Stat(filepath.Join(dir, ".git"))
|
|
return err == nil
|
|
}
|
|
|
|
// extractEmbeddedTools writes the embedded tool scripts to a temp directory
|
|
// and returns its path.
|
|
func extractEmbeddedTools() (string, error) {
|
|
dir, err := os.MkdirTemp("", "ollie-tools-*")
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
err = fs.WalkDir(embeddedTools, "tools", func(path string, d fs.DirEntry, err error) error {
|
|
if err != nil {
|
|
return err
|
|
}
|
|
rel, _ := filepath.Rel("tools", path)
|
|
if rel == "." {
|
|
return nil
|
|
}
|
|
target := filepath.Join(dir, rel)
|
|
|
|
if d.IsDir() {
|
|
return os.MkdirAll(target, 0755)
|
|
}
|
|
|
|
data, err := embeddedTools.ReadFile(path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return os.WriteFile(target, data, 0755)
|
|
})
|
|
if err != nil {
|
|
os.RemoveAll(dir)
|
|
return "", err
|
|
}
|
|
return dir, nil
|
|
}
|