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)
	}
}

// --- 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"+
			"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
}