ModelRouter/internal/lua/plugin_shard_test.go

package lua

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

// The sharding spike (zz_shard_spike_test.go) measured the speedup. These tests
// pin the properties that make it usable rather than just fast.

// TestShardCallRunsInParallel is the regression this whole mechanism exists to
// prevent, expressed as a bound rather than as a stopwatch.
//
// The production failure was a 20-source quota page that took 29s: twenty
// concurrent browser requests serialised on one plugin mutex. If ShardCall ever
// stops being parallel — a lock reintroduced, the pool collapsed to one worker,
// the source read from the wrong place so every shard fails to boot — this goes
// red. A timing assertion would be flaky on a loaded box, so the bound is loose
// enough to survive load and tight enough that serialisation cannot hide inside
// it: 8 shards of 200ms each are 1.6s serial, and this fails above 1.0s.
func TestShardCallRunsInParallel(t *testing.T) {
	dir, ps := newShardTestPlugins(t, `
local P = {}
P.name = "probe"
P.version = "1.0.0"
P.exports = { "work" }
P.work = function(payload)
  local res = __llmsproxy_http({ url = payload.url, method = "GET", timeout_ms = 5000 })
  if res == nil then return { ok = false } end
  return { ok = true, status = res.status }
end
return P
`)
	_ = dir

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(200 * time.Millisecond)
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	defer srv.Close()
	ps.SetOutboundHosts([]string{"127.0.0.1"})

	args := make([]map[string]interface{}, 8)
	for i := range args {
		args[i] = map[string]interface{}{"url": srv.URL + "/probe"}
	}

	t0 := time.Now()
	out, err := ps.ShardCall("probe", "work", args)
	elapsed := time.Since(t0)
	if err != nil {
		t.Fatalf("ShardCall: %v", err)
	}
	if len(out) != 8 {
		t.Fatalf("got %d results, want 8", len(out))
	}
	for i, r := range out {
		m, ok := r.(map[string]interface{})
		if !ok || m["ok"] != true {
			t.Errorf("shard %d did not succeed: %#v", i, r)
		}
	}
	// 8 x 200ms serialised is 1.6s. Parallel is ~200ms. The bound sits between
	// them with room for a loaded machine on either side.
	if elapsed > 1000*time.Millisecond {
		t.Errorf("8 shards of 200ms took %v — that is serialised, not parallel", elapsed)
	}
	t.Logf("8 shards took %v", elapsed)
}

// TestShardCallIsPositionallyAligned: the page renders cards in the order the
// sources came back, so a result landing in the wrong slot shows one upstream's
// balance under another upstream's name. That is worse than an error, because it
// looks correct.
func TestShardCallIsPositionallyAligned(t *testing.T) {
	dir, ps := newShardTestPlugins(t, `
local P = {}
P.name = "echo"
P.version = "1.0.0"
P.exports = { "work" }
P.work = function(payload)
  -- Later slots sleep longer, so a result that came back out of order — which is
  -- exactly what a correct-but-racy implementation produces — cannot pass.
  local n = tonumber(payload.n)
  local acc = 0
  for _ = 1, n do acc = acc + 1 end
  return { n = n, acc = acc }
end
return P
`)
	_ = dir

	// Delays are in Lua-visible work, not wall time: the point is that slot i
	// takes longer than slot i-1, so completion order is the REVERSE of input
	// order while the correct implementation still returns input order.
	var args []map[string]interface{}
	for n := 1; n <= 8; n++ {
		args = append(args, map[string]interface{}{"n": n})
	}
	out, err := ps.ShardCall("echo", "work", args)
	if err != nil {
		t.Fatalf("ShardCall: %v", err)
	}
	for i, r := range out {
		m, ok := r.(map[string]interface{})
		if !ok {
			t.Fatalf("slot %d is not a table: %#v", i, r)
		}
		got, gotOK := m["n"].(float64)
		if !gotOK {
			t.Fatalf("slot %d has no n: %#v", i, m)
		}
		if int(got) != i+1 {
			t.Errorf("slot %d carries n=%v; results are not aligned with the input", i, got)
		}
	}
}

