cletus-plugins/sdk/pluginsdk/serve_test.go

package pluginsdk

import (
	"bufio"
	"context"
	"encoding/json"
	"io"
	"strings"
	"testing"
	"time"
)

// pipeHost drives a worker over in-memory pipes, as Cletus does.
type pipeHost struct {
	t      *testing.T
	toW    *io.PipeWriter
	fromW  *bufio.Scanner
	done   chan error
	cancel context.CancelFunc
}

func startWorker(t *testing.T, h Handlers) *pipeHost {
	t.Helper()
	inR, inW := io.Pipe()
	outR, outW := io.Pipe()
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- ServeIO(ctx, inR, outW, h); outW.Close() }()
	sc := bufio.NewScanner(outR)
	sc.Buffer(make([]byte, 1<<20), MaxMessageBytes)
	p := &pipeHost{t: t, toW: inW, fromW: sc, done: done, cancel: cancel}
	t.Cleanup(func() { cancel(); inW.Close() })
	return p
}

func (p *pipeHost) send(m message) {
	p.t.Helper()
	data, _ := json.Marshal(m)
	if _, err := p.toW.Write(append(data, '\n')); err != nil {
		p.t.Fatal(err)
	}
}

func (p *pipeHost) read() message {
	p.t.Helper()
	if !p.fromW.Scan() {
		p.t.Fatalf("worker closed: %v", p.fromW.Err())
	}
	var m message
	if err := json.Unmarshal(p.fromW.Bytes(), &m); err != nil {
		p.t.Fatal(err)
	}
	return m
}

func TestHandshakeAndShutdown(t *testing.T) {
	p := startWorker(t, Handlers{})
	p.send(message{ID: "1", Method: "plugin.handshake"})
	if got := p.read(); string(got.Result) != `{"api_version":1}` {
		t.Fatalf("handshake = %s", got.Result)
	}
	p.send(message{ID: "2", Method: "plugin.shutdown"})
	if got := p.read(); got.ID != "2" || got.Error != nil {
		t.Fatalf("shutdown reply = %+v", got)
	}
	select {
	case err := <-p.done:
		if err != nil {
			t.Fatalf("ServeIO = %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("worker did not stop after plugin.shutdown")
	}
}

func TestToolWithHostCall(t *testing.T) {
	tool := Tool{Name: "demo_read", Description: "Reads a file.", Run: func(ctx context.Context, h *Host, call ToolCall) ToolResult {
		var args struct{ Path string }
		if err := call.Decode(&args); err != nil {
			return Fail("%v", err)
		}
		path, err := h.ResolvePath(ctx, args.Path)
		if err != nil {
			return Fail("%v", err)
		}
		return Text("resolved %s", path)
	}}
	p := startWorker(t, Handlers{Tools: []Tool{tool}})

	p.send(message{ID: "1", Method: "tools.list"})
	var list []toolDescriptor
	if err := json.Unmarshal(p.read().Result, &list); err != nil || len(list) != 1 || list[0].Def.Function.Name != "demo_read" {
		t.Fatalf("tools.list = %+v, %v", list, err)
	}

	params, _ := json.Marshal(toolInvocation{Name: "demo_read", Args: `{"Path":"a.txt"}`})
	p.send(message{ID: "2", Method: "tools.invoke", Params: params})
	call := p.read()
	if call.Method != "host.files.resolve" || call.ParentID != "2" || !strings.Contains(string(call.Params), "a.txt") {
		t.Fatalf("host call = %+v", call)
	}
	p.send(message{ID: call.ID, Result: json.RawMessage(`"/work/a.txt"`)})
	var result ToolResult
	if err := json.Unmarshal(p.read().Result, &result); err != nil || result.Text != "resolved /work/a.txt" {
		t.Fatalf("result = %+v, %v", result, err)
	}
}

func TestCancelStopsHandler(t *testing.T) {
	started := make(chan struct{})
	p := startWorker(t, Handlers{UI: map[string]Func{"wait": func(ctx context.Context, h *Host, params json.RawMessage) (any, error) {
		close(started)
		<-ctx.Done()
		return nil, ctx.Err()
	}}})
	p.send(message{ID: "9", Method: "ui.wait", Params: json.RawMessage(`[]`)})
	<-started
	cancelID, _ := json.Marshal("9")
	p.send(message{Method: "$/cancel", Params: cancelID})
	if got := p.read(); got.ID != "9" || got.Error == nil {
		t.Fatalf("cancelled reply = %+v", got)
	}
}

func TestPanicBecomesError(t *testing.T) {
	p := startWorker(t, Handlers{UI: map[string]Func{"boom": func(context.Context, *Host, json.RawMessage) (any, error) {
		panic("bad input")
	}}})
	p.send(message{ID: "1", Method: "ui.boom"})
	if got := p.read(); got.Error == nil || !strings.Contains(got.Error.Message, "bad input") {
		t.Fatalf("reply = %+v", got)
	}
}

func TestServiceAndUnknownMethod(t *testing.T) {
	p := startWorker(t, Handlers{Services: map[string]ServiceFunc{"gauge.check": func(_ context.Context, _ *Host, call ServiceCall) (any, error) {
		return map[string]string{"caller": call.Caller}, nil
	}}})
	p.send(message{ID: "1", Method: "svc.gauge.check", Params: json.RawMessage(`{"caller":"acme.other","args":{}}`)})
	if got := p.read(); string(got.Result) != `{"caller":"acme.other"}` {
		t.Fatalf("service = %+v", got)
	}
	p.send(message{ID: "2", Method: "nope"})
	if got := p.read(); got.Error == nil {
		t.Fatalf("unknown method answered %+v", got)
	}
}

func TestArgs(t *testing.T) {
	var name string
	var count int
	if err := Args(json.RawMessage(`["x", 3]`), &name, &count); err != nil || name != "x" || count != 3 {
		t.Fatalf("Args = %q %d %v", name, count, err)
	}
	if err := Args(json.RawMessage(`{"a":1}`), &name); err == nil {
		t.Fatal("Args accepted an object")
	}
}

func TestNewTestHost(t *testing.T) {
	h := NewTestHost(func(method string, params json.RawMessage) (any, error) {
		return method == "host.files.can_write", nil
	})
	ok, err := h.CanWrite(context.Background(), "x")
	if err != nil || !ok {
		t.Fatalf("CanWrite = %v %v", ok, err)
	}
}