ModelRouter/internal/lua/fire_parallel_test.go

package lua

import (
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"
)

// Fire runs same-stage hooks in parallel across plugins. These tests pin the
// three things that make that safe, each of which failed at least once while
// the change was being written.

// twoHookPlugins loads two plugins that both hook request_end.
func twoHookPlugins(t *testing.T, codeA, codeB string) *Plugins {
	t.Helper()
	vm := NewVM(filepath.Join(t.TempDir(), "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatalf("vm start: %v", err)
	}
	t.Cleanup(vm.Stop)
	ps := NewPlugins(vm, filepath.Join(t.TempDir(), "plugins"))
	ps.save.setWakeDelayForTest(time.Millisecond)
	t.Cleanup(ps.Close)
	ps.DisableStatePersistence()
	if err := ps.LoadSource("aaa", codeA); err != nil {
		t.Fatalf("load aaa: %v", err)
	}
	if err := ps.LoadSource("bbb", codeB); err != nil {
		t.Fatalf("load bbb: %v", err)
	}
	return ps
}

func counterPlugin(name string, key string) string {
	// Built by concatenation rather than fmt: %q is not a Go verb, and the first
	// attempt produced Lua source with a stray '=' that only failed at compile
	// time inside three different tests.
	return `
local p = { name = "` + name + `" }
p.state = { n = 0 }
p.hooks = { request_end = "bump" }
function p.bump(payload)
  p.state.n = p.state.n + 1
  if "` + key + `" ~= "" then return { who = "from-` + name + `" } end
  return nil
end
return p
`
}

// TestSingleAndParallelPathsMergeIdentically: the single-plugin fast path and
// the multi-plugin parallel path must produce the same payload.
//
// This is not theoretical. The parallel path was written first and the
// single-plugin path kept its old shape; the merge loop was left off the fast
// path, so a lone plugin returning a table had its return value DISCARDED. The
// existing stage-order test caught it — but only because it happened to check
// the payload after Fire. A plugin that returned fields nobody read would have
// broken silently.
func TestSingleAndParallelPathsMergeIdentically(t *testing.T) {
	single := twoHookPlugins(t, counterPlugin("aaa", "who"), `
local p = { name = "zzz" }
p.state = { n = 0 }
p.hooks = { request_end = "bump" }
function p.bump(payload) p.state.n = p.state.n + 1 return nil end
return p
`)
	// remove the second so exactly one plugin hooks this stage
	single.mu.Lock()
	single.plugins = single.plugins[:1]
	single.mu.Unlock()
	single.rebuild()

	ps := twoHookPlugins(t, counterPlugin("aaa", "who"), counterPlugin("bbb", "who"))

	singleOut := single.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})
	psOut := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})

	if singleOut["who"] != "from-aaa" {
		t.Errorf("single-plugin path dropped the return value: %v", singleOut)
	}
	// Two plugins both return `who`; the merge order decides the winner, and it
	// must be deterministic (load order), not scheduler-dependent.
	if psOut["who"] != "from-bbb" {
		t.Errorf("merge order is not load order: who = %v, want from-bbb", psOut["who"])
	}
}

// TestMergeOrderIsDeterministic: goroutine completion order must not leak into
// the result. Running Fire repeatedly must always yield the same payload, or a
// gateway's behaviour changes run to run with the same plugins installed.
func TestMergeOrderIsDeterministic(t *testing.T) {
	// The plugins differ in cost so a scheduler-dependent merge is visible.
	slow := `
local p = { name = "aaa" }
p.state = { n = 0 }
p.hooks = { request_end = "slow" }
function p.slow(payload)
  local acc = 0
  for i = 1, 3000 do acc = acc + i % 7 end
  p.state.n = p.state.n + 1
  return { winner = "aaa", acc = acc }
end
return p
`
	fast := `
local p = { name = "bbb" }
p.state = { n = 0 }
p.hooks = { request_end = "fast" }
function p.fast(payload)
  p.state.n = p.state.n + 1
  return { winner = "bbb" }
end
return p
`
	ps := twoHookPlugins(t, slow, fast)
	first := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})["winner"]
	for i := 0; i < 60; i++ {
		got := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})["winner"]
		if got != first {
			t.Fatalf("winner changed between runs: %q then %q — merge depends on goroutine "+
				"scheduling, so the same configuration behaves differently run to run", first, got)
		}
	}
	// Load order decides: bbb is loaded second, so it wins.
	if first != "bbb" {
		t.Errorf("winner = %q, want bbb (load order decides, not timing)", first)
	}
}

