ModelRouter/internal/gateway/plugin_wiring_test.go

package gateway

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"llmsproxy/internal/config"
	"llmsproxy/internal/core"
)

// newFailingServer serves handler and returns its test server.
func newFailingServer(h http.Handler) *httptest.Server { return httptest.NewServer(h) }

// THE TEST THAT WAS MISSING.
//
// Every other plugin test calls Plugins.Fire() directly. That proves the Lua
// runtime works, but it says nothing about whether the GATEWAY actually fires
// the stages — the wiring in handleChat / singleChat / writeRec could be
// deleted and the whole suite would stay green while no hook ever ran in
// production.
//
// So these tests drive real HTTP requests through the real handler and assert
// on what the plugin observed.

// observed is what a spy plugin recorded. It is written from Lua, so the
// counters live in plugin state and are read back through the state API.
type spyPayload struct {
	Stage  string
	Type   string
	Model  string
	Source string
	OK     bool
	Tokens int64
}

// spyPlugin returns Lua source for a plugin that appends every stage it sees to
// `state.seen`. It is generated per test so each test gets a clean log.
func spyPlugin() string {
	return `
local plugin = { name = "spy", version = "1.0.0" }
plugin.state = { seen = {} }
plugin.hooks = {
  request_start = "on_start",
  routed = "on_routed",
  request_end = "on_end",
}
local function rec(p, stage)
  table.insert(plugin.state.seen, {
    stage = stage, type = p.type or "", model = p.model or "",
    source = p.source or "", ok = p.ok and true or false,
    tokens = p.prompt_tokens or 0,
  })
end
function plugin.on_start(p) rec(p, "request_start") return nil end
function plugin.on_routed(p) rec(p, "routed") return nil end
function plugin.on_end(p) rec(p, "request_end") return nil end
return plugin
`
}

// gatewayWithSpy boots a gateway with one mock upstream and the spy plugin
// loaded, and returns the gateway plus the upstream URL.
func gatewayWithSpy(t *testing.T) (*Gateway, string) {
	t.Helper()
	up := mockUpstream()
	t.Cleanup(up.Close)

	dir := t.TempDir()
	cfgPath := filepath.Join(dir, "config.yaml")
	cfgBody := "listen: :0\n" +
		"adapter_dir: " + filepath.Join(dir, "adapters") + "\n" +
		"plugin_dir: " + filepath.Join(dir, "plugins") + "\n" +
		"runtime_file: " + filepath.Join(dir, "runtime.json") + "\n" +
		"gateway_keys:\n  - sk-test\n"
	if err := os.WriteFile(cfgPath, []byte(cfgBody), 0600); err != nil {
		t.Fatal(err)
	}
	cfg, err := config.Load(cfgPath)
	if err != nil {
		t.Fatal(err)
	}
	cfg.Sources = []config.Source{{
		Name:    "spyup",
		BaseURL: up.URL,
		Adapter: "openai",
		APIKey:  "sk-x",
		Models:  []config.Model{{ID: "spy-model", Kind: "chat"}},
	}}
	if err := cfg.ApplyDefaults(); err != nil {
		t.Fatal(err)
	}
	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatalf("core: %v", err)
	}
	t.Cleanup(c.Close)
	if err := c.Plugins().LoadSource("spy", spyPlugin()); err != nil {
		t.Fatalf("load spy: %v", err)
	}
	g, err := New(c)
	if err != nil {
		t.Fatalf("gateway: %v", err)
	}
	return g, up.URL
}

// seenStages reads the spy's log out of its published state.
func seenStages(t *testing.T, g *Gateway) []map[string]interface{} {
	t.Helper()
	raw := g.core.Plugins().State("spy")
	if raw == nil {
		return nil
	}
	b, err := json.Marshal(raw)
	if err != nil {
		t.Fatal(err)
	}
	var st struct {
		Seen []map[string]interface{} `json:"seen"`
	}
	if err := json.Unmarshal(b, &st); err != nil {
		t.Fatal(err)
	}
	return st.Seen
}

func stageNames(seen []map[string]interface{}) []string {
	out := make([]string, 0, len(seen))
	for _, s := range seen {
		out = append(out, s["stage"].(string))
	}
	return out
}

