ModelRouter/internal/lua/persist_test.go

package lua

import (
	"encoding/json"
	"os"
	"path/filepath"
	"testing"
	"time"
)

// These cover plugin state persistence: a plugin's totals live in the Lua VM,
// which dies with the process. Measured on the real gateway before this change:
// 218,241 requests and 26.2M prompt tokens gone after one systemctl restart.

func persistTestPlugin() string {
	return `
local plugin = { name = "counter", version = "1.0" }
plugin.state = { n = 0, prices_seen = 0 }
plugin.prices = { rate = 1 }
plugin.hooks = { request_end = "bump" }
function plugin.bump(payload)
  plugin.state.n = plugin.state.n + 1
  return nil
end
return plugin
`
}

// newPersistVM wires a plugin registry on a fresh dir, mirroring newPluginVM.
func newPersistVM(t *testing.T) (*Plugins, string) {
	t.Helper()
	dir := t.TempDir()
	vm := NewVM(filepath.Join(dir, "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatalf("vm start: %v", err)
	}
	t.Cleanup(vm.Stop)
	pdir := filepath.Join(dir, "plugins")
	ps := NewPlugins(vm, pdir)
	// Shorten the saver debounce: Close() waits for the saver goroutine, so
	// with the production 2s interval every test would pay 2s on teardown.
	ps.save.setWakeDelayForTest(2 * time.Millisecond)
	t.Cleanup(ps.Close)
	return ps, pdir
}

func loadCounter(t *testing.T, ps *Plugins) {
	t.Helper()
	if err := ps.LoadSource("counter", persistTestPlugin()); err != nil {
		t.Fatalf("load: %v", err)
	}
}

func stateN(t *testing.T, ps *Plugins) float64 {
	t.Helper()
	st := ps.State("counter")
	m, ok := st.(map[string]interface{})
	if !ok {
		t.Fatalf("state is %T, want map", st)
	}
	n, ok := m["n"].(float64)
	if !ok {
		t.Fatalf("state.n is %T, want float64", m["n"])
	}
	return n
}

// TestStateSurvivesRestart is the defect itself: a rebuilt registry over the
// same plugin dir must come up with the previous totals, not at zero.
func TestStateSurvivesRestart(t *testing.T) {
	ps, pdir := newPersistVM(t)
	loadCounter(t, ps)
	for i := 0; i < 25; i++ {
		ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
	}
	if got := stateN(t, ps); got != 25 {
		t.Fatalf("in-process total = %v, want 25", got)
	}
	// Force the write the saver would do, so the test does not depend on timing.
	ps.save.flush()
	ps.Close()

	// A brand-new registry over the same dir: this is the restart.
	vm := NewVM(filepath.Join(t.TempDir(), "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	ps2 := NewPlugins(vm, pdir)
	defer ps2.Close()
	loadCounter(t, ps2)

	if got := stateN(t, ps2); got != 25 {
		t.Fatalf("total after restart = %v, want 25 — this is the 'restart loses the books' bug", got)
	}
}

// TestPricesAndStateAreSeparate: restoring configuration over history (or the
// reverse) would either erase the totals or resurrect stale prices.
func TestPricesAndStateAreSeparate(t *testing.T) {
	ps, pdir := newPersistVM(t)
	loadCounter(t, ps)
	for i := 0; i < 7; i++ {
		ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
	}
	ps.save.flush()
	ps.Close()

	b, err := os.ReadFile(ps.stateFile("counter"))
	if err != nil {
		t.Fatalf("no state file: %v", err)
	}
	var saved struct {
		State  map[string]interface{} `json:"state"`
		Prices map[string]interface{} `json:"prices"`
	}
	if err := json.Unmarshal(b, &saved); err != nil {
		t.Fatalf("state file is not valid JSON: %v", err)
	}
	if saved.Prices == nil || saved.Prices["rate"] != float64(1) {
		t.Errorf("prices were not persisted separately: %v", saved.Prices)
	}
	if saved.State["n"] != float64(7) {
		t.Errorf("state.n = %v, want 7", saved.State["n"])
	}
	_ = pdir
}

// TestPricesOnlyUpdatePersists guards a real hole: SetState returns early for a
// prices-only payload (correctly leaving state alone), and the first version
// returned before the persistence write — so a reprice was durable in memory
// only and a restart silently reverted to the old prices.
func TestPricesOnlyUpdatePersists(t *testing.T) {
	ps, _ := newPersistVM(t)
	loadCounter(t, ps)
	for i := 0; i < 5; i++ {
		ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
	}
	ps.save.flush()

	err := ps.SetState("counter", map[string]interface{}{
		"prices": map[string]interface{}{"rate": 42},
	})
	if err != nil {
		t.Fatalf("SetState: %v", err)
	}
	if got := stateN(t, ps); got != 5 {
		t.Errorf("a prices-only payload must not touch state: n = %v, want 5", got)
	}

	b, err := os.ReadFile(ps.stateFile("counter"))
	if err != nil {
		t.Fatalf("no state file after a prices-only update: %v", err)
	}
	var saved struct {
		State  map[string]interface{} `json:"state"`
		Prices map[string]interface{} `json:"prices"`
	}
	if err := json.Unmarshal(b, &saved); err != nil {
		t.Fatal(err)
	}
	if saved.Prices["rate"] != float64(42) {
		t.Errorf("new price not persisted: %v — a restart would revert to the old price", saved.Prices)
	}
	if saved.State["n"] != float64(5) {
		t.Errorf("state was clobbered by a price change: %v", saved.State["n"])
	}
}

// TestHookReturnValueSurvivesSnapshot guards the ordering bug: snapshotting
// plugin.state resets the Lua stack, so doing it BEFORE reading the hook's
// return value silently turned every opinionated plugin into a silent one —
// breaking the documented "return a table to merge into payload" contract with
// no error anywhere.
func TestHookReturnValueSurvivesSnapshot(t *testing.T) {
	ps, _ := newPersistVM(t)
	code := `
local plugin = { name = "opinionated" }
plugin.state = { n = 0 }
plugin.hooks = { request_end = "tag" }
function plugin.tag(payload)
  plugin.state.n = plugin.state.n + 1
  return { cost_usd = 1.25, verdict = "billed" }
end
return plugin
`
	if err := ps.LoadSource("opinionated", code); err != nil {
		t.Fatal(err)
	}
	out := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
	if out["cost_usd"] != 1.25 {
		t.Errorf("hook return value lost: payload = %v", out)
	}
	if out["verdict"] != "billed" {
		t.Errorf("merged field lost: %v", out)
	}
	ps.save.flush()
	st := ps.State("opinionated").(map[string]interface{})
	if st["n"] != float64(1) {
		t.Errorf("state.n = %v, want 1 (the hook still ran)", st["n"])
	}
}

// TestCorruptStateFileIsNotFatal: a truncated write must degrade to compiled-in
// defaults, never to a gateway that refuses to start.
func TestCorruptStateFileIsNotFatal(t *testing.T) {
	ps, _ := newPersistVM(t)
	loadCounter(t, ps)
	for i := 0; i < 3; i++ {
		ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
	}
	ps.save.flush()
	ps.Close()

	if err := os.WriteFile(ps.stateFile("counter"), []byte("{not json"), 0o644); err != nil {
		t.Fatal(err)
	}

	vm := NewVM(filepath.Join(t.TempDir(), "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	ps2 := NewPlugins(vm, ps.dir)
	defer ps2.Close()
	var warned bool
	ps2.logf = func(string, ...interface{}) { warned = true }
	loadCounter(t, ps2)

	if got := stateN(t, ps2); got != 0 {
		t.Errorf("with a corrupt file the plugin must fall back to its defaults, got n = %v", got)
	}
	if !warned {
		t.Error("a corrupt state file must warn the operator, not fail silently")
	}
	// And it must still forward.
	ps2.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
}

// TestFlushIsCoalesced: N mutations must not become N writes. A synchronous
// per-request write would put a file write on the hot path.
func TestFlushIsCoalesced(t *testing.T) {
	ps, _ := newPersistVM(t)
	loadCounter(t, ps)
	for i := 0; i < 500; i++ {
		ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
	}
	ps.save.flush()
	// 500 mutations must collapse to a handful of writes, not 500. The bound is
	// loose on purpose: the saver also flushes on a timer, so how many writes
	// happen during the loop depends on how long 500 hook calls take. What must
	// never happen is one write per mutation.
	if w := ps.save.writes.Load(); w > 5 {
		t.Errorf("500 hook calls produced %d file writes; the saver must coalesce", w)
	}
	if got := stateN(t, ps); got != 500 {
		t.Fatalf("n = %v, want 500", got)
	}
}

// TestCloseFlushesTail: the shutdown path must not lose the last interval,
// which would reintroduce the same defect in a smaller window.
func TestCloseFlushesTail(t *testing.T) {
	ps, _ := newPersistVM(t)
	loadCounter(t, ps)
	ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
	ps.Close()
	// Close() already waits for the final flush; no sleep needed.

	b, err := os.ReadFile(ps.stateFile("counter"))
	if err != nil {
		t.Fatalf("Close() did not flush: %v — a shutdown loses the tail", err)
	}
	var saved struct {
		State map[string]interface{} `json:"state"`
	}
	if err := json.Unmarshal(b, &saved); err != nil {
		t.Fatal(err)
	}
	if saved.State["n"] != float64(1) {
		t.Errorf("flushed n = %v, want 1", saved.State["n"])
	}
}

// TestOnDemandCallIsPersisted guards the persistence gap that the quota memo fell
// into, and it needs a plugin with NO hooks to see it.
//
// markDirty fired from two places: hooks and SetState. `Call` — the on-demand path
// a plugin page uses — fired from neither, so a plugin that does all its work in
// exported functions (quota: it declares no hooks at all) mutated state that
// nothing ever wrote. The visible symptom was a health probe that looked inert
// from outside: the memo changed in memory and the file kept the previous
// generation until some unrelated configuration write happened to flush it.
//
// A hookless plugin is the whole point of the fixture: the counter above
// accumulates through a hook, which is the path that already persisted, so it
// cannot see this defect.
func TestOnDemandCallIsPersisted(t *testing.T) {
	ps, _ := newPersistVM(t)
	if err := ps.LoadSource("ondemand", `
local plugin = { name = "ondemand", version = "1.0" }
plugin.state = { n = 0 }
plugin.exports = { "bump" }
function plugin.bump(payload)
  plugin.state.n = plugin.state.n + 1
  return { n = plugin.state.n }
end
return plugin
`); err != nil {
		t.Fatal(err)
	}

	for i := 1; i <= 3; i++ {
		out, err := ps.Call("ondemand", "bump", map[string]interface{}{})
		if err != nil {
			t.Fatalf("call %d: %v", i, err)
		}
		if m, _ := out.(map[string]interface{}); m["n"] != float64(i) {
			t.Fatalf("call %d returned %v, want n=%d", i, out, i)
		}
	}
	// Close() waits for the final flush, so no timing is involved: anything the
	// calls did NOT mark dirty cannot appear here.
	ps.Close()

	b, err := os.ReadFile(ps.stateFile("ondemand"))
	if err != nil {
		t.Fatalf("an exported call did not mark the plugin dirty: %v — a plugin "+
			"with no hooks can never persist its work", err)
	}
	var saved struct {
		State map[string]interface{} `json:"state"`
	}
	if err := json.Unmarshal(b, &saved); err != nil {
		t.Fatal(err)
	}
	if saved.State["n"] != float64(3) {
		t.Errorf("persisted n = %v after three exported calls, want 3", saved.State["n"])
	}
}