274 lines
6.3 KiB
Go
274 lines
6.3 KiB
Go
// 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 tools.Server that
|
|
// forwards execution calls over the RPC channel.
|
|
package remote
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"os/exec"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"syscall"
|
|
|
|
"ollie/pkg/tools"
|
|
)
|
|
|
|
// Server implements tools.Server by forwarding calls to a remote
|
|
// ollie-remote process over SSH.
|
|
type Server struct {
|
|
mu sync.Mutex
|
|
stdin io.WriteCloser
|
|
stdout io.ReadCloser
|
|
cmd *exec.Cmd
|
|
enc *json.Encoder
|
|
dec *json.Decoder
|
|
nextID atomic.Int64
|
|
}
|
|
|
|
// Config holds the parameters for connecting to a remote host.
|
|
type Config struct {
|
|
// SSHTarget is the SSH destination (e.g., "user@host" or an SSH config alias).
|
|
SSHTarget string
|
|
// CWD is the working directory on the remote host.
|
|
CWD string
|
|
// ToolsPath overrides the remote tool scripts path.
|
|
ToolsPath string
|
|
// Yolo disables sandbox enforcement on the remote.
|
|
Yolo bool
|
|
}
|
|
|
|
// Dial opens an SSH connection to the remote host, launches ollie-remote,
|
|
// and returns a Server ready for tool calls.
|
|
func Dial(ctx context.Context, cfg Config) (*Server, error) {
|
|
// Build the remote command — use full path since SSH non-interactive
|
|
// shells don't have ~/.cache/ollie/bin in PATH.
|
|
remoteCmd := fmt.Sprintf("~/.cache/ollie/bin/ollie-remote serve --cwd %s", shellEscape(cfg.CWD))
|
|
if cfg.ToolsPath != "" {
|
|
remoteCmd += fmt.Sprintf(" --tools %s", shellEscape(cfg.ToolsPath))
|
|
}
|
|
if cfg.Yolo {
|
|
remoteCmd += " --yolo"
|
|
}
|
|
|
|
// Parse optional port from SSHTarget (user@host:port or host:port)
|
|
sshArgs := []string{}
|
|
sshTarget := cfg.SSHTarget
|
|
if host, port, ok := strings.Cut(sshTarget, ":"); ok && port != "" {
|
|
sshTarget = host
|
|
sshArgs = append(sshArgs, "-p", port)
|
|
}
|
|
sshArgs = append(sshArgs, sshTarget, remoteCmd)
|
|
|
|
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)
|
|
}
|
|
// stderr goes to /dev/null for now; could log it
|
|
cmd.Stderr = nil
|
|
|
|
if err := cmd.Start(); err != nil {
|
|
return nil, fmt.Errorf("ssh start: %w", err)
|
|
}
|
|
|
|
s := &Server{
|
|
stdin: stdin,
|
|
stdout: stdout,
|
|
cmd: cmd,
|
|
enc: json.NewEncoder(stdin),
|
|
dec: json.NewDecoder(bufio.NewReader(stdout)),
|
|
}
|
|
|
|
// Verify connectivity with a ping
|
|
if err := s.ping(ctx); err != nil {
|
|
s.Close()
|
|
return nil, fmt.Errorf("remote ping failed: %w", err)
|
|
}
|
|
|
|
return s, nil
|
|
}
|
|
|
|
// Close shuts down the SSH connection.
|
|
func (s *Server) 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
|
|
}
|
|
|
|
// --- tools.Server interface ---
|
|
|
|
func (s *Server) ListTools() ([]tools.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 []tools.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 *Server) CallTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
|
|
// 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)
|
|
}
|
|
|
|
var resp rpcResponse
|
|
if err := s.dec.Decode(&resp); err != nil {
|
|
s.mu.Unlock()
|
|
return nil, fmt.Errorf("remote rpc read: %w", err)
|
|
}
|
|
s.mu.Unlock()
|
|
|
|
if resp.Error != nil {
|
|
return nil, fmt.Errorf("remote error: %s", resp.Error.Message)
|
|
}
|
|
|
|
return resp.Result, nil
|
|
}
|
|
|
|
// ping verifies the remote server is responsive.
|
|
func (s *Server) 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, "'", "'\\''") + "'"
|
|
}
|
|
|
|
// Decl returns a factory function compatible with tools.NewDispatcherFunc.
|
|
// It dials the remote on first call and returns the Server.
|
|
func Decl(cfg Config) func() tools.Server {
|
|
var (
|
|
once sync.Once
|
|
server *Server
|
|
err error
|
|
)
|
|
return func() tools.Server {
|
|
once.Do(func() {
|
|
server, err = Dial(context.Background(), cfg)
|
|
if err != nil {
|
|
// Return a stub that errors on every call
|
|
server = nil
|
|
}
|
|
})
|
|
if server == nil {
|
|
return &errServer{err: err}
|
|
}
|
|
return server
|
|
}
|
|
}
|
|
|
|
// errServer is a tools.Server that returns an error for every call.
|
|
type errServer struct {
|
|
err error
|
|
}
|
|
|
|
func (e *errServer) ListTools() ([]tools.ToolInfo, error) {
|
|
return nil, e.err
|
|
}
|
|
|
|
func (e *errServer) CallTool(_ context.Context, _ string, _ json.RawMessage) (json.RawMessage, error) {
|
|
return nil, e.err
|
|
}
|