mesh-catalog/modules/mesh-knowledge/cmd/mesh-knowledge/bus_test.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/nats-io/nats-server/v2/server"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
)

// Over a real bus: a JetStream server is linked into the test and started for it alone; the module is this test
// binary started again as the bundle, speaking MCP on its standard streams to a runtime played here, which does
// what the tool runner does with a bundle: a tool it calls is a request on the bus, an event a publish on the
// module's event subject, an ask a publish on the router's subject, the router's record a key of a bucket only
// the router writes. On the other side of the bus stand the issue tracker's reserve, the agent's judgement, the
// records module (reading the fixture), and the router, played by the test.

// busWorld is one test's bus and what answers on it.
type busWorld struct {
	t      *testing.T
	nc     *nats.Conn
	router jetstream.KeyValue
	mu     sync.Mutex
	events []map[string]any
	asked  map[string][]json.RawMessage
}

func newBusWorld(t *testing.T) *busWorld {
	t.Helper()
	s, err := server.NewServer(&server.Options{Host: "127.0.0.1", Port: server.RANDOM_PORT, JetStream: true, StoreDir: t.TempDir(), NoLog: true, NoSigs: true})
	if err != nil {
		t.Fatal(err)
	}
	go s.Start()
	if !s.ReadyForConnections(30 * time.Second) {
		t.Fatal("no bus within 30s")
	}
	t.Cleanup(func() { s.Shutdown(); s.WaitForShutdown() })
	nc, err := nats.Connect(s.ClientURL())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(nc.Close)
	js, _ := jetstream.New(nc)
	kv, err := js.CreateKeyValue(context.Background(), jetstream.KeyValueConfig{Bucket: "router"})
	if err != nil {
		t.Fatal(err)
	}
	w := &busWorld{t: t, nc: nc, router: kv, asked: map[string][]json.RawMessage{}}
	// The module's events, as anyone on the bus hears them.
	if _, err := nc.Subscribe("mesh.mod.mesh-knowledge.event.>", func(m *nats.Msg) {
		var body map[string]any
		_ = json.Unmarshal(m.Data, &body)
		body["event"] = strings.TrimPrefix(m.Subject, "mesh.mod.mesh-knowledge.event.")
		w.mu.Lock()
		w.events = append(w.events, body)
		w.mu.Unlock()
	}); err != nil {
		t.Fatal(err)
	}
	// What stands on the other side: the issue tracker, the records module, the agent's judgement.
	records := fixture()
	answer := func(key string, do func(body json.RawMessage) (any, error)) {
		if _, err := nc.Subscribe("ask."+key, func(m *nats.Msg) {
			w.mu.Lock()
			w.asked[key] = append(w.asked[key], m.Data)
			w.mu.Unlock()
			out, err := do(m.Data)
			if err != nil {
				_ = m.Respond([]byte(`{"error":` + strconv.Quote(err.Error()) + `}`))
				return
			}
			raw, _ := json.Marshal(map[string]any{"result": out})
			_ = m.Respond(raw)
		}); err != nil {
			t.Fatal(err)
		}
	}
	next := 300
	answer("seat:issue-tracker.reserve", func(body json.RawMessage) (any, error) {
		var a struct{ Kind, For string }
		_ = json.Unmarshal(body, &a)
		next++
		return map[string]any{"kind": a.Kind, "number": next - 1}, nil
	})
	answer("records.records_status", func(json.RawMessage) (any, error) {
		c, err := records.Standing()
		return map[string]any{"commit": c}, err
	})
	answer("records.records_list", func(body json.RawMessage) (any, error) {
		var a struct{ Folder string }
		_ = json.Unmarshal(body, &a)
		f, d, err := records.List(a.Folder)
		return map[string]any{"folders": f, "documents": d}, err
	})
	answer("records.records_read", func(body json.RawMessage) (any, error) {
		var a struct{ Path string }
		_ = json.Unmarshal(body, &a)
		c, commit, err := records.Read(a.Path)
		return map[string]any{"content": c, "commit": commit}, err
	})
	j := &judge{}
	answer("claude-code.claude_code_judge@test", func(body json.RawMessage) (any, error) {
		var a struct {
			Question string         `json:"question"`
			Schema   map[string]any `json:"schema"`
			Model    string         `json:"model"`
		}
		_ = json.Unmarshal(body, &a)
		out, by, err := j.judge(a.Question, a.Schema, a.Model)
		return map[string]any{"answer": out, "model": a.Model, "node": "test", "by": by}, err
	})
	return w
}

