ModelRouter/internal/gateway/plugins_api_test.go

package gateway

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

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

// gatewayWithBilling boots a gateway with the bundled billing plugin loaded, so
// the UI-injection endpoint is exercised against a real plugin rather than a
// hand-written stub. The other plugin tests in this package assert on the
// WebUI source; this one asserts on the HTTP contract the browser consumes.
func gatewayWithBilling(t *testing.T) *Gateway {
	t.Helper()
	dir := t.TempDir()
	cfgPath := filepath.Join(dir, "config.yaml")
	body := "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(body), 0600); err != nil {
		t.Fatal(err)
	}
	cfg, err := config.Load(cfgPath)
	if err != nil {
		t.Fatal(err)
	}
	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatalf("core: %v", err)
	}
	t.Cleanup(c.Close)
	// Load the shipped plugin explicitly: seeding only runs for a directory that
	// does not exist yet, and this test wants a known plugin regardless.
	src, err := lua.ReadBundledPlugin("billing")
	if err != nil {
		t.Fatalf("read bundled billing: %v", err)
	}
	if err := c.Plugins().LoadSource("billing", src); err != nil {
		t.Fatalf("load billing: %v", err)
	}
	g, err := New(c)
	if err != nil {
		t.Fatalf("gateway: %v", err)
	}
	return g
}