// TestThrowingHookDoesNotBreakSiblings covers the failure path that golua
// actually produces: a hook raising a Lua error.
//
// The test name used to claim it covered panics, and it used error() to do it.
// Probing the six ways a Lua program can fault (error(), indexing nil, calling
// nil, concatenating nil, arithmetic on nil, unbounded recursion) showed golua
// converts ALL of them into an error RETURN, not a Go panic — so the recover()
// in Fire was untested by that case, and deleting recover() still passed. The
// test was renamed to say what it verifies.
//
// recover() is kept anyway: it guards the Go side of Fire (a nil map write, a
// future change to how the payload is prepared), which is cheap and cannot be
// triggered from Lua today. Claiming it is covered by a Lua test would be the
// kind of assurance that evaporates the first time someone checks.
func TestThrowingHookDoesNotBreakSiblings(t *testing.T) {
	throwing := `
local p = { name = "aaa" }
p.state = {}
p.hooks = { request_end = "boom" }
function p.boom(payload) error("intentional failure") end
return p
`
	healthy := counterPlugin("bbb", "")
	ps := twoHookPlugins(t, throwing, healthy)

	out := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"}) // must not crash
	if out == nil {
		t.Error("Fire returned nil")
	}
	st, ok := ps.State("bbb").(map[string]interface{})
	if !ok {
		t.Fatal("healthy plugin has no state")
	}
	if n := st["n"]; n != float64(1) {
		t.Errorf("healthy plugin did not run alongside the throwing one: n = %v", n)
	}
	he := ps.HookErrors()
	if len(he) == 0 {
		t.Error("a throwing hook was not recorded in hook_errors — the failure would be invisible")
	}
}

// TestParallelHooksAllRunOnce: every plugin must be invoked exactly once per
// Fire. A lost or duplicated goroutine shows up as a wrong total, which for the
// billing plugin means a wrong bill.
func TestParallelHooksAllRunOnce(t *testing.T) {
	ps := twoHookPlugins(t, counterPlugin("aaa", ""), counterPlugin("bbb", ""))
	const fires = 100
	for i := 0; i < fires; i++ {
		ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})
	}
	for _, name := range []string{"aaa", "bbb"} {
		st, ok := ps.State(name).(map[string]interface{})
		if !ok {
			t.Fatalf("%s has no state", name)
		}
		if n := st["n"]; n != float64(fires) {
			t.Errorf("%s counted %v hooks, want %d", name, n, fires)
		}
	}
}

// TestConcurrentFireIsSafe drives Fire from many goroutines at once. Each
// plugin has its own Lua state and its own mutex, so this must hold; the race
// detector is what proves it, not the assertions.
func TestConcurrentFireIsSafe(t *testing.T) {
	ps := twoHookPlugins(t, counterPlugin("aaa", ""), counterPlugin("bbb", ""))
	var wg sync.WaitGroup
	for g := 0; g < 8; g++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < 50; i++ {
				ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
			}
		}()
	}
	wg.Wait()
	for _, name := range []string{"aaa", "bbb"} {
		st, ok := ps.State(name).(map[string]interface{})
		if !ok {
			t.Fatalf("%s has no state", name)
		}
		if n := st["n"]; n != float64(400) {
			t.Errorf("%s counted %v, want 400", name, n)
		}
	}
}

