diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index 30154ba..ff8c8e0 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -386,6 +386,19 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) { if model == "" { model = g.core.DefaultModel() } + // request_start fires for EVERY chat request, on both the AUTO and the + // direct path, and it fires BEFORE the quota / model-scope gates on + // purpose: a plugin that counts volume or audits traffic must also see the + // requests the gateway rejected, otherwise "requests accepted" would be all + // it could ever report. It sits after authentication (so the key and role in + // the payload are real) and after the messages check (a body with no + // messages is not a chat request at all). + // + // Calling it here rather than inside each branch is what keeps the two paths + // honest: an earlier version called it only from the AUTO branch, so every + // direct (model-pinned) request silently skipped it. That was caught by + // TestHooksFireOnRealDirectChat, not by reading the code. + g.fireStart(r.Context(), &req, "chat", model, len(req.Messages), len(req.Tools)) if isAuto(model) { chain := g.core.AutoChain() if chain == nil || len(chain.Tiers) == 0 { @@ -415,7 +428,6 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) { Type: "chat", OK: false, } - g.fireStart(ctx, &req, "chat", "AUTO", len(req.Messages), len(req.Tools)) // quotaExhausted reports a slot whose token window has been used up; // exhausted slots are dropped from scheduling without penalty. quotaExhausted := func(sl *scheduler.Slot) bool { @@ -873,6 +885,31 @@ func (g *Gateway) fireStart(ctx context.Context, req *chatRequest, kind, model s }) } +// fireImageStart dispatches request_start for /v1/images/generations. +// +// It is a separate function rather than a call to fireStart with a nil +// chatRequest because the image body has no messages and no tools: passing +// zeroes through a struct built for chat would invite someone to read a field +// that simply does not exist on this path. +func (g *Gateway) fireImageStart(ctx context.Context, model string) { + ps := g.core.Plugins() + if ps == nil || ps.Count() == 0 { + return + } + ps.Fire(lua.StageRequestStart, map[string]interface{}{ + "stage": string(lua.StageRequestStart), + "type": "image", + "model": model, + "key": keyID(reqKey(ctx)), + "role": reqRole(ctx), + "source": "", + "stream": false, + "messages_count": 0, + "tools_count": 0, + "ts": time.Now().Unix(), + }) +} + // fireRouted dispatches the plugin routed stage once a (source, model) slot has // been selected. tier is the AUTO tier index, or -1 on the direct path, so a // plugin can tell "this came from tier 1" from "this bypassed the chain". @@ -1208,6 +1245,11 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) { if model == "" { model = g.core.DefaultModel() } + // Same rule as the chat path, and for the same reason: an image request is + // billable traffic, so a cost plugin must see it. It fires before the + // quota/scope gates so rejected image requests are visible too. + // messages_count/tools_count are 0: the image request has neither. + g.fireImageStart(r.Context(), model) if isAuto(model) { if chain := g.core.AutoImageChain(); chain != nil && len(chain.Tiers) > 0 { if q := g.checkQuota(r.Context(), "AUTO"); q != nil { @@ -1231,6 +1273,7 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) { if usedModel != "" { rec.Model = usedModel // actual image model served, not "AUTO" } + g.fireRouted(r.Context(), "image", usedSrc, rec.Model, -2, false) rec.OK = true rec.Status = http.StatusOK // Image generation has no token concept. Recording len(ImageData) @@ -1283,6 +1326,7 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) { if resp.Model != "" { rec.Model = resp.Model // record the actual model served, not the raw request id } + g.fireRouted(r.Context(), "image", rec.Source, rec.Model, -1, false) rec.OK = true rec.Status = http.StatusOK // Image generation has no token concept — see the AUTO path above. diff --git a/internal/gateway/plugin_wiring_test.go b/internal/gateway/plugin_wiring_test.go new file mode 100644 index 0000000..306dea0 --- /dev/null +++ b/internal/gateway/plugin_wiring_test.go @@ -0,0 +1,502 @@ +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) + } +}