ModelRouter/internal/lua/quota_backoff_test.go

package lua

import (
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

// The backoff lives in quota.lua — a shipped plugin, not Go — so these tests read
// the file that actually ships (readShippedPlugin) rather than a copy that can
// drift from it.

// TestQuotaBackoffIsAppliedInSource guards the wiring.
//
// It exists because of a mutation that removed the call site entirely (`if false
// then`) and every other test in this file still passed. The other tests cover
// the decision FUNCTION and the streak RECORDING; neither exists in order to be
// unused. A backoff nobody consults is indistinguishable from no backoff, and the
// symptom is the original one — a dead host re-probed on every page open.
func TestQuotaBackoffIsAppliedInSource(t *testing.T) {
	src := readShippedPlugin(t, "quota.lua")
	if !strings.Contains(src, "if backoffSeconds(seen, known.streak) then") {
		t.Error("quota.lua never consults backoffSeconds for a host with a streak — " +
			"the backoff is computed but not applied")
	}
	if !strings.Contains(src, `reason = "backing off after "`) {
		t.Error("the skip path does not say why it skipped; a silently skipped host " +
			"looks exactly like a host that simply has no balance endpoint")
	}
	// The recording guard has to REPLACE the memo, not merge into it. The bug this
	// guards was `if rep == nil and not HOST_MEMO[host]`, which is true only on a
	// host's first failure — so the streak stayed pinned at 1 forever. Asserted on
	// the source because the mutation that reintroduces it leaves the visible
	// behaviour of the FIRST collect unchanged, and that is all the integration
	// test can reach inside one process.
	if !strings.Contains(src, "if rep == nil then") {
		t.Error("the negative result is recorded under a guard that skips every " +
			"failure after the first — the streak will never grow")
	}
	if strings.Contains(src, "and not HOST_MEMO[host] then") {
		t.Error("the recording guard excludes a host that already has a memo, so " +
			"consecutive failures are never counted")
	}
}

// TestQuotaBackoffWindowIsCapped guards the failure mode an exponential backoff
// has when it is not capped: a host that is down for a week stops being probed
// and the page permanently misreports it as having no balance endpoint.
func TestQuotaBackoffWindowIsCapped(t *testing.T) {
	src := readShippedPlugin(t, "quota.lua")
	for _, want := range []string{"BACKOFF_BASE", "BACKOFF_FACTOR", "BACKOFF_MAX"} {
		if !strings.Contains(src, want) {
			t.Errorf("quota.lua no longer declares %s — the backoff was weakened or removed", want)
		}
	}
}

// TestQuotaBackoffWindows walks the whole policy through the plugin's own
// decision function: for each consecutive failure, how long the host must be left
// alone, that it is still skipped halfway through that window, that it is retried
// once the window passes, and that an absurdly old host is always due.
//
// Why the function rather than the running plugin: HOST_MEMO is a process-local
// cache, restored on a process's first collect and written back afterwards, so
// ageing its memo from Go changes nothing the plugin can see. Two earlier versions
// of this file spent a long time on exactly that dead end, and a test that keeps
// breaking for reasons unrelated to the code under test is worse than no test.
//
// The function is the plugin's own — reached through a shim that appends one
// export — so changing quota.lua's arithmetic moves this test. It is not a
// reimplementation that would agree with a broken plugin.
func TestQuotaBackoffWindows(t *testing.T) {
	ps := quotaWithDecision(t)

	// streak, wait seconds. The last rows are at and past the ceiling.
	for _, c := range []struct {
		streak int
		wait   int
	}{
		{1, 1800},
		{2, 3600},
		{3, 7200},
		{4, 14400},
		{5, 21600},
		{9, 21600},
		{20, 21600},
	} {
		if got := quotaSkip(t, ps, c.streak, c.wait/2); !got {
			t.Logf("  streak=%d age=%d → %v", c.streak, c.wait/2, got)
			t.Errorf("streak=%d: not skipping %ds into a %ds window — the wait is too short",
				c.streak, c.wait/2, c.wait)
		}
		if quotaSkip(t, ps, c.streak, c.wait+60) {
			t.Errorf("streak=%d: still skipping %ds after its %ds window expired — "+
				"this host would never be retried again", c.streak, c.wait+60, c.wait)
		}
		if quotaSkip(t, ps, c.streak, 1<<30) {
			t.Errorf("streak=%d: a host silent for decades is still being skipped", c.streak)
		}
	}
}

// TestQuotaBackoffReturnsOnlyBooleans pins a type rule that produced two silent
// bugs.
//
// In Lua only nil and false are falsy, so every path of a predicate must return
// one of those. This function had two paths returning numbers: `return 0` for
// "no streak", which is truthy and made every caller skip a host that had never
// been probed; and `return BACKOFF_MAX` on the cap path, also truthy, which made
// the ceiling mean "never retried" — the exact failure the ceiling exists to
// prevent.
//
// Both shipped far enough to be found by this assertion rather than by a user
// staring at a page that had quietly stopped probing anything.
func TestQuotaBackoffReturnsOnlyBooleans(t *testing.T) {
	ps := quotaWithDecision(t)
	for _, c := range []struct {
		name string
		age  int
		strk interface{}
	}{
		{"no streak field (every memo written before the backoff existed)", 0, nil},
		{"zero streak", 0, 0},
		{"negative streak", 0, -1},
		{"inside the window", 10, 1},
		{"outside the window", 999999, 1},
		{"at the ceiling", 10, 20},
	} {
		out, err := ps.Call("quota", "__decision", map[string]interface{}{
			"at":     float64(time.Now().Unix() - int64(c.age)),
			"streak": c.strk,
		})
		if err != nil {
			t.Fatalf("%s: %v", c.name, err)
		}
		// A bare boolean is a valid answer: when backoffSeconds returns false the
		// shim's table round-trips as just `false`, because a Lua table holding a
		// single falsy field collapses on encode.
		switch v := out.(type) {
		case bool:
			_ = v // a boolean is exactly what this asserts
		case map[string]interface{}:
			if _, ok := v["r"].(bool); !ok {
				t.Errorf("%s: backoffSeconds returned %T (%v) — in Lua a number is "+
					"TRUTHY, so this caller would skip the host", c.name, v["r"], v["r"])
			}
		default:
			t.Errorf("%s: decision returned %T (%v); want a boolean", c.name, out, out)
		}
	}
}

// TestQuotaBackoffRecordsConsecutiveFailures covers the recording guard through
// the running plugin, because that part IS observable from outside.
//
// The guard used to be `not HOST_MEMO[host]`, true only on a host's first
// failure. After that the memo held a table, the guard was false, nothing was
// written, and the streak stayed pinned at 1 forever — measured by reading back
// {at:... streak:1} after four consecutive failures.
//
// Only two things are asserted, because only two are reachable in one process:
// a first failure records streak 1, and a SKIP does not grow it. A skipped probe
// is not a failure, and a skip that counted would let a never-probed host grow
// without end and never be retried.
func TestQuotaBackoffRecordsConsecutiveFailures(t *testing.T) {
	srv := newCountingServer(t)
	ps := newQuotaPlugins(t, srv.URL)

	args := map[string]interface{}{
		"sources": []interface{}{map[string]interface{}{
			"name": "dead", "base_url": srv.URL + "/v1", "api_key": "k",
		}},
		"name": "dead",
	}
	streak := func() float64 {
		st, _ := ps.State("quota").(map[string]interface{})
		memo, _ := st["probe_memo"].(map[string]interface{})
		for _, v := range memo {
			if e, ok := v.(map[string]interface{}); ok {
				if f, ok := e["streak"].(float64); ok {
					return f
				}
			}
		}
		t.Fatalf("no negative memo recorded: %#v", memo)
		return 0
	}

	if _, err := ps.Call("quota", "collect_source", args); err != nil {
		t.Fatal(err)
	}
	if got := streak(); got != 1 {
		t.Fatalf("streak after the first failure = %.0f, want 1", got)
	}

	before := srv.Hits()
	if _, err := ps.Call("quota", "collect_source", args); err != nil {
		t.Fatal(err)
	}
	if served := srv.Hits() - before; served != 0 {
		t.Errorf("a host inside its backoff window still made %d requests — the "+
			"whole feature is to not make any", served)
	}
	if got := streak(); got != 1 {
		t.Errorf("a SKIP changed the streak to %.0f — a skipped probe is not a failure", got)
	}
}

// TestQuotaMemoKeyIsTheHost pins the key the memo is written and read under.
//
// This is the foundation the whole memoisation rests on, and it was broken in the
// most total way possible: hostOf was written with string.gmatch and a capture
// group, gmatch yields whole matches rather than captures, and with `^` the
// iterator never advanced — so it returned nil and every lookup and every write
// went to HOST_MEMO[nil]. The memo therefore never returned anything for any
// host, and the negative memo, the TTL and the backoff on top of them were all
// dead code that looked alive.
//
// Measured on this gateway before the fix: hostOf(key) matched 0 of the 19
// entries in production state.
//
// The check is on the shipped source plus its runtime behaviour, because the
// failure mode is invisible from outside — a memo that never matches anything
// behaves exactly like a memo that has nothing to remember.
func TestQuotaMemoKeyIsTheHost(t *testing.T) {
	src := readShippedPlugin(t, "quota.lua")
	if !strings.Contains(src, "string.match(u,") {
		t.Error("hostOf does not use string.match — string.gmatch yields whole " +
			"matches, so the capture group is ignored and the function returns nil")
	}

	ps := quotaWithDecision(t)
	for _, c := range []struct{ in, want string }{
		{"https://qijiar.top/v1", "qijiar.top"},
		{"https://api.deepseek.com", "api.deepseek.com"},
		{"http://127.0.0.1:19900", "127.0.0.1:19900"},
		{"https://s2a.frankhan.top:14514/v1", "s2a.frankhan.top:14514"},
	} {
		out, err := ps.Call("quota", "__hostof", map[string]interface{}{"u": c.in})
		if err != nil {
			t.Fatalf("hostOf(%s): %v", c.in, err)
		}
		mm, _ := out.(map[string]interface{})
		got, _ := mm["h"].(string)
		if got != c.want {
			t.Errorf("hostOf(%q) = %q, want %q", c.in, got, c.want)
		}
		if got == "" {
			t.Errorf("hostOf(%q) returned nothing — every lookup would be "+
				"HOST_MEMO[nil]", c.in)
		}
	}

	// The same host under two paths must produce the SAME key, which is the whole
	// reason the memo is keyed this way.
	out, err := ps.Call("quota", "__hostof", map[string]interface{}{"u": "https://x.example/a/v1"})
	if err != nil {
		t.Fatal(err)
	}
	out2, err := ps.Call("quota", "__hostof", map[string]interface{}{"u": "https://x.example/b/v1"})
	if err != nil {
		t.Fatal(err)
	}
	m1, _ := out.(map[string]interface{})
	m2, _ := out2.(map[string]interface{})
	if m1["h"] != m2["h"] {
		t.Errorf("two paths on one host produced different keys: %v and %v — "+
			"they would never share a memo", m1["h"], m2["h"])
	}
}

// --- helpers ---

// readShippedPlugin returns the bundled source, so the tests exercise the file
// that actually ships rather than a copy that can drift from it.
func readShippedPlugin(t *testing.T, name string) string {
	t.Helper()
	b, err := bundledAdapters.ReadFile("plugins/" + name)
	if err != nil {
		t.Fatalf("read embedded %s: %v", name, err)
	}
	return string(b)
}

// quotaWithDecision loads quota.lua with its local backoffSeconds exposed as an
// export. The plugin's own exports list is appended to, so the export allowlist is
// satisfied exactly as it is in production.
func quotaWithDecision(t *testing.T) *Plugins {
	t.Helper()
	src := readShippedPlugin(t, "quota.lua")
	shim := strings.Replace(src, "return PLUGIN",
		"PLUGIN.exports[#PLUGIN.exports+1] = '__decision'\n"+
			// Wrapped, not exported raw: backoffSeconds takes (at, streak)
			// positionally, and the test calls it with a table because that is the
			// shape a Go caller can pass. Exporting the raw function made every
			// assertion read `at = <the argument table>, streak = nil`, so the
			// function took its "no streak" path and returned false for every input
			// — which looked exactly like a backoff that never fires.
			"PLUGIN.__decision = function(q) return { r = backoffSeconds(q.at, q.streak) } end\n"+
			"PLUGIN.exports[#PLUGIN.exports+1] = '__hostof'\n"+
			"PLUGIN.__hostof = function(q) return { h = hostOf(q.u) } end\n"+
			"return PLUGIN", 1)
	if shim == src {
		t.Fatal("could not inject the decision shim — quota.lua no longer ends with 'return PLUGIN'")
	}
	return newQuotaPluginsWith(t, shim, "127.0.0.1")
}

// quotaSkip asks the plugin whether a host that failed `streak` times and was
// last heard from `age` seconds ago should be skipped.
func quotaSkip(t *testing.T, ps *Plugins, streak, age int) bool {
	out, err := ps.Call("quota", "__decision", map[string]interface{}{
		"at":     float64(time.Now().Unix() - int64(age)),
		"streak": float64(streak),
	})
	if err != nil {
		t.Fatalf("decision: %v", err)
	}
	// A single-field table collapses to the bare value on encode, so both shapes
	// mean the same thing here.
	switch v := out.(type) {
	case bool:
		return v
	case map[string]interface{}:
		b, ok := v["r"].(bool)
		if !ok {
			t.Fatalf("decision returned %T, not a boolean: %#v", v["r"], v)
		}
		return b
	default:
		t.Fatalf("decision returned %T (%v), not a boolean", out, out)
		return false
	}
}

// newQuotaPlugins boots the shipped quota.lua against a temp plugin dir.
func newQuotaPlugins(t *testing.T, baseURL string) *Plugins {
	t.Helper()
	return newQuotaPluginsWith(t, readShippedPlugin(t, "quota.lua"), baseURL)
}

func newQuotaPluginsWith(t *testing.T, src, baseURL string) *Plugins {
	t.Helper()
	dir := t.TempDir()
	pluginDir := filepath.Join(dir, "plugins")
	if err := os.MkdirAll(pluginDir, 0o755); err != nil {
		t.Fatal(err)
	}
	vm := NewVM(filepath.Join(dir, "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(vm.Stop)
	ps := NewPlugins(vm, pluginDir)
	t.Cleanup(ps.Close)
	ps.SetOutboundHosts([]string{"127.0.0.1"})
	if err := os.WriteFile(filepath.Join(pluginDir, "quota.lua"), []byte(src), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := ps.LoadDir(); err != nil {
		t.Fatal(err)
	}
	p := ps.find("quota")
	if p == nil {
		t.Fatal("quota plugin did not load")
	}
	if p.LoadError != "" {
		t.Fatalf("quota plugin failed to compile: %s", p.LoadError)
	}
	_ = baseURL
	return ps
}

// countingServer answers 404 on every path and counts what it is asked, so a test
// can tell "the plugin skipped" (no requests) from "the plugin probed and found
// nothing" (requests, no data).
//
// It used to answer 200 with a non-balance body, and the commandcode adapter
// accepted that — so the host was memoised as a SUCCESS and the negative path was
// never reached. The test then failed for the right reason, and would have
// passed for the wrong one had the assertion been looser.
type countingServer struct {
	*httptest.Server
	hits int
}

func (c *countingServer) Hits() int { return c.hits }

func newCountingServer(t *testing.T) *countingServer {
	t.Helper()
	cs := &countingServer{}
	cs.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cs.hits++
		w.WriteHeader(http.StatusNotFound)
		_, _ = w.Write([]byte(`{"error":"not found"}`))
	}))
	t.Cleanup(cs.Close)
	return cs
}

// TestQuotaDueSourcesAndMemoRecordCoverTheShardGap pins the split that makes the
// backoff work at all under sharding.
//
// The conflict, measured on a live 20-source gateway:
//   - probing is network-bound and Plugins.Call holds the plugin mutex for the
//     whole Lua call, so /call/ serialises (29s)
//   - so the page uses ShardCall — and a shard has its own PLUGIN.state
//   - and the backoff lives in PLUGIN.state: fanout left the memo at 0 entries,
//     so dead hosts were re-probed forever
//
// The resolution is to put the DECISION and the RECORDING back on the shared VM
// (microseconds, no network) and shard only the network wait. These two exports
// are that seam, and without them the fanout path silently defeats the backoff.
func TestQuotaDueSourcesAndMemoRecordCoverTheShardGap(t *testing.T) {
	src := readShippedPlugin(t, "quota.lua")
	for _, want := range []string{"due_sources", "memo_record"} {
		if !strings.Contains(src, want) {
			t.Errorf("quota.lua does not export %s — a sharded pass cannot see the "+
				"backoff state, nor leave any behind", want)
		}
		if !strings.Contains(src, `"`+want+`"`) {
			t.Errorf("%s is not listed in PLUGIN.exports, so the page cannot call it", want)
		}
	}
	// due_sources must not probe. If it did, the whole point of asking it first
	// would be lost: the cost would be paid before any shard runs.
	if strings.Contains(src, "due_sources = function(payload)") &&
		strings.Contains(src, "due_sources") {
		// The body is small; check it does not call collectSource.
		i := strings.Index(src, "due_sources = function(payload)")
		j := strings.Index(src[i:], "\nend\n")
		if j > 0 {
			body := src[i : i+j]
			if strings.Contains(body, "collectSource(") {
				t.Error("due_sources calls collectSource — the filter must make no " +
					"network request, or it costs the same as probing everything")
			}
		}
	}
}

// TestQuotaMemoRecordGrowsTheStreakFromShards covers the recording half: a host
// that failed inside a shard must end up in the SHARED memo with a streak, or the
// next refresh has nothing to back off from.
func TestQuotaMemoRecordGrowsTheStreakFromShards(t *testing.T) {
	srv := newCountingServer(t)
	ps := newQuotaPlugins(t, srv.URL)
	host := hostOnly(srv.URL)

	read := func() map[string]interface{} {
		st, _ := ps.State("quota").(map[string]interface{})
		memo, _ := st["probe_memo"].(map[string]interface{})
		return memo
	}

	// A shard reports a host that answered with nothing.
	for i := 1; i <= 3; i++ {
		out, err := ps.Call("quota", "memo_record", map[string]interface{}{
			"results": []interface{}{
				map[string]interface{}{"host": srv.URL + "/v1"},
			},
		})
		if err != nil {
			t.Fatalf("memo_record %d: %v", i, err)
		}
		if out != true {
			t.Fatalf("memo_record returned %#v", out)
		}
	}
	memo := read()
	e, ok := memo[host].(map[string]interface{})
	if !ok {
		t.Fatalf("memo has no entry for %s: %#v", host, memo)
	}
	if got, _ := e["streak"].(float64); got != 3 {
		t.Errorf("streak after three shard failures = %.0f, want 3", got)
	}

	// A host that answers is recorded as a positive memo, which resets the streak.
	if _, err := ps.Call("quota", "memo_record", map[string]interface{}{
		"results": []interface{}{
			map[string]interface{}{"host": srv.URL + "/v1", "vendor": "deepseek"},
		},
	}); err != nil {
		t.Fatal(err)
	}
	if got := read()[host]; got != "deepseek" {
		t.Errorf("a host that answered is memoised as %#v, want the vendor name — "+
			"the streak must not survive a success", got)
	}
}

// TestQuotaDueSourcesSkipsBackedOffHosts is the behaviour the page depends on:
// the filter must keep a backed-off host out of the sharded set.
func TestQuotaDueSourcesSkipsBackedOffHosts(t *testing.T) {
	srv := newCountingServer(t)
	ps := newQuotaPlugins(t, srv.URL)
	host := hostOnly(srv.URL)

	// One failure, which is enough to put the host inside its window.
	if _, err := ps.Call("quota", "memo_record", map[string]interface{}{
		"results": []interface{}{map[string]interface{}{"host": srv.URL + "/v1"}},
	}); err != nil {
		t.Fatal(err)
	}

	// A SECOND host, so the two cases are genuinely independent. The memo is keyed
	// by host, so two paths on one host are one entry — using the same server here
	// would have asserted that a per-host memo behaves per-path.
	other := newCountingServer(t)
	srcs := []interface{}{
		map[string]interface{}{"name": "dead", "base_url": srv.URL + "/v1", "api_key": "k"},
		map[string]interface{}{"name": "fresh", "base_url": other.URL + "/v1", "api_key": "k"},
	}
	out, err := ps.Call("quota", "due_sources", map[string]interface{}{"sources": srcs})
	if err != nil {
		t.Fatal(err)
	}
	m, ok := out.(map[string]interface{})
	if !ok {
		t.Fatalf("due_sources returned %#v", out)
	}
	due, _ := m["due"].([]interface{})
	backing, _ := m["backing_off"].([]interface{})
	if len(backing) != 1 {
		t.Errorf("backing_off = %#v, want exactly the host that failed", backing)
	} else if b, _ := backing[0].(map[string]interface{}); b["host"] != host {
		t.Errorf("backing off %v, want %s — the key must be the host, or it will "+
			"never match on the next refresh", b["host"], host)
	}
	if len(due) != 1 {
		t.Errorf("due = %#v, want only the unprobed host", due)
	} else if d, _ := due[0].(map[string]interface{}); d["name"] != "fresh" {
		t.Errorf("due contains %v, want fresh", d["name"])
	}

	// And it must make no network request — the whole saving depends on it.
	if n := srv.Hits(); n != 0 {
		t.Errorf("due_sources made %d requests; it must decide without probing", n)
	}
}

// hostOnly mirrors the plugin's own hostOf. Used only to CHECK what the plugin
// wrote — seeding a memo under a different key produces a memo the plugin never
// looks up, which cost this file a couple of rounds before it was noticed.
func hostOnly(baseURL string) string {
	u := strings.TrimPrefix(strings.TrimPrefix(baseURL, "https://"), "http://")
	if i := strings.Index(u, "/"); i >= 0 {
		u = u[:i]
	}
	return u
}