so-many-crows/dispatch_test.go

package main

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// fakeDispatchJev answers kind/mode as choice and restates_earlier as noul.
func fakeDispatchJev(t *testing.T) *JevClient {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body struct {
			Questions map[string]any `json:"questions"`
		}
		json.NewDecoder(r.Body).Decode(&body)
		out := map[string]any{}
		for q := range body.Questions {
			switch q {
			case "kind":
				out[q] = map[string]any{"type": "choice", "choice": "say", "confidence": 0.5}
			case "mode":
				out[q] = map[string]any{"type": "choice", "choice": "conspiracy", "confidence": 0.9,
					"probabilities": map[string]float64{"conspiracy": 0.9, "solo": 0.1}}
			case "restates_earlier":
				out[q] = map[string]any{"type": "noul", "noul": 0.85}
			}
		}
		json.NewEncoder(w).Encode(map[string]any{"model": "jev-fake", "answers": out,
			"usage": map[string]int{"input_tokens": 50, "output_tokens": 0}})
	}))
	t.Cleanup(srv.Close)
	return &JevClient{url: srv.URL, key: "tk", model: "jev-fake", threshold: 0.7, client: srv.Client()}
}

// peakBrain tracks concurrent model calls.
type peakBrain struct {
	mu        sync.Mutex
	peak, now int32
}

func (b *peakBrain) Name() string { return "peak" }
func (b *peakBrain) Ask(ctx context.Context, system, prompt string) (string, int, error) {
	c := atomic.AddInt32(&b.now, 1)
	if c > atomic.LoadInt32(&b.peak) {
		atomic.StoreInt32(&b.peak, c)
	}
	time.Sleep(30 * time.Millisecond)
	atomic.AddInt32(&b.now, -1)
	return "KIND: deliver\n[a] and [b] matter", 10, nil
}

func newDispatchServer(t *testing.T, g Guardrails) (*Server, string, *peakBrain) {
	t.Helper()
	dir := t.TempDir()
	ledger, _ := LoadLedger(filepath.Join(dir, "ledger.json"))
	flog, err := OpenFlightLog(filepath.Join(dir, "flight.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { flog.Close() })
	nests, _ := LoadNests(filepath.Join(dir, "nests.json"))
	msgs, _ := LoadMessageStore(filepath.Join(dir, "messages"))
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	pb := &peakBrain{}
	appr := NewWebApprover(time.Minute)
	p := &Pipeline{Flock: NewFlock(pb, g, flog), Ledger: ledger, Log: flog, G: g, Approver: appr}
	p.Flock.SetClassifier(fakeDispatchJev(t))
	s := NewServer(ctx, p, appr, nests).WithMessages(msgs)
	return s, filepath.Join(dir, "flight.jsonl"), pb
}

func TestDispatchCallsGoThroughFlock(t *testing.T) {
	g := DefaultGuardrails()
	flog := ""
	dir := t.TempDir()
	lg, err := OpenFlightLog(filepath.Join(dir, "flight.jsonl"))
	if err != nil {
		t.Fatal(err)
	}
	defer lg.Close()
	flog = filepath.Join(dir, "flight.jsonl")
	f := NewFlock(MockBrain{}, g, lg)
	f.SetClassifier(fakeDispatchJev(t))

	mode, ok := f.ClassifyMode(context.Background(), "is this risky")
	if !ok || mode != "conspiracy" {
		t.Fatalf("want conspiracy from the fake, got %q ok=%v", mode, ok)
	}
	evs, _ := ReadEvents(flog)
	probLogged := false
	for _, e := range evs {
		if e.Kind == "mode" && e.Prob != nil {
			probLogged = true
		}
	}
	if !probLogged {
		t.Fatal("ClassifyMode must write a flight-log event with probabilities")
	}

	// Saturate the flock semaphore: a wrapper call must wait for a slot,
	// proving it goes through Flock (never a direct client call).
	for i := 0; i < g.MaxCrowsInFlight; i++ {
		f.sem <- struct{}{}
	}
	cctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond)
	defer cancel()
	_, blockedOk := f.ClassifyMode(cctx, "risky?")
	if blockedOk {
		t.Fatal("a saturated flock must not answer (semaphore not enforced)")
	}
	for i := 0; i < g.MaxCrowsInFlight; i++ {
		<-f.sem
	}
}

func TestConspiracyInviteCapInFlight(t *testing.T) {
	g := DefaultGuardrails()
	g.MaxCrowsInFlight = 1
	s, _, _ := newDispatchServer(t, g)
	id := newNest(t, serve(s), "Wide", `"Hugin","Munin","Soot"`)
	promote(s, "Hugin", "Munin", "Soot")
	// inviteCount caps at MaxCrowsInFlight=1 -> at most one invited
	if n := s.inviteCount(mustNest(s, id), 4); n != 1 {
		t.Fatalf("want invite capped to 1 by in-flight limit, got %d", n)
	}
}

func TestDispatchGuardrailsEndToEnd(t *testing.T) {
	g := DefaultGuardrails()
	g.MaxCrowsInFlight = 2
	s, logPath, pb := newDispatchServer(t, g)
	ts := serve(s)
	id := newNest(t, ts, "Lab", `"Hugin","Munin"`)
	promote(s, "Hugin", "Munin")
	code, req := post(t, ts.URL+"/api/nests/"+id+"/messages",
		`{"author":"owner","kind":"request","text":"is this launch risky"}`)
	if code != http.StatusCreated || req["mode"] != "conspiracy" || req["tagged_by"] != "header" {
		t.Fatalf("post: %d %v", code, req)
	}
	awaitMessages(t, ts, id, 3) // request + 2 lens claims

	evs, err := ReadEvents(logPath)
	if err != nil {
		t.Fatal(err)
	}
	kinds := map[string]int{}
	for _, e := range evs {
		kinds[e.Kind]++
	}
	if kinds["mode"] == 0 || kinds["tag"] == 0 || kinds["flight"] == 0 {
		t.Fatalf("dispatch decisions must be logged (mode/tag/flight), got %v", kinds)
	}
	if pb.peak > 2 {
		t.Fatalf("flock cap violated: peak concurrent calls %d > 2", pb.peak)
	}
}

func mustNest(s *Server, id string) Nest {
	n, _ := s.nests.Get(id)
	return n
}