264 lines
7.0 KiB
Go
264 lines
7.0 KiB
Go
package toolsrv
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"sync"
|
|
"syscall"
|
|
|
|
"ollie/paths"
|
|
)
|
|
|
|
// RPCResponse is a JSON-RPC 2.0 response.
|
|
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"`
|
|
}
|
|
|
|
// lockedEncoder serializes JSON writes to a single writer.
|
|
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)
|
|
}
|
|
|
|
// ServeSocket accepts JSON-RPC connections on a Unix socket. Blocks until
|
|
// ctx is cancelled.
|
|
func (s *Server) ServeSocket(ctx context.Context, sockPath string) error {
|
|
os.Remove(sockPath)
|
|
os.MkdirAll(filepath.Dir(sockPath), 0700)
|
|
|
|
ln, err := net.Listen("unix", sockPath)
|
|
if err != nil {
|
|
return fmt.Errorf("listen %s: %w", sockPath, err)
|
|
}
|
|
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()
|
|
s.ServeRPC(ctx, conn, conn)
|
|
mu.Lock()
|
|
for i, c := range activeConns {
|
|
if c == conn {
|
|
activeConns = append(activeConns[:i], activeConns[i+1:]...)
|
|
break
|
|
}
|
|
}
|
|
mu.Unlock()
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
return nil
|
|
}
|
|
|
|
// ServeRPC reads JSON-RPC 2.0 requests from in and writes responses to out.
|
|
// Blocks until in is closed or ctx is cancelled.
|
|
func (s *Server) ServeRPC(ctx context.Context, 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 s.handleRPC(ctx, req, enc)
|
|
}
|
|
}
|
|
|
|
// CWD returns the server's current working directory.
|
|
func (s *Server) CWD() string {
|
|
s.wdMu.RLock()
|
|
defer s.wdMu.RUnlock()
|
|
return s.cwd
|
|
}
|
|
|
|
// handleRPC dispatches a single JSON-RPC request.
|
|
//
|
|
// Reserved method names (control plane):
|
|
// - list_tools, host_info, ping, set_env, set_cwd
|
|
// - detach, list_detached, signal_detached, get_detached_output, dismiss_detached
|
|
// - tool_load
|
|
//
|
|
// Any other method name is treated as a tool invocation (data plane).
|
|
// A tool with the same name as a reserved method is unreachable.
|
|
func (s *Server) handleRPC(ctx context.Context, req RPCRequest, enc *lockedEncoder) {
|
|
switch req.Method {
|
|
case "list_tools":
|
|
tools, err := s.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 := HostInfo{
|
|
Platform: runtime.GOOS,
|
|
Arch: runtime.GOARCH,
|
|
IsGitRepo: paths.IsGitRepo(s.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
|
|
}
|
|
s.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
|
|
}
|
|
s.SetCWD(params.Dir)
|
|
enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: true})
|
|
|
|
case "detach":
|
|
ok := s.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: s.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 := s.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 "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 := s.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)})
|
|
}
|
|
|
|
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 := s.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 := s.DismissDetached(params.PID)
|
|
enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: ok})
|
|
|
|
default:
|
|
// Data plane: method name is the tool name.
|
|
streamCtx := WithOutputStream(ctx, func(data string) {
|
|
enc.Encode(RPCResponse{
|
|
JSONRPC: "2.0",
|
|
ID: req.ID,
|
|
Result: OutputNotification{Data: data},
|
|
Stream: true,
|
|
})
|
|
})
|
|
result, err := s.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)})
|
|
}
|
|
}
|
|
}
|