415 lines
11 KiB
Go
415 lines
11 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]
|
|
//
|
|
// When --listen is provided, accepts connections on a Unix socket.
|
|
// Otherwise, serves a single session over stdin/stdout (bootstrap mode).
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"os"
|
|
"os/signal"
|
|
"path/filepath"
|
|
"runtime"
|
|
"sync"
|
|
"syscall"
|
|
|
|
"ollie/env"
|
|
"ollie/toolsrv"
|
|
)
|
|
|
|
var (
|
|
cwd = flag.String("cwd", ".", "working directory for execution")
|
|
listenPath = flag.String("listen", "", "Unix socket path to listen on (default: stdio mode)")
|
|
yolo = flag.Bool("yolo", false, "skip sandbox enforcement")
|
|
)
|
|
|
|
func main() {
|
|
if len(os.Args) < 2 || os.Args[1] != "serve" {
|
|
fmt.Fprintln(os.Stderr, "usage: ollie-remote serve --cwd <path> [--listen <socket>]")
|
|
os.Exit(1)
|
|
}
|
|
flag.CommandLine.Parse(os.Args[2:])
|
|
|
|
// Load env file and set defaults (same as olliesrv).
|
|
env.EnsureDefaults()
|
|
|
|
// Prepend our bin dir to PATH so landrun is found (deployed alongside us)
|
|
home, _ := os.UserHomeDir()
|
|
binDir := filepath.Join(home, ".config", "ollie", "bin")
|
|
os.Setenv("PATH", binDir+":"+os.Getenv("PATH"))
|
|
|
|
// Build the tool server with tool registry.
|
|
toolReg, _ := toolsrv.NewRegistry()
|
|
|
|
// OLLIE_SESSION_ID may be set in the inherited environment (local spawn)
|
|
// or arrive later via set_env RPC (SSH bootstrap). Accept both paths.
|
|
sessionID := os.Getenv("OLLIE_SESSION_ID")
|
|
|
|
var opts []toolsrv.Option
|
|
if *yolo {
|
|
opts = append(opts, toolsrv.WithYolo())
|
|
}
|
|
if toolReg != nil && sessionID != "" {
|
|
opts = append(opts, toolsrv.WithToolRegistry(toolReg, sessionID))
|
|
}
|
|
server := toolsrv.New(*cwd)
|
|
for _, o := range opts {
|
|
o(server)
|
|
}
|
|
|
|
// If tool registry is available but session ID was not set yet,
|
|
// wire up OnEnvSet to attach it when OLLIE_SESSION_ID arrives.
|
|
server.OnEnvSet = func(key, value string) {
|
|
switch key {
|
|
case "OLLIE_SESSION_ID":
|
|
os.Setenv("OLLIE_SESSION_ID", value)
|
|
if toolReg != nil {
|
|
server.SetToolRegistry(toolReg, value)
|
|
}
|
|
case "OLLIE_UNAME":
|
|
os.Setenv("OLLIE_UNAME", value)
|
|
}
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
sigCh := make(chan os.Signal, 1)
|
|
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
|
go func() {
|
|
<-sigCh
|
|
cancel()
|
|
}()
|
|
|
|
if *listenPath != "" {
|
|
serveSocket(ctx, server, *listenPath)
|
|
} else {
|
|
serveRPC(ctx, server, os.Stdin, os.Stdout)
|
|
}
|
|
}
|
|
|
|
func serveSocket(ctx context.Context, srv *toolsrv.Server, sockPath string) {
|
|
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)
|
|
|
|
fmt.Fprintln(os.Stderr, "ListenReady")
|
|
|
|
var mu sync.Mutex
|
|
var activeConns []net.Conn
|
|
var wg sync.WaitGroup
|
|
|
|
go func() {
|
|
<-ctx.Done()
|
|
ln.Close()
|
|
mu.Lock()
|
|
for _, c := range activeConns {
|
|
c.Close()
|
|
}
|
|
mu.Unlock()
|
|
}()
|
|
|
|
for {
|
|
conn, err := ln.Accept()
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
break
|
|
}
|
|
continue
|
|
}
|
|
mu.Lock()
|
|
activeConns = append(activeConns, conn)
|
|
mu.Unlock()
|
|
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
defer conn.Close()
|
|
serveRPC(ctx, srv, conn, conn)
|
|
mu.Lock()
|
|
for i, c := range activeConns {
|
|
if c == conn {
|
|
activeConns = append(activeConns[:i], activeConns[i+1:]...)
|
|
break
|
|
}
|
|
}
|
|
mu.Unlock()
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
}
|
|
|
|
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"`
|
|
Stream bool `json:"stream,omitempty"`
|
|
}
|
|
|
|
type rpcError struct {
|
|
Code int `json:"code"`
|
|
Message string `json:"message"`
|
|
}
|
|
|
|
type outputNotif struct {
|
|
Data string `json:"data"`
|
|
}
|
|
|
|
type lockedEncoder struct {
|
|
mu sync.Mutex
|
|
enc *json.Encoder
|
|
}
|
|
|
|
func (e *lockedEncoder) Encode(v any) error {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
return e.enc.Encode(v)
|
|
}
|
|
|
|
func serveRPC(ctx context.Context, srv *toolsrv.Server, in io.Reader, out io.Writer) {
|
|
dec := json.NewDecoder(in)
|
|
enc := &lockedEncoder{enc: json.NewEncoder(out)}
|
|
|
|
for {
|
|
var req rpcRequest
|
|
if err := dec.Decode(&req); err != nil {
|
|
if err == io.EOF || ctx.Err() != nil {
|
|
return
|
|
}
|
|
enc.Encode(rpcResponse{
|
|
JSONRPC: "2.0",
|
|
ID: req.ID,
|
|
Error: &rpcError{Code: -32700, Message: "parse error: " + err.Error()},
|
|
})
|
|
continue
|
|
}
|
|
|
|
go handleRPC(ctx, srv, req, enc)
|
|
}
|
|
}
|
|
|
|
func handleRPC(ctx context.Context, srv *toolsrv.Server, req rpcRequest, enc *lockedEncoder) {
|
|
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()}})
|
|
return
|
|
}
|
|
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()}})
|
|
return
|
|
}
|
|
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()}})
|
|
return
|
|
}
|
|
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()}})
|
|
return
|
|
}
|
|
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()}})
|
|
return
|
|
}
|
|
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()}})
|
|
return
|
|
}
|
|
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()}})
|
|
return
|
|
}
|
|
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()}})
|
|
return
|
|
}
|
|
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()}})
|
|
return
|
|
}
|
|
tier := srv.ResultTierArgs(params.Name, params.Args)
|
|
enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: tier})
|
|
|
|
case "tool_load":
|
|
var params struct {
|
|
Name string `json:"name"`
|
|
}
|
|
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
|
enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}})
|
|
return
|
|
}
|
|
if err := srv.LoadTool(params.Name); err != nil {
|
|
enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}})
|
|
} else {
|
|
res, _ := json.Marshal(map[string]string{"result": "loaded: " + params.Name})
|
|
enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(res)})
|
|
}
|
|
|
|
default:
|
|
streamCtx := toolsrv.WithOutputStream(ctx, func(data string) {
|
|
enc.Encode(rpcResponse{
|
|
JSONRPC: "2.0",
|
|
ID: req.ID,
|
|
Result: outputNotif{Data: data},
|
|
Stream: true,
|
|
})
|
|
})
|
|
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),
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func isGitRepo(dir string) bool {
|
|
_, err := os.Stat(filepath.Join(dir, ".git"))
|
|
return err == nil
|
|
} |