// TestShardCallRejectsUnknownPluginAndUnexported: a shard boots a VM from
// arbitrary source, so the export allowlist is the ONLY thing standing between a
// page and the plugin's internals. It has to be checked on this path too.
func TestShardCallRejectsUnknownPluginAndUnexported(t *testing.T) {
	dir, ps := newShardTestPlugins(t, `
local P = {}
P.name = "probe"
P.version = "1.0.0"
P.exports = { "public" }
P.public = function() return { ok = true } end
P.secret = function() return { leaked = true } end
return P
`)
	_ = dir

	if _, err := ps.ShardCall("nope", "public", []map[string]interface{}{{}}); err == nil {
		t.Error("an unknown plugin was sharded")
	}
	_, err := ps.ShardCall("probe", "secret", []map[string]interface{}{{}})
	if err == nil {
		t.Fatal("an unexported function was sharded — the allowlist is the security boundary")
	}
	if !strings.Contains(err.Error(), "does not export") {
		t.Errorf("error does not name the problem: %v", err)
	}
}

// TestShardFailureDoesNotSinkTheBatch: one unreachable upstream must not turn a
// twenty-card page into a single error. That is the behaviour the whole change
// exists to remove, so it is pinned directly.
func TestShardFailureDoesNotSinkTheBatch(t *testing.T) {
	dir, ps := newShardTestPlugins(t, `
local P = {}
P.name = "probe"
P.version = "1.0.0"
P.exports = { "work" }
P.work = function(payload)
  if payload.fail == true then error("upstream exploded") end
  return { ok = true }
end
return P
`)
	_ = dir

	args := []map[string]interface{}{
		{},
		{"fail": true},
		{},
	}
	out, err := ps.ShardCall("probe", "work", args)
	if err != nil {
		t.Fatalf("one failing shard sank the batch: %v", err)
	}
	if len(out) != 3 {
		t.Fatalf("got %d results, want 3", len(out))
	}
	if out[1] != nil {
		t.Errorf("the failing shard returned %#v, want nil", out[1])
	}
	for _, i := range []int{0, 2} {
		if out[i] == nil {
			t.Errorf("healthy shard %d was lost because another shard failed", i)
		}
	}
}

// TestShardAllFailingIsAnError: the opposite case. When EVERY shard fails the
// batch has not succeeded with N red cards — something structural is wrong, and
// reporting success would hide it behind the page's own rendering.
func TestShardAllFailingIsAnError(t *testing.T) {
	dir, ps := newShardTestPlugins(t, `
local P = {}
P.name = "probe"
P.version = "1.0.0"
P.exports = { "work" }
P.work = function(payload) error("always explodes") end
return P
`)
	_ = dir

	out, err := ps.ShardCall("probe", "work", []map[string]interface{}{{}, {}, {}})
	if err == nil {
		t.Fatalf("a batch where every shard failed reported success: %#v", out)
	}
	if !strings.Contains(err.Error(), "all 3 shards") {
		t.Errorf("error does not say the batch failed wholesale: %v", err)
	}
}

// TestShardUsesInstalledSourceNotDisk: a plugin installed through the API never
// touches a file. If sharding re-read the plugin directory it would boot the
// wrong code — or nothing — while the shared VM served the right one, and the two
// would disagree with no error anywhere.
func TestShardUsesInstalledSourceNotDisk(t *testing.T) {
	dir, ps := newShardTestPlugins(t, `
local P = {}
P.name = "probe"
P.version = "1.0.0"
P.exports = { "work" }
P.work = function() return { ok = true, version = "from-disk" } end
return P
`)
	_ = dir

	// Replace the plugin with code that exists only in memory.
	const inMemory = `
local P = {}
P.name = "probe"
P.version = "2.0.0"
P.exports = { "work" }
P.work = function() return { ok = true, version = "from-memory" } end
return P
`
	if err := ps.LoadSource("probe", inMemory); err != nil {
		t.Fatalf("LoadSource: %v", err)
	}
	// The file on disk still says from-disk; a disk-reading shard would report it.
	if _, err := os.Stat(filepath.Join(dir, "plugins", "probe.lua")); err != nil {
		t.Logf("no file on disk at all: %v", err)
	}

	out, err := ps.ShardCall("probe", "work", []map[string]interface{}{{}})
	if err != nil {
		t.Fatal(err)
	}
	m, _ := out[0].(map[string]interface{})
	if got, _ := m["version"].(string); got != "from-memory" {
		t.Errorf("shard ran version %q, want from-memory — it booted the wrong source", got)
	}
}