func (w *busWorld) of(event string) []map[string]any {
	w.mu.Lock()
	defer w.mu.Unlock()
	var out []map[string]any
	for _, e := range w.events {
		if e["event"] == event {
			out = append(out, e)
		}
	}
	return out
}

// bundle is the knowledge base running as the runtime launches it, and the runtime's half of its conversation.
type bundle struct {
	t    *testing.T
	w    *busWorld
	in   io.WriteCloser
	mu   sync.Mutex
	next int
	wait map[string]chan json.RawMessage
}

func (w *busWorld) start(t *testing.T) (*bundle, string) {
	t.Helper()
	dir := t.TempDir()
	settings := filepath.Join(dir, "settings.json")
	_ = os.WriteFile(settings, []byte(`{"judge-model": "sonnet", "drafting": "off"}`), 0o600)
	status := filepath.Join(dir, "health.json")
	cmd := exec.Command(os.Args[0], "-test.run=^$")
	cmd.Env = append(os.Environ(), bundleVar+"=1", "MESH_KNOWLEDGE_SETTINGS="+settings, "MESH_KNOWLEDGE_STATUS="+status,
		"MESH_SERVED_MODULE="+Module, "MESH_MODULE="+Module, "MESH_NODE=test")
	in, _ := cmd.StdinPipe()
	out, _ := cmd.StdoutPipe()
	cmd.Stderr = logWriter{t}
	if err := cmd.Start(); err != nil {
		t.Fatal(err)
	}
	b := &bundle{t: t, w: w, in: in, wait: map[string]chan json.RawMessage{}}
	go b.read(out)
	t.Cleanup(func() { _ = in.Close(); _ = cmd.Wait() })
	return b, status
}

type logWriter struct{ t *testing.T }

func (l logWriter) Write(p []byte) (int, error) { l.t.Logf("bundle: %s", p); return len(p), nil }

func (b *bundle) say(m map[string]any) {
	m["jsonrpc"] = "2.0"
	raw, _ := json.Marshal(m)
	b.mu.Lock()
	defer b.mu.Unlock()
	_, _ = b.in.Write(append(raw, '\n'))
}

func (b *bundle) read(out io.Reader) {
	sc := bufio.NewScanner(out)
	sc.Buffer(make([]byte, 1<<20), 1<<24)
	for sc.Scan() {
		var m struct {
			ID     json.RawMessage `json:"id"`
			Method string          `json:"method"`
			Params json.RawMessage `json:"params"`
		}
		if json.Unmarshal(sc.Bytes(), &m) != nil {
			continue
		}
		if m.Method == "" {
			var id string
			_ = json.Unmarshal(m.ID, &id)
			b.mu.Lock()
			ch := b.wait[id]
			delete(b.wait, id)
			b.mu.Unlock()
			if ch != nil {
				ch <- append([]byte(nil), sc.Bytes()...)
			}
			continue
		}
		go b.runtime(m.ID, m.Method, m.Params)
	}
}

