ollie/toolsrv/conn.go

package toolsrv

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"syscall"
)

// Conn is a unified JSON-RPC client connected to a tool server subprocess
// (ollie-remote). It satisfies Runner and all extended interfaces that the
// agent layer uses via type assertions. Both local and remote connections
// return the same *Conn — the only difference is the underlying transport.
type Conn struct {
	writeMu   sync.Mutex
	pendingMu sync.Mutex
	pending   map[int64]*pendingCall
	closed    chan struct{}
	closeOnce sync.Once
	readErr   error

	enc    *json.Encoder
	dec    *json.Decoder
	nextID atomic.Int64
	Info   HostInfo

	// cleanup is called on Close (kills subprocess, closes SSH, etc.)
	cleanup func()

	// onToolsChanged is a callback for tool directory change notifications.
	onToolsChanged func(string)
}

// NewConn wraps an io.ReadWriteCloser as a JSON-RPC connection.
// The cleanup function is called when Close() is invoked.
func NewConn(rwc io.ReadWriteCloser, cleanup func()) *Conn {
	c := &Conn{
		enc:     json.NewEncoder(rwc),
		dec:     json.NewDecoder(bufio.NewReader(rwc)),
		cleanup: cleanup,
		pending: make(map[int64]*pendingCall),
		closed:  make(chan struct{}),
	}
	go c.readLoop()
	return c
}

// NewConnSplit creates a Conn from separate reader/writer (e.g. stdin/stdout pipes).
func NewConnSplit(r io.Reader, w io.WriteCloser, cleanup func()) *Conn {
	c := &Conn{
		enc:     json.NewEncoder(w),
		dec:     json.NewDecoder(bufio.NewReader(r)),
		cleanup: cleanup,
		pending: make(map[int64]*pendingCall),
		closed:  make(chan struct{}),
	}
	go c.readLoop()
	return c
}

// --- Runner interface ---

func (c *Conn) ListTools() ([]ToolInfo, error) {
	resp, err := c.call("list_tools", nil)
	if err != nil {
		return nil, err
	}
	var infos []ToolInfo
	if err := json.Unmarshal(resp, &infos); err != nil {
		return nil, fmt.Errorf("list_tools unmarshal: %w", err)
	}
	return infos, nil
}

func (c *Conn) CallTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
	return c.callStreaming(ctx, tool, args)
}

// --- Extended interfaces (used by agent via type assertion) ---

func (c *Conn) SetEnv(key, value string) {
	params, _ := json.Marshal(map[string]string{"key": key, "value": value})
	c.call("set_env", params)
}

func (c *Conn) SetCWD(dir string) {
	params, _ := json.Marshal(map[string]string{"dir": dir})
	c.call("set_cwd", params)
}

func (c *Conn) SetAllowTools(names []string) {
	params, _ := json.Marshal(map[string]any{"names": names})
	c.call("set_allow_tools", params)
}

func (c *Conn) SetOnToolsChanged(fn func(string)) {
	c.pendingMu.Lock()
	c.onToolsChanged = fn
	c.pendingMu.Unlock()
}

func (c *Conn) Detach() bool {
	resp, err := c.call("detach", nil)
	if err != nil {
		return false
	}
	var ok bool
	json.Unmarshal(resp, &ok)
	return ok
}

func (c *Conn) ListDetachedRaw() []any {
	resp, err := c.call("list_detached", nil)
	if err != nil {
		return nil
	}
	var out []any
	json.Unmarshal(resp, &out)
	return out
}

func (c *Conn) SignalDetached(pid int, sig syscall.Signal) error {
	params, _ := json.Marshal(map[string]any{"pid": pid, "signal": int(sig)})
	_, err := c.call("signal_detached", params)
	return err
}

func (c *Conn) GetDetachedOutput(pid int) (string, error) {
	params, _ := json.Marshal(map[string]any{"pid": pid})
	resp, err := c.call("get_detached_output", params)
	if err != nil {
		return "", err
	}
	var out string
	json.Unmarshal(resp, &out)
	return out, nil
}