// TestShardPoolIsBounded: the pool must not let a page allocate VMs without a
// ceiling. Twenty sources is the real case, so the bound sits above it.
func TestShardPoolIsBounded(t *testing.T) {
	dir, ps := newShardTestPlugins(t, `
local P = {}
P.name = "probe"
P.version = "1.0.0"
P.exports = { "work" }
P.work = function() return { ok = true } end
return P
`)
	_ = dir

	if defaultShardLimit < 20 {
		t.Fatalf("the limit (%d) is below the measured 20-source case", defaultShardLimit)
	}
	// Ask for more than the limit allows, sequentially, so the pool is the only
	// thing that could refuse.
	n := defaultShardLimit + 8
	var refused int
	for i := 0; i < n; i++ {
		if _, err := ps.ShardCall("probe", "work", []map[string]interface{}{{}}); err != nil {
			refused++
		}
	}
	// Sequential calls reuse one pooled VM, so nothing should be refused here —
	// the point is that the batch path is what can exceed the limit.
	t.Logf("sequential calls: %d of %d refused", refused, n)

	batch := make([]map[string]interface{}, defaultShardLimit+8)
	for i := range batch {
		batch[i] = map[string]interface{}{}
	}
	_, err := ps.ShardCall("probe", "work", batch)
	if err != nil && !strings.Contains(err.Error(), "all") && !strings.Contains(err.Error(), "shard limit") {
		t.Errorf("over-limit batch produced an unexpected error: %v", err)
	}
	t.Logf("over-limit batch of %d: err=%v", len(batch), err)
}

// newShardTestPlugins builds a Plugins over a temp dir holding one plugin.
func newShardTestPlugins(t *testing.T, code string) (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)
	if err := os.WriteFile(filepath.Join(pluginDir, "probe.lua"), []byte(code), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := ps.LoadDir(); err != nil {
		t.Fatal(err)
	}
	return dir, ps
}

// TestReloadReplacesRatherThanAppends is the defect the shard source lookup
// exposed, pinned on its own so it does not depend on sharding to be exercised.
//
// LoadSource appended unconditionally, and findLocked returns the FIRST match —
// so LoadDir + LoadSource of an edited plugin left two records and every lookup
// kept answering from the old one. Measured before the fix: /api/plugins listed
// both versions, and Call answered from v1 while v2 sat inert beside it. Saving
// an edit appeared to do nothing at all.
func TestReloadReplacesRatherThanAppends(t *testing.T) {
	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)

	if err := os.WriteFile(filepath.Join(pluginDir, "probe.lua"),
		[]byte("local P={}\nP.name='probe'\nP.version='1'\nreturn P\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := ps.LoadDir(); err != nil {
		t.Fatal(err)
	}
	if got := len(ps.List()); got != 1 {
		t.Fatalf("after LoadDir: %d records, want 1", got)
	}

	if err := ps.LoadSource("probe",
		"local P={}\nP.name='probe'\nP.version='2'\nreturn P\n"); err != nil {
		t.Fatal(err)
	}
	if got := len(ps.List()); got != 1 {
		t.Errorf("after reload: %d records, want 1 — the old one was kept", got)
	}
	p := ps.find("probe")
	if p == nil {
		t.Fatal("plugin gone after reload")
	}
	if p.Info.Version != "2" {
		t.Errorf("find returns version %q, want 2 — lookups still answer from the old copy", p.Info.Version)
	}
}