// runtime is what the tool runner does for the bundle, over the bus.
func (b *bundle) runtime(id json.RawMessage, method string, params json.RawMessage) {
	var p struct {
		Key     string            `json:"key"`
		Body    json.RawMessage   `json:"body"`
		Subject string            `json:"subject"`
		Headers map[string]string `json:"headers"`
	}
	_ = json.Unmarshal(params, &p)
	result, err := func() (any, error) {
		switch method {
		case "mesh/ask":
			msg, err := b.w.nc.Request("ask."+p.Key, p.Body, 30*time.Second)
			if err != nil {
				return nil, fmt.Errorf("nothing answers %s: %v", p.Key, err)
			}
			var r struct {
				Result json.RawMessage `json:"result"`
				Error  string          `json:"error"`
			}
			_ = json.Unmarshal(msg.Data, &r)
			if r.Error != "" {
				return nil, errors.New(r.Error)
			}
			return r.Result, nil
		case "mesh/publish":
			if p.Headers["x-source"] != Module {
				return nil, fmt.Errorf("an event not said as %s", Module)
			}
			return map[string]any{}, b.w.nc.Publish("mesh.mod."+Module+".event."+p.Key, p.Body)
		case "mesh/seat.publish":
			return map[string]any{}, b.w.nc.Publish(p.Subject, p.Body)
		case "mesh/seat.record":
			e, err := b.w.router.Get(context.Background(), p.Key)
			if errors.Is(err, jetstream.ErrKeyNotFound) {
				return nil, nil
			}
			if err != nil {
				return nil, err
			}
			return map[string]any{"key": p.Key, "value": json.RawMessage(e.Value()), "revision": e.Revision()}, nil
		case "mesh/subscribe":
			return map[string]any{}, nil
		}
		return nil, fmt.Errorf("the runtime offers no %s", method)
	}()
	if err != nil {
		b.say(map[string]any{"id": id, "error": map[string]any{"code": -32000, "message": err.Error()}})
		return
	}
	b.say(map[string]any{"id": id, "result": result})
}

// call is one tools/call, as the tool runner makes it.
func (b *bundle) call(tool string, args map[string]any) (json.RawMessage, error) {
	b.mu.Lock()
	b.next++
	id := "runtime-" + strconv.Itoa(b.next)
	ch := make(chan json.RawMessage, 1)
	b.wait[id] = ch
	b.mu.Unlock()
	b.say(map[string]any{"id": id, "method": "tools/call", "params": map[string]any{"name": tool, "arguments": args}})
	select {
	case raw := <-ch:
		var m struct {
			Result struct {
				Content []struct {
					Text string `json:"text"`
				} `json:"content"`
				IsError bool `json:"isError"`
			} `json:"result"`
		}
		if err := json.Unmarshal(raw, &m); err != nil || len(m.Result.Content) == 0 {
			return nil, fmt.Errorf("an answer that is not one: %s", raw)
		}
		if m.Result.IsError {
			return nil, errors.New(m.Result.Content[0].Text)
		}
		return json.RawMessage(m.Result.Content[0].Text), nil
	case <-time.After(60 * time.Second):
		return nil, fmt.Errorf("%s did not answer in 60 s", tool)
	}
}

