338 lines
9.0 KiB
Go
338 lines
9.0 KiB
Go
// Package pluginsdk is the worker side of a Cletus plugin: the message
|
|
// protocol, chat tools, screen and service handlers, and calls to the host.
|
|
// See docs/specification.md and docs/host-api.md.
|
|
package pluginsdk
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"strings"
|
|
"sync"
|
|
)
|
|
|
|
// APIVersion is the worker protocol version the host expects at handshake.
|
|
const APIVersion = 1
|
|
|
|
// MaxMessageBytes is the largest message (one line) either side may send.
|
|
const MaxMessageBytes = 64 << 20
|
|
|
|
// Func handles one screen (ui.*) request. params is the JSON array of the
|
|
// screen's arguments; decode it with Args.
|
|
type Func func(ctx context.Context, h *Host, params json.RawMessage) (any, error)
|
|
|
|
// ServiceFunc handles one service method another plugin calls.
|
|
type ServiceFunc func(ctx context.Context, h *Host, call ServiceCall) (any, error)
|
|
|
|
// ServiceCall is a service request the host routed from another plugin.
|
|
type ServiceCall struct {
|
|
Caller string `json:"caller"`
|
|
Mutating bool `json:"mutating"`
|
|
Args json.RawMessage `json:"args"`
|
|
}
|
|
|
|
// Handlers is everything a worker answers.
|
|
type Handlers struct {
|
|
// Tools are the chat tools. Each must also be listed in plugin.json.
|
|
Tools []Tool
|
|
// UI maps a screen method name ("check", answered as ui.check) to its handler.
|
|
UI map[string]Func
|
|
// Services maps "<service>.<method>" ("gauge.check") to its handler.
|
|
// Each must be declared under "provides" in plugin.json.
|
|
Services map[string]ServiceFunc
|
|
// Busy reports whether work is running that must not be interrupted.
|
|
// Nil means never busy.
|
|
Busy func() bool
|
|
}
|
|
|
|
// ID returns the plugin id the host started the worker as.
|
|
func ID() string { return os.Getenv("CLETUS_PLUGIN_ID") }
|
|
|
|
// DataDir returns the worker's data folder for the signed-in account. It is
|
|
// also the worker's working directory. Empty when not started by a host.
|
|
func DataDir() string { return os.Getenv("CLETUS_PLUGIN_DATA") }
|
|
|
|
// Serve runs the worker on standard input and output until the host shuts
|
|
// it down or closes the pipe. Standard output belongs to the protocol, so
|
|
// Serve points os.Stdout at standard error: a stray fmt.Println lands in
|
|
// the worker log instead of breaking the connection.
|
|
func Serve(h Handlers) error {
|
|
out := os.Stdout
|
|
os.Stdout = os.Stderr
|
|
return ServeIO(context.Background(), os.Stdin, out, h)
|
|
}
|
|
|
|
type message struct {
|
|
ID string `json:"id,omitempty"`
|
|
ParentID string `json:"parent_id,omitempty"`
|
|
Method string `json:"method,omitempty"`
|
|
Params json.RawMessage `json:"params,omitempty"`
|
|
Result json.RawMessage `json:"result,omitempty"`
|
|
Error *wireError `json:"error,omitempty"`
|
|
}
|
|
|
|
type wireError struct {
|
|
Message string `json:"message"`
|
|
}
|
|
|
|
type server struct {
|
|
handlers Handlers
|
|
tools map[string]Tool
|
|
out io.Writer
|
|
writeMu sync.Mutex
|
|
|
|
mu sync.Mutex
|
|
active map[string]context.CancelFunc
|
|
pending map[string]chan message
|
|
next int
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
// ServeIO runs the worker on the given streams. It returns when ctx ends,
|
|
// the input closes, or the host sends plugin.shutdown.
|
|
func ServeIO(ctx context.Context, in io.Reader, out io.Writer, h Handlers) error {
|
|
s := &server{handlers: h, tools: map[string]Tool{}, out: out,
|
|
active: map[string]context.CancelFunc{}, pending: map[string]chan message{}}
|
|
for _, t := range h.Tools {
|
|
if _, dup := s.tools[t.Name]; dup {
|
|
return fmt.Errorf("tool %q is registered twice", t.Name)
|
|
}
|
|
s.tools[t.Name] = t
|
|
}
|
|
ctx, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
defer s.wg.Wait()
|
|
defer s.cancelAll()
|
|
|
|
lines := make(chan []byte)
|
|
readErr := make(chan error, 1)
|
|
go func() {
|
|
reader := bufio.NewReaderSize(in, 64<<10)
|
|
for {
|
|
line, err := readLine(reader)
|
|
if len(line) > 0 {
|
|
select {
|
|
case lines <- line:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
if err != nil {
|
|
readErr <- err
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case err := <-readErr:
|
|
if errors.Is(err, io.EOF) {
|
|
return nil
|
|
}
|
|
return err
|
|
case line := <-lines:
|
|
var m message
|
|
if err := json.Unmarshal(line, &m); err != nil {
|
|
return fmt.Errorf("invalid message from host: %w", err)
|
|
}
|
|
if stop := s.dispatch(ctx, m); stop {
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func readLine(r *bufio.Reader) ([]byte, error) {
|
|
var line []byte
|
|
for {
|
|
chunk, err := r.ReadSlice('\n')
|
|
line = append(line, chunk...)
|
|
if len(line) > MaxMessageBytes+1 {
|
|
return nil, errors.New("message from host is too large")
|
|
}
|
|
if err == bufio.ErrBufferFull {
|
|
continue
|
|
}
|
|
return []byte(strings.TrimSpace(string(line))), err
|
|
}
|
|
}
|
|
|
|
// dispatch handles one incoming message and reports whether to stop.
|
|
func (s *server) dispatch(ctx context.Context, m message) bool {
|
|
switch {
|
|
case m.Method == "$/cancel":
|
|
var id string
|
|
if json.Unmarshal(m.Params, &id) == nil {
|
|
s.mu.Lock()
|
|
if cancel := s.active[id]; cancel != nil {
|
|
cancel()
|
|
}
|
|
s.mu.Unlock()
|
|
}
|
|
case m.Method == "plugin.shutdown":
|
|
s.reply(m.ID, map[string]bool{"ok": true}, nil)
|
|
return true
|
|
case m.Method != "":
|
|
if m.ID == "" {
|
|
return false // a notification this worker doesn't use
|
|
}
|
|
reqCtx, cancel := context.WithCancel(ctx)
|
|
s.mu.Lock()
|
|
s.active[m.ID] = cancel
|
|
s.mu.Unlock()
|
|
s.wg.Add(1)
|
|
go func() {
|
|
defer s.wg.Done()
|
|
defer func() {
|
|
s.mu.Lock()
|
|
delete(s.active, m.ID)
|
|
s.mu.Unlock()
|
|
cancel()
|
|
}()
|
|
result, err := s.handle(reqCtx, m)
|
|
s.reply(m.ID, result, err)
|
|
}()
|
|
case m.ID != "":
|
|
s.mu.Lock()
|
|
ch := s.pending[m.ID]
|
|
delete(s.pending, m.ID)
|
|
s.mu.Unlock()
|
|
if ch != nil {
|
|
ch <- m
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (s *server) handle(ctx context.Context, m message) (result any, err error) {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
result, err = nil, fmt.Errorf("the plugin failed while handling %s: %v", m.Method, r)
|
|
}
|
|
}()
|
|
h := &Host{call: func(ctx context.Context, method string, params any) (json.RawMessage, error) {
|
|
return s.callHost(ctx, m.ID, method, params)
|
|
}}
|
|
switch {
|
|
case m.Method == "plugin.handshake":
|
|
return map[string]int{"api_version": APIVersion}, nil
|
|
case m.Method == "plugin.busy":
|
|
return s.handlers.Busy != nil && s.handlers.Busy(), nil
|
|
case m.Method == "tools.list":
|
|
return s.listTools(), nil
|
|
case m.Method == "tools.invoke":
|
|
return s.invokeTool(ctx, h, m.Params)
|
|
case strings.HasPrefix(m.Method, "ui."):
|
|
fn := s.handlers.UI[strings.TrimPrefix(m.Method, "ui.")]
|
|
if fn == nil {
|
|
return nil, fmt.Errorf("unknown method %s", m.Method)
|
|
}
|
|
return fn(ctx, h, m.Params)
|
|
case strings.HasPrefix(m.Method, "svc."):
|
|
fn := s.handlers.Services[strings.TrimPrefix(m.Method, "svc.")]
|
|
if fn == nil {
|
|
return nil, fmt.Errorf("unknown method %s", m.Method)
|
|
}
|
|
var call ServiceCall
|
|
if err := json.Unmarshal(m.Params, &call); err != nil {
|
|
return nil, fmt.Errorf("invalid service request: %w", err)
|
|
}
|
|
return fn(ctx, h, call)
|
|
}
|
|
return nil, fmt.Errorf("unknown method %s", m.Method)
|
|
}
|
|
|
|
func (s *server) reply(id string, result any, err error) {
|
|
if id == "" {
|
|
return
|
|
}
|
|
m := message{ID: id}
|
|
if err != nil {
|
|
m.Error = &wireError{Message: err.Error()}
|
|
} else {
|
|
raw, marshalErr := json.Marshal(result)
|
|
if marshalErr != nil {
|
|
m.Error = &wireError{Message: "the plugin returned a result that is not JSON: " + marshalErr.Error()}
|
|
} else {
|
|
m.Result = raw
|
|
}
|
|
}
|
|
s.send(m)
|
|
}
|
|
|
|
func (s *server) send(m message) {
|
|
data, err := json.Marshal(m)
|
|
if err != nil || len(data) > MaxMessageBytes {
|
|
if m.ID != "" && m.Method == "" {
|
|
data, _ = json.Marshal(message{ID: m.ID, Error: &wireError{Message: "the plugin's reply is too large"}})
|
|
} else {
|
|
return
|
|
}
|
|
}
|
|
s.writeMu.Lock()
|
|
defer s.writeMu.Unlock()
|
|
_, _ = s.out.Write(append(data, '\n'))
|
|
}
|
|
|
|
func (s *server) callHost(ctx context.Context, parent, method string, params any) (json.RawMessage, error) {
|
|
raw, err := json.Marshal(params)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%s: %w", method, err)
|
|
}
|
|
ch := make(chan message, 1)
|
|
s.mu.Lock()
|
|
s.next++
|
|
id := fmt.Sprintf("w%d", s.next)
|
|
s.pending[id] = ch
|
|
s.mu.Unlock()
|
|
s.send(message{ID: id, ParentID: parent, Method: method, Params: raw})
|
|
select {
|
|
case reply := <-ch:
|
|
if reply.Error != nil {
|
|
return nil, errors.New(reply.Error.Message)
|
|
}
|
|
return reply.Result, nil
|
|
case <-ctx.Done():
|
|
s.mu.Lock()
|
|
delete(s.pending, id)
|
|
s.mu.Unlock()
|
|
cancel, _ := json.Marshal(id)
|
|
s.send(message{Method: "$/cancel", Params: cancel})
|
|
return nil, ctx.Err()
|
|
}
|
|
}
|
|
|
|
func (s *server) cancelAll() {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
for _, cancel := range s.active {
|
|
cancel()
|
|
}
|
|
}
|
|
|
|
// Args decodes a screen request's argument array into targets, in order.
|
|
// Missing trailing arguments leave their targets unchanged.
|
|
func Args(params json.RawMessage, targets ...any) error {
|
|
var items []json.RawMessage
|
|
if len(params) > 0 && string(params) != "null" {
|
|
if err := json.Unmarshal(params, &items); err != nil {
|
|
return fmt.Errorf("screen arguments must be a JSON array: %w", err)
|
|
}
|
|
}
|
|
for i, target := range targets {
|
|
if i >= len(items) {
|
|
break
|
|
}
|
|
if err := json.Unmarshal(items[i], target); err != nil {
|
|
return fmt.Errorf("argument %d: %w", i+1, err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|