// TestHooksFireOnRealDirectChat: a plain non-streaming chat through the real
// handler must produce request_start, routed and request_end — in that order,
// with the real source and model filled in.
func TestHooksFireOnRealDirectChat(t *testing.T) {
	g, _ := gatewayWithSpy(t)
	rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
		`{"model":"spy-model","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != http.StatusOK {
		t.Fatalf("chat status=%d body=%s", rr.Code, rr.Body.String())
	}
	seen := seenStages(t, g)
	got := stageNames(seen)
	want := []string{"request_start", "routed", "request_end"}
	if len(got) != 3 {
		t.Fatalf("stages = %v, want %v (a hook that never fires is the failure "+
			"this test exists for)", got, want)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("stage order = %v, want %v", got, want)
		}
	}
	// The payload must be real data, not empty scaffolding.
	end := seen[2]
	if end["source"] != "spyup" {
		t.Errorf("request_end source = %v, want spyup", end["source"])
	}
	if end["model"] != "spy-model" {
		t.Errorf("request_end model = %v, want spy-model", end["model"])
	}
	if end["ok"] != true {
		t.Errorf("request_end ok = %v, want true", end["ok"])
	}
	if toInt(t, end["tokens"]) <= 0 {
		t.Errorf("request_end prompt_tokens = %v, want the mock upstream's real usage", end["tokens"])
	}
}

// TestHooksFireOnRealStreamChat: streaming is a DIFFERENT code path (the record
// is written from a deferred call after the stream ends), so it needs its own
// assertion. A hook that only worked on the non-streaming path would be a real
// bug for agent clients, which are the main streaming users.
func TestHooksFireOnRealStreamChat(t *testing.T) {
	g, _ := gatewayWithSpy(t)
	rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
		`{"model":"spy-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != http.StatusOK {
		t.Fatalf("stream status=%d body=%s", rr.Code, rr.Body.String())
	}
	got := stageNames(seenStages(t, g))
	if len(got) != 3 {
		t.Fatalf("stream stages = %v, want all three", got)
	}
	if got[0] != "request_start" || got[2] != "request_end" {
		t.Errorf("stream stage order = %v", got)
	}
}

// TestHooksFireOnAutoRequest: the AUTO chain path is a third code path. It also
// must fire, and routed must report tier -2 (resolved by the chain) rather than
// the direct path's -1.
func TestHooksFireOnAutoRequest(t *testing.T) {
	g, _ := gatewayWithSpy(t)
	// Register the model in the AUTO chain via the admin API.
	put := doReq(t, g, http.MethodPut, "/api/auto", `{"rules":[{"model":"spy-model"}]}`)
	if put.Code != http.StatusOK {
		t.Fatalf("save auto: %d %s", put.Code, put.Body.String())
	}
	rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
		`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != http.StatusOK {
		t.Fatalf("auto chat status=%d body=%s", rr.Code, rr.Body.String())
	}
	seen := seenStages(t, g)
	if len(seen) != 3 {
		t.Fatalf("AUTO stages = %v, want three", stageNames(seen))
	}
	if seen[0]["model"] != "AUTO" {
		t.Errorf("request_start model = %v, want the requested \"AUTO\"", seen[0]["model"])
	}
	if seen[1]["model"] != "spy-model" {
		t.Errorf("routed model = %v, want the RESOLVED model", seen[1]["model"])
	}
}

// TestHooksFireOnFailedRequest is the one that matters most for accounting: a
// request that 500s upstream must still reach request_end, or a billing plugin
// would never see failed traffic.
func TestHooksFireOnFailedRequest(t *testing.T) {
	// Point at an upstream that always fails.
	bad := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		_, _ = w.Write([]byte(`{"error":"boom"}`))
	})
	up := newFailingServer(bad)
	defer up.Close()

	dir := t.TempDir()
	cfgPath := filepath.Join(dir, "config.yaml")
	cfgBody := "listen: :0\n" +
		"adapter_dir: " + filepath.Join(dir, "adapters") + "\n" +
		"plugin_dir: " + filepath.Join(dir, "plugins") + "\n" +
		"runtime_file: " + filepath.Join(dir, "runtime.json") + "\n" +
		"gateway_keys:\n  - sk-test\n"
	os.WriteFile(cfgPath, []byte(cfgBody), 0600)
	cfg, err := config.Load(cfgPath)
	if err != nil {
		t.Fatal(err)
	}
	cfg.Sources = []config.Source{{
		Name: "badup", BaseURL: up.URL, Adapter: "openai", APIKey: "sk-x",
		Models: []config.Model{{ID: "bad-model", Kind: "chat"}},
	}}
	if err := cfg.ApplyDefaults(); err != nil {
		t.Fatal(err)
	}
	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatalf("core: %v", err)
	}
	defer c.Close()
	if err := c.Plugins().LoadSource("spy", spyPlugin()); err != nil {
		t.Fatal(err)
	}
	g, err := New(c)
	if err != nil {
		t.Fatal(err)
	}

	rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
		`{"model":"bad-model","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code == http.StatusOK {
		t.Fatal("expected the upstream failure to surface as a non-200")
	}
	seen := seenStages(t, g)
	// request_start must fire; routed must NOT (no source was ever selected);
	// request_end MUST (otherwise billing misses failed traffic).
	got := stageNames(seen)
	if len(got) == 0 || got[0] != "request_start" {
		t.Fatalf("stages = %v, want request_start first", got)
	}
	hasEnd := false
	for _, s := range seen {
		if s["stage"] == "request_end" {
			hasEnd = true
			if s["ok"] != false {
				t.Error("request_end reported ok=true for a failed request")
			}
		}
	}
	if !hasEnd {
		t.Error("a FAILED request never reached request_end; a billing plugin " +
			"would not see it at all")
	}
	if contains(got, "routed") {
		t.Error("routed fired for a request that never selected a source")
	}
}