func TestTheKnowledgeBaseOverTheBus(t *testing.T) {
	w := newBusWorld(t)
	b, _ := w.start(t)
	must := func(raw json.RawMessage, err error) json.RawMessage {
		t.Helper()
		if err != nil {
			t.Fatal(err)
		}
		return raw
	}
	// The import, read through the records module over the bus; run twice.
	var imp Imported
	_ = json.Unmarshal(must(b.call(Seat+".import", map[string]any{})), &imp)
	if imp.Created != 11 || imp.Refused != 0 {
		t.Fatalf("import over the bus: %+v", imp)
	}
	_ = json.Unmarshal(must(b.call(Seat+".import", map[string]any{})), &imp)
	if imp.Created != 0 || !strings.Contains(imp.Said, "Nothing changed") {
		t.Fatalf("a second import: %s", imp.Said)
	}
	// A number from the issue tracker's reserve, asked with its contract.
	var r Reserved
	_ = json.Unmarshal(must(b.call(Seat+".reserve", map[string]any{"kind": "decision", "for": "a merge window that closes sooner"})), &r)
	if r.Number != 300 {
		t.Fatalf("reserved %+v", r)
	}
	w.mu.Lock()
	var asked map[string]any
	_ = json.Unmarshal(w.asked["seat:issue-tracker.reserve"][0], &asked)
	w.mu.Unlock()
	if asked["kind"] != "decision" || asked["for"] != "a merge window that closes sooner" || len(asked) != 2 {
		t.Fatalf("the tracker was asked %v", asked)
	}
	// A proposal, its sections and links as an agent sends them; refused naming the rule, then taken.
	p := valid(300)
	links := []map[string]any{}
	for _, l := range p.Links {
		links = append(links, map[string]any{"kind": l.Kind, "to": l.To, "from": l.From})
	}
	sections := []map[string]any{}
	for _, s := range p.Sections {
		sections = append(sections, map[string]any{"name": s.Name, "text": s.Text})
	}
	args := map[string]any{"number": 300, "title": p.Title, "topic": "gardening", "date": p.Date, "deciders": []string{"jochen"},
		"clauses": p.Clauses, "sections": sections, "links": links}
	if _, err := b.call(Seat+".propose", args); err == nil || !strings.Contains(err.Error(), "- topic:") {
		t.Fatalf("a proposal with an unknown topic: %v", err)
	}
	if _, err := b.call(Seat+".propose", map[string]any{"number": 300, "title": "x", "by": "someone@else"}); err == nil ||
		!strings.Contains(err.Error(), "does not take by") {
		t.Fatalf("an argument naming the caller: %v", err)
	}
	args["topic"] = "the mesh"
	var prop Proposed
	_ = json.Unmarshal(must(b.call(Seat+".propose", args)), &prop)
	if prop.Status != StatusProposed {
		t.Fatalf("%+v", prop)
	}
	// No one reviews a record: the review is the knowledge base's own, and the verb is gone.
	if _, err := b.call(Seat+".review", map[string]any{"number": 300, "verdict": "approve"}); err == nil {
		t.Fatal("a review verb is served")
	}
	// This proposal's caller is unproven over this runtime: it is never reviewed, and says why.
	if !strings.Contains(prop.Says, "proved no caller") {
		t.Fatalf("propose says %q", prop.Says)
	}
	// What holds, in one call within its bound.
	raw := must(b.call(Seat+".ask", map[string]any{"question": "what holds about the merge window, and why"}))
	if len(raw) > DefaultBound || !strings.Contains(string(raw), "ADR 0276") {
		t.Fatalf("%d bytes: %s", len(raw), raw)
	}
	var g Got
	_ = json.Unmarshal(must(b.call(Seat+".get", map[string]any{"number": 300})), &g)
	if g.Status != StatusProposed || len(g.Clauses) != 2 || g.Clauses[0].Holds != "proposed: not in force" {
		t.Fatalf("get: %+v", g)
	}
	var h Historied
	_ = json.Unmarshal(must(b.call(Seat+".history", map[string]any{"number": 300, "from": 1, "to": 2})), &h)
	if len(h.Versions) != 2 || len(h.Diff) == 0 {
		t.Fatalf("history: %+v", h)
	}
	// Every change said on the bus, as the module, with who acted.
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) && len(w.of(EventProposed)) == 0 {
		time.Sleep(50 * time.Millisecond)
	}
	if len(w.of(EventImported)) != 11 || len(w.of(EventReserved)) != 1 || len(w.of(EventProposed)) != 1 {
		t.Fatalf("events: %d imported, %d reserved, %d proposed", len(w.of(EventImported)), len(w.of(EventReserved)), len(w.of(EventProposed)))
	}
	if e := w.of(EventProposed)[0]; e["by"] != Unproven || e["record"] != "ADR 0300" {
		t.Fatalf("the proposed event: %v", e)
	}
	// The health tool says what the watch last saw: here it has seen nothing.
	var health HealthAnswer
	_ = json.Unmarshal(must(b.call(HealthTool, map[string]any{})), &health)
	if health.Healthy || !strings.Contains(health.Why, "no verdict") {
		t.Fatalf("health: %+v", health)
	}
}