// TestUIInjectServesPluginUI: GET /api/ui-inject is the single call the WebUI
// makes at boot, and it must carry BOTH a contributed page and contributed
// elements — the browser builds the sidebar from the page and mounts the
// elements into existing panes from the same payload, so a partial response
// would produce a page with no body or a missing tile.
func TestUIInjectServesPluginUI(t *testing.T) {
	g := gatewayWithBilling(t)
	rr := doReq(t, g, http.MethodGet, "/api/ui-inject", "")
	if rr.Code != http.StatusOK {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	var view struct {
		// Decoded into the real types so the test cannot drift from the wire
		// contract. An inline copy missed the `pages` field when the payload
		// shape changed and failed to compile, which is at least loud — but the
		// same copy also went on asserting the OLD single-page shape for a
		// release, silently.
		UI     lua.UIExtension `json:"ui"`
		Stages []string        `json:"stages"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &view); err != nil {
		t.Fatalf("decode: %v", err)
	}
	// The payload carries every page in one list; ui.page is no longer a
	// separate slot (it was, and a second plugin contributing a page overwrote
	// whatever was there).
	var billing *lua.UIPage
	for i, pg := range view.UI.Pages {
		if pg != nil && pg.PageID == "billing" {
			billing = view.UI.Pages[i]
		}
	}
	if billing == nil {
		t.Fatalf("no billing page in the inject payload; pages=%d", len(view.UI.Pages))
	}
	if !strings.Contains(billing.Mount, "billing-root") {
		t.Error("the page mount came back empty")
	}
	if len(view.UI.Elements) == 0 {
		t.Error("billing contributes an element to the status page but it is missing")
	}
	for _, e := range view.UI.Elements {
		if e.Target != "status" {
			t.Errorf("element target = %q, want \"status\"", e.Target)
		}
	}
	// Derived from AllStages, not hardcoded: the previous assertion of "3"
	// is exactly what let chain_step go missing from this payload unnoticed.
	if len(view.Stages) != len(lua.AllStages) {
		t.Errorf("stages = %v, want %d (one per AllStages entry)", view.Stages, len(lua.AllStages))
	}
	for i, st := range lua.AllStages {
		if i >= len(view.Stages) || view.Stages[i] != string(st) {
			t.Errorf("stages[%d] = %v, want %q", i, view.Stages, string(st))
		}
	}
	if !containsStr(view.Stages, string(lua.StageChainStep)) {
		t.Error("the discovery payload does not advertise chain_step; a plugin " +
			"author would conclude the stage does not exist")
	}
}

// TestUIInjectIsEmptyWithoutPlugins: a gateway with no plugins must still answer
// 200 with an empty (not missing, not null) payload. The WebUI calls this
// unconditionally at boot, so a 404 or a null `ui` would break every dashboard.
func TestUIInjectIsEmptyWithoutPlugins(t *testing.T) {
	g := newTestGateway(t)
	rr := doReq(t, g, http.MethodGet, "/api/ui-inject", "")
	if rr.Code != http.StatusOK {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	if !strings.Contains(rr.Body.String(), `"ui"`) {
		t.Error("no ui key in the response; the WebUI would have nothing to read")
	}
}

// containsStr reports whether list has s.
func containsStr(list []string, s string) bool {
	for _, x := range list {
		if x == s {
			return true
		}
	}
	return false
}

// newRecorderFor pushes a request through the full handler chain.
func newRecorderFor(t *testing.T, g *Gateway, req *http.Request) *httptest.ResponseRecorder {
	t.Helper()
	rr := httptest.NewRecorder()
	g.Handler().ServeHTTP(rr, req)
	return rr
}

// TestPluginsListAndStateAPI covers the management surface the plugin docs
// promise: listing, and reading a plugin's own published state.
func TestPluginsListAndStateAPI(t *testing.T) {
	g := gatewayWithBilling(t)
	rr := doReq(t, g, http.MethodGet, "/api/plugins", "")
	if rr.Code != http.StatusOK {
		t.Fatalf("GET /api/plugins = %d: %s", rr.Code, rr.Body.String())
	}
	for _, want := range []string{"billing", "hook_errors", "plugin_dir", "request_end"} {
		if !strings.Contains(rr.Body.String(), want) {
			t.Errorf("/api/plugins response lacks %q", want)
		}
	}
	// state read: the plugin published its (empty) state, so the key exists.
	rr = doReq(t, g, http.MethodGet, "/api/plugins/billing/state", "")
	if rr.Code != http.StatusOK {
		t.Fatalf("GET state = %d: %s", rr.Code, rr.Body.String())
	}
	if !strings.Contains(rr.Body.String(), `"total"`) {
		t.Errorf("billing state lacks the total bucket: %s", rr.Body.String())
	}
}

// TestPluginStatePUTIsAdminOnly: only the WRITE side is gated. A user key must
// be able to READ its own billing widget's data, but must not be able to
// rewrite the price table.
func TestPluginStatePUTIsAdminOnly(t *testing.T) {
	g := gatewayWithBilling(t)

	// Build a user-role key and remember its secret.
	rec, err := g.core.CreateKey("viewer", "user", nil, "")
	if err != nil {
		t.Fatal(err)
	}
	userKey := rec.Key

	// A user key may read the state.
	req, _ := http.NewRequest(http.MethodGet, "/api/plugins/billing/state", nil)
	req.Header.Set("Authorization", "Bearer "+userKey)
	rr := newRecorderFor(t, g, req)
	if rr.Code != http.StatusOK {
		t.Errorf("user GET state = %d, want 200 (the billing widget must render for users)", rr.Code)
	}

	// A user key may NOT write it.
	put, _ := http.NewRequest(http.MethodPut, "/api/plugins/billing/state",
		strings.NewReader(`{"prices":{"default":{"prompt":0}}}`))
	put.Header.Set("Authorization", "Bearer "+userKey)
	put.Header.Set("Content-Type", "application/json")
	prr := newRecorderFor(t, g, put)
	if prr.Code != http.StatusForbidden {
		t.Errorf("user PUT state = %d, want 403 (a user must not rewrite the price table)", prr.Code)
	}

	// An admin key may.
	adm := doReq(t, g, http.MethodPut, "/api/plugins/billing/state",
		`{"prices":{"default":{"prompt":1e-6}}}`)
	if adm.Code != http.StatusOK {
		t.Errorf("admin PUT state = %d: %s", adm.Code, adm.Body.String())
	}
}