ollie-core/pkg/remote/remote.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"
	}

	cmd := exec.CommandContext(ctx, "ssh", cfg.SSHTarget, remoteCmd)
	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) {
	// The remote server has the same tools as what's configured.
	// We return a minimal set — the real tool list comes from the
	// local execute.Decl during prompt assembly.
	return nil, nil
}

func (s *Server) CallTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
	// Forward the call as an "execute" RPC
	id := s.nextID.Add(1)

	req := rpcRequest{
		JSONRPC: "2.0",
		ID:      id,
		Method:  "execute",
		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)
	}

	// The result is the execute server's JSON output
	resultJSON, err := json.Marshal(resp.Result)
	if err != nil {
		return nil, fmt.Errorf("remote result marshal: %w", err)
	}
	return resultJSON, 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
}