// TestBrokenPluginDoesNotBreakForwarding is the safety property, asserted
// end-to-end: a plugin that throws in EVERY stage must not stop a chat from
// being served.
func TestBrokenPluginDoesNotBreakForwarding(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	dir := t.TempDir()
	cfgPath := filepath.Join(dir, "config.yaml")
	os.WriteFile(cfgPath, []byte("listen: :0\n"+
		"adapter_dir: "+filepath.Join(dir, "adapters")+"\n"+
		"plugin_dir: "+filepath.Join(dir, "plugins")+"\n"+
		"runtime_file: "+filepath.Join(dir, "runtime.json")+"\n"+
		"gateway_keys:\n  - sk-test\n"), 0600)
	cfg, err := config.Load(cfgPath)
	if err != nil {
		t.Fatal(err)
	}
	cfg.Sources = []config.Source{{
		Name: "u", BaseURL: up.URL, Adapter: "openai", APIKey: "sk-x",
		Models: []config.Model{{ID: "m", Kind: "chat"}},
	}}
	if err := cfg.ApplyDefaults(); err != nil {
		t.Fatal(err)
	}
	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	broken := `
local plugin = { name = "broken", version = "1.0.0" }
plugin.hooks = { request_start = "f", routed = "f", request_end = "f" }
function plugin.f(p) error("plugin is broken") end
return plugin
`
	if err := c.Plugins().LoadSource("broken", broken); err != nil {
		t.Fatalf("load broken: %v", err)
	}
	g, err := New(c)
	if err != nil {
		t.Fatal(err)
	}
	rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
		`{"model":"m","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != http.StatusOK {
		t.Fatalf("a broken plugin broke forwarding: status=%d body=%s", rr.Code, rr.Body.String())
	}
	if !strings.Contains(rr.Body.String(), "pong") {
		t.Errorf("the response is not the upstream's: %s", rr.Body.String())
	}
	// And the failure must be visible rather than silent.
	if len(c.Plugins().HookErrors()) == 0 {
		t.Error("a plugin that threw in every stage left no error record")
	}
}

// TestPluginListIsServedWithStages is the operator-facing check: an admin
// debugging a plugin needs to see which stages are live and whether the hook
// file is where the kernel thinks it is.
func TestPluginListIsServedWithStages(t *testing.T) {
	g, _ := gatewayWithSpy(t)
	rr := doReq(t, g, http.MethodGet, "/api/plugins", "")
	if rr.Code != http.StatusOK {
		t.Fatalf("status=%d", rr.Code)
	}
	for _, want := range []string{"spy", "request_start", "routed", "request_end", "plugin_dir"} {
		if !strings.Contains(rr.Body.String(), want) {
			t.Errorf("/api/plugins response lacks %q", want)
		}
	}
}

func toInt(t *testing.T, v interface{}) int64 {
	t.Helper()
	f, ok := v.(float64)
	if !ok {
		return 0
	}
	return int64(f)
}

func contains(list []string, s string) bool {
	for _, x := range list {
		if x == s {
			return true
		}
	}
	return false
}

// imageUpstream serves an OpenAI-shaped image response.
func imageUpstream() *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"created":1,"data":[{"b64_json":"aGk="}]}`))
	}))
}

