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) } }