467 lines
12 KiB
Go
467 lines
12 KiB
Go
// ollie-remote — tool execution server.
|
|
//
|
|
// Accepts JSON-RPC 2.0 connections. Executes tool calls in a sandboxed
|
|
// environment. Used both locally (spawned by olliesrv) and remotely
|
|
// (deployed via SSH bootstrap).
|
|
//
|
|
// Usage:
|
|
//
|
|
// ollie-remote serve --cwd /path/to/project [--listen /path/to/sock] [--yolo]
|
|
// 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 (
|
|
"context"
|
|
"embed"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"io/fs"
|
|
"net"
|
|
"os"
|
|
"os/signal"
|
|
"path/filepath"
|
|
"runtime"
|
|
"sync"
|
|
"syscall"
|
|
|
|
"ollie/skills"
|
|
"ollie/toolsrv"
|
|
"ollie/tools/builtin"
|
|
)
|
|
|
|
//go:embed sandbox/default.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)")
|
|
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> [--listen <socket>]")
|
|
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 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)
|
|
}
|
|
|
|
// 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()
|
|
}()
|
|
|
|
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 ---
|
|
|
|
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 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 "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",
|
|
})
|
|
|
|
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) {
|
|
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
|
|
}
|