// TestHooksFireOnRealImageRequest: the image path is a THIRD entry point. It
// has its own handler and its own scheduler call, so "chat works" says nothing
// about it — and image traffic is billable, so a cost plugin that never sees
// it under-reports.
func TestHooksFireOnRealImageRequest(t *testing.T) {
	up := imageUpstream()
	defer up.Close()

	dir := t.TempDir()
	cfgPath := filepath.Join(dir, "config.yaml")
	os.WriteFile(cfgPath, []byte("listen: :0\n"+
		"adapter_dir: "+filepath.Join(dir, "adapters")+"\n"+
		"plugin_dir: "+filepath.Join(dir, "plugins")+"\n"+
		"runtime_file: "+filepath.Join(dir, "runtime.json")+"\n"+
		"gateway_keys:\n  - sk-test\n"), 0600)
	cfg, err := config.Load(cfgPath)
	if err != nil {
		t.Fatal(err)
	}
	cfg.Sources = []config.Source{{
		Name: "imgup", BaseURL: up.URL, Adapter: "openai", APIKey: "sk-x",
		Models: []config.Model{{ID: "img-model", Kind: "image"}},
	}}
	if err := cfg.ApplyDefaults(); err != nil {
		t.Fatal(err)
	}
	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatalf("core: %v", err)
	}
	defer c.Close()
	if err := c.Plugins().LoadSource("spy", spyPlugin()); err != nil {
		t.Fatal(err)
	}
	g, err := New(c)
	if err != nil {
		t.Fatal(err)
	}

	rr := doReq(t, g, http.MethodPost, "/v1/images/generations",
		`{"model":"img-model","prompt":"a cat"}`)
	if rr.Code != http.StatusOK {
		t.Fatalf("image status=%d body=%s", rr.Code, rr.Body.String())
	}
	seen := seenStages(t, g)
	got := stageNames(seen)
	if len(got) != 3 {
		t.Fatalf("image stages = %v, want all three (image is a separate handler)", got)
	}
	if seen[0]["type"] != "image" {
		t.Errorf("request_start type = %v, want \"image\"", seen[0]["type"])
	}
	if seen[1]["source"] != "imgup" {
		t.Errorf("routed source = %v, want imgup", seen[1]["source"])
	}
}

// TestRejectedChatStillFiresRequestStart pins the ordering decision: request_start
// fires BEFORE the quota/scope gates, so a plugin auditing traffic sees rejected
// requests too. A plugin that only ever saw served requests could not compute a
// true request rate.
func TestRejectedChatStillFiresRequestStart(t *testing.T) {
	up := mockUpstream()
	defer up.Close()
	dir := t.TempDir()
	cfgPath := filepath.Join(dir, "config.yaml")
	os.WriteFile(cfgPath, []byte("listen: :0\n"+
		"adapter_dir: "+filepath.Join(dir, "adapters")+"\n"+
		"plugin_dir: "+filepath.Join(dir, "plugins")+"\n"+
		"runtime_file: "+filepath.Join(dir, "runtime.json")+"\n"+
		"gateway_keys:\n  - sk-test\n"), 0600)
	cfg, err := config.Load(cfgPath)
	if err != nil {
		t.Fatal(err)
	}
	cfg.Sources = []config.Source{{
		Name: "u", BaseURL: up.URL, Adapter: "openai", APIKey: "sk-x",
		Models: []config.Model{{ID: "m", Kind: "chat"}},
	}}
	if err := cfg.ApplyDefaults(); err != nil {
		t.Fatal(err)
	}
	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	if err := c.Plugins().LoadSource("spy", spyPlugin()); err != nil {
		t.Fatal(err)
	}
	g, err := New(c)
	if err != nil {
		t.Fatal(err)
	}
	// A model the gateway does not know => 404, rejected before any routing.
	rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
		`{"model":"no-such-model","messages":[{"role":"user","content":"hi"}]}`)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("expected 404 for an unknown model, got %d", rr.Code)
	}
	got := stageNames(seenStages(t, g))
	if !contains(got, "request_start") {
		t.Errorf("stages = %v; a REJECTED request must still fire request_start so a "+
			"plugin can count real traffic, not just served traffic", got)
	}
}