// TestFireDoesNotHoldRegistryLockWhileRunningHooks: a hook that installs or
// disables a plugin takes ps.mu for write. If Fire held ps.mu across the hook,
// that would self-deadlock — the exact failure TrueAgent recorded for its own
// parallel stop path.
func TestFireDoesNotHoldRegistryLockWhileRunningHooks(t *testing.T) {
	reentrant := `
local p = { name = "aaa" }
p.state = {}
p.hooks = { request_end = "reenter" }
function p.reenter(payload)
  -- Reading the registry from inside a hook is the read half of the same lock.
  local _ = #payload
  return nil
end
return p
`
	vm := NewVM(filepath.Join(t.TempDir(), "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	ps := NewPlugins(vm, filepath.Join(t.TempDir(), "plugins"))
	defer ps.Close()
	ps.DisableStatePersistence()
	if err := ps.LoadSource("aaa", reentrant); err != nil {
		t.Fatal(err)
	}
	if err := ps.LoadSource("bbb", counterPlugin("bbb", "")); err != nil {
		t.Fatal(err)
	}

	done := make(chan struct{})
	go func() {
		defer close(done)
		ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})
		// Touch the registry the way an admin request would, right after.
		_ = ps.Count()
		_ = ps.List()
	}()
	select {
	case <-done:
	case <-time.After(10 * time.Second):
		t.Fatal("Fire deadlocked against the registry lock")
	}
}

// TestBillingSurvivesParallelFire runs the REAL bundled plugin next to another
// one. A synthetic counter cannot catch a mismatch between the documented
// payload shape and what the plugin actually reads.
func TestBillingSurvivesParallelFire(t *testing.T) {
	dir := t.TempDir()
	pdir := filepath.Join(dir, "plugins")
	os.MkdirAll(pdir, 0o755)
	src, err := os.ReadFile("plugins/billing.lua")
	if err != nil {
		t.Fatalf("read billing.lua: %v", err)
	}
	if err := os.WriteFile(filepath.Join(pdir, "billing.lua"), src, 0o644); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(pdir, "observer.lua"),
		[]byte(`local p={name="observer"} p.state={n=0}
p.hooks={request_end="see"}
function p.see(payload) p.state.n=p.state.n+1 return nil end
return p`), 0o644); err != nil {
		t.Fatal(err)
	}

	vm := NewVM(filepath.Join(dir, "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	ps := NewPlugins(vm, pdir)
	ps.save.setWakeDelayForTest(time.Millisecond)
	defer ps.Close()
	ps.DisableStatePersistence()
	if err := ps.LoadDir(); err != nil {
		t.Fatal(err)
	}
	if ps.Count() != 2 {
		t.Fatalf("loaded %d plugins, want 2", ps.Count())
	}

	payload := map[string]interface{}{
		"model": "deepseek-v4.1-flash", "source": "commandcode", "key": "k", "ok": true,
		"usage": map[string]interface{}{
			"prompt_tokens": float64(1000), "completion_tokens": float64(100),
			"cache_hit_tokens": float64(0),
		},
	}
	ps.Fire(StageRequestEnd, payload)

	st := ps.State("billing")
	if st == nil {
		t.Fatal("billing produced no state")
	}
	m := st.(map[string]interface{})
	total, ok := m["total"].(map[string]interface{})
	if !ok {
		t.Fatalf("billing.total is %T", m["total"])
	}
	if total["requests"] != float64(1) {
		t.Errorf("billing counted %v requests, want 1 — the plugin was starved by the parallel path",
			total["requests"])
	}
	if n := ps.State("observer").(map[string]interface{})["n"]; n != float64(1) {
		t.Errorf("observer counted %v, want 1", n)
	}
	if len(ps.HookErrors()) != 0 {
		t.Errorf("hook errors under parallel Fire: %v", ps.HookErrors())
	}
}