func (c *Conn) DismissDetached(pid int) bool {
	params, _ := json.Marshal(map[string]any{"pid": pid})
	resp, err := c.call("dismiss_detached", params)
	if err != nil {
		return false
	}
	var ok bool
	json.Unmarshal(resp, &ok)
	return ok
}

// Close shuts down the connection and subprocess.
func (c *Conn) Close() {
	c.fail(fmt.Errorf("connection closed"))
	if c.cleanup != nil {
		c.cleanup()
	}
}

// Ping verifies the remote server is responsive.
func (c *Conn) Ping() error {
	_, err := c.call("ping", nil)
	return err
}

// FetchHostInfo retrieves environment details from the server.
func (c *Conn) FetchHostInfo() (HostInfo, error) {
	resp, err := c.call("host_info", nil)
	if err != nil {
		return HostInfo{}, err
	}
	var info HostInfo
	if err := json.Unmarshal(resp, &info); err != nil {
		return HostInfo{}, err
	}
	return info, nil
}

type pendingCall struct {
	response chan rpcResponse
	ctx      context.Context
}

func (c *Conn) readLoop() {
	for {
		var resp rpcResponse
		if err := c.dec.Decode(&resp); err != nil {
			c.fail(fmt.Errorf("rpc read: %w", err))
			return
		}

		c.pendingMu.Lock()
		call := c.pending[resp.ID]
		c.pendingMu.Unlock()
		if call == nil {
			continue
		}

		// Tool output is sent as an intermediate response with the request ID.
		var notif outputNotification
		if resp.Stream && resp.Error == nil && json.Unmarshal(resp.Result, &notif) == nil {
			if notif.Data != "" {
				StreamOutput(call.ctx, notif.Data)
			}
			continue
		}

		c.pendingMu.Lock()
		delete(c.pending, resp.ID)
		c.pendingMu.Unlock()
		call.response <- resp
	}
}

func (c *Conn) fail(err error) {
	c.closeOnce.Do(func() {
		c.pendingMu.Lock()
		c.readErr = err
		close(c.closed)
		for id, call := range c.pending {
			delete(c.pending, id)
			call.response <- rpcResponse{Error: &rpcError{Code: -32000, Message: err.Error()}}
		}
		c.pendingMu.Unlock()
	})
}

// --- JSON-RPC internals ---

func (c *Conn) call(method string, params json.RawMessage) (json.RawMessage, error) {
	return c.callStreaming(context.Background(), method, params)
}

func (c *Conn) callStreaming(ctx context.Context, method string, params json.RawMessage) (json.RawMessage, error) {
	id := c.nextID.Add(1)
	call := &pendingCall{response: make(chan rpcResponse, 1), ctx: ctx}
	c.pendingMu.Lock()
	select {
	case <-c.closed:
		c.pendingMu.Unlock()
		return nil, fmt.Errorf("rpc %s: connection closed", method)
	default:
	}
	c.pending[id] = call
	c.pendingMu.Unlock()

	c.writeMu.Lock()
	err := c.enc.Encode(rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params})
	c.writeMu.Unlock()
	if err != nil {
		c.pendingMu.Lock()
		delete(c.pending, id)
		c.pendingMu.Unlock()
		return nil, fmt.Errorf("rpc write %s: %w", method, err)
	}

	select {
	case resp := <-call.response:
		if resp.Error != nil {
			return nil, fmt.Errorf("rpc %s: %s", method, resp.Error.Message)
		}
		return resp.Result, nil
	case <-ctx.Done():
		c.pendingMu.Lock()
		delete(c.pending, id)
		c.pendingMu.Unlock()
		return nil, ctx.Err()
	case <-c.closed:
		c.pendingMu.Lock()
		err := c.readErr
		c.pendingMu.Unlock()
		if err == nil {
			err = fmt.Errorf("connection closed")
		}
		return nil, fmt.Errorf("rpc %s: %w", method, err)
	}
}