ModelRouter/internal/gateway/billing_rules_api_test.go

package gateway

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

	"llmsproxy/internal/config"
)

// rulesGateway needs a real billing plugin (the API injects prices into it)
// AND two sources whose base_urls the rules will match, so it cannot reuse
// either existing helper: gatewayWithBilling has no sources, and
// newTestGateway loads no plugins.
func rulesGateway(t *testing.T) *Gateway {
	t.Helper()
	g := gatewayWithBilling(t)
	cfg := g.core.Config()
	cfg.Sources = []config.Source{
		{Name: "localzen", BaseURL: "https://free.example.com/v1"},
		{Name: "commandcode", BaseURL: "https://api.commandcode.ai/v1"},
	}
	return g
}

func rulesDo(t *testing.T, g *Gateway, method, body string) *httptest.ResponseRecorder {
	t.Helper()
	return doReq(t, g, method, "/api/plugins/billing/rules", body)
}

// TestBillingRulesCRUD covers the whole editor loop against the real HTTP
// surface: read, add, replace-by-URL, delete. The replace-by-URL case is the
// one worth pinning — two rules for one URL would make the first unreachable
// (first match wins) without any error, which is exactly the kind of silent
// ambiguity an editor must not be able to create.
func TestBillingRulesCRUD(t *testing.T) {
	g := rulesGateway(t)
	seed := `{"billing":{"active":"p1","profiles":[{"id":"p1","rules":[
      {"url":"https://free.example.com/v1","mode":"free"}]}]}}`
	if rr := rulesDo(t, g, http.MethodPut, seed); rr.Code != http.StatusOK {
		t.Fatalf("seed rules = %d: %s", rr.Code, rr.Body.String())
	}

	// Add a token-priced rule for the second URL.
	add := `{"profile":"p1","rule":{"url":"https://api.commandcode.ai/v1","mode":"token",
      "models":{"deepseek-v4.1-flash":{"prompt":"0.15","completion":"0.60"}}}}`
	if rr := rulesDo(t, g, http.MethodPost, add); rr.Code != http.StatusOK {
		t.Fatalf("add rule = %d: %s", rr.Code, rr.Body.String())
	}
	var got billingRulesPayload
	if err := json.Unmarshal(rulesDo(t, g, http.MethodGet, "").Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if len(got.DSL.Profiles[0].Rules) != 2 {
		t.Fatalf("after add there are %d rules, want 2", len(got.DSL.Profiles[0].Rules))
	}

	// Adding the SAME url again must replace, not append.
	if rr := rulesDo(t, g, http.MethodPost, `{"profile":"p1","rule":{"url":"https://api.commandcode.ai/v1","mode":"free"}}`); rr.Code != http.StatusOK {
		t.Fatalf("re-add same url = %d: %s", rr.Code, rr.Body.String())
	}
	if err := json.Unmarshal(rulesDo(t, g, http.MethodGet, "").Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if n := len(got.DSL.Profiles[0].Rules); n != 2 {
		t.Fatalf("re-adding an existing url appended instead of replacing: %d rules", n)
	}
	var mode string
	for _, r := range got.DSL.Profiles[0].Rules {
		if r.URL == "https://api.commandcode.ai/v1" {
			mode = r.Mode
		}
	}
	if mode != "free" {
		t.Errorf("re-added rule mode = %q, want free (the replacement must win)", mode)
	}

	// Delete by URL.
	if rr := rulesDo(t, g, http.MethodDelete, `{"profile":"p1","url":"https://api.commandcode.ai/v1"}`); rr.Code != http.StatusOK {
		t.Fatalf("delete rule = %d: %s", rr.Code, rr.Body.String())
	}
	if err := json.Unmarshal(rulesDo(t, g, http.MethodGet, "").Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if len(got.DSL.Profiles[0].Rules) != 1 {
		t.Errorf("after delete there are %d rules, want 1", len(got.DSL.Profiles[0].Rules))
	}
}

// TestBillingRulesRejectUnknownURLWithWarning is the silent-failure guard: a
// rule whose URL matches no source prices nothing, so every request on it
// falls into unpriced. The API must say so instead of accepting it silently.
func TestBillingRulesRejectUnknownURLWithWarning(t *testing.T) {
	g := rulesGateway(t)
	if rr := rulesDo(t, g, http.MethodPut, `{"billing":{"active":"p","profiles":[{"id":"p",
      "rules":[{"url":"https://typo.example.com/v1","mode":"free"}]}]}}`); rr.Code != http.StatusOK {
		t.Fatalf("seed = %d: %s", rr.Code, rr.Body.String())
	}
	var got billingRulesPayload
	if err := json.Unmarshal(rulesDo(t, g, http.MethodGet, "").Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if len(got.Warnings) == 0 {
		t.Fatal("a rule matching no configured source was accepted with no warning — " +
			"it will price nothing and every request will be recorded as unpriced")
	}
	if !strings.Contains(got.Warnings[0], "typo.example.com") {
		t.Errorf("warning does not name the offending url: %v", got.Warnings)
	}
	// The editor needs the real URL list to offer suggestions.
	if len(got.URLs) != 2 {
		t.Errorf("urls offered to the editor = %v, want the 2 configured base_urls", got.URLs)
	}
}

// TestBillingRulesInvalidIsRejectedNotPersisted checks the write path refuses
// bad input and leaves the previous rules in place — an editor that clears the
// table on a typo loses the price table.
func TestBillingRulesInvalidIsRejectedNotPersisted(t *testing.T) {
	g := rulesGateway(t)
	seed := `{"billing":{"active":"p","profiles":[{"id":"p","rules":[
      {"url":"https://free.example.com/v1","mode":"free"}]}]}}`
	if rr := rulesDo(t, g, http.MethodPut, seed); rr.Code != http.StatusOK {
		t.Fatalf("seed = %d: %s", rr.Code, rr.Body.String())
	}
	// An unknown mode must be refused.
	if rr := rulesDo(t, g, http.MethodPost, `{"profile":"p","rule":{"url":"https://x/v1","mode":"banana"}}`); rr.Code != http.StatusBadRequest {
		t.Errorf("invalid mode accepted with %d, want 400", rr.Code)
	}
	var got billingRulesPayload
	if err := json.Unmarshal(rulesDo(t, g, http.MethodGet, "").Body.Bytes(), &got); err != nil {
		t.Fatal(err)
	}
	if len(got.DSL.Profiles[0].Rules) != 1 {
		t.Errorf("a rejected rule was persisted anyway: %d rules", len(got.DSL.Profiles[0].Rules))
	}
	// An unparseable price must be refused too, not silently priced at zero.
	if rr := rulesDo(t, g, http.MethodPost, `{"profile":"p","rule":{"url":"https://y/v1","mode":"token",
      "models":{"m":{"prompt":"free","completion":"0.6"}}}}`); rr.Code != http.StatusBadRequest {
		t.Errorf("non-numeric price accepted with %d, want 400", rr.Code)
	}
}

// TestBillingRulesPersistAndReachThePlugin is the point of the whole
// endpoint, and mutation verification showed the earlier CRUD tests missed it
// twice: dropping cfg.Save() and dropping the SetState injection both left
// every test green. Both failures are invisible in a GET — the response echoes
// the in-memory rules either way — so this test checks the two consequences
// that actually matter:
//
//  1. persistence: the rule must be in the config file on disk, because an
//     edit that evaporates on restart is a lie the UI told the operator;
//  2. effect: the billing plugin must now PRICE the source (its published
//     prices table has a non-zero entry), because a rule that is stored but
//     never injected prices nothing at all.
func TestBillingRulesPersistAndReachThePlugin(t *testing.T) {
	g := rulesGateway(t)
	cfg := g.core.Config()

	seed := `{"billing":{"active":"p","profiles":[{"id":"p","rules":[
      {"url":"https://free.example.com/v1","mode":"free"},
      {"url":"https://api.commandcode.ai/v1","mode":"token",
       "models":{"deepseek-v4.1-flash":{"prompt":"0.15","completion":"0.60"}}}]}]}}`
	if rr := rulesDo(t, g, http.MethodPut, seed); rr.Code != http.StatusOK {
		t.Fatalf("seed = %d: %s", rr.Code, rr.Body.String())
	}

	// 1. persisted to the file the process loaded from.
	raw, err := os.ReadFile(cfg.Path)
	if err != nil {
		t.Fatalf("read config: %v", err)
	}
	if !strings.Contains(string(raw), "commandcode.ai") {
		t.Errorf("the rule is not in config.yaml on disk — an edit that does not "+
			"survive a restart:\n%s", raw)
	}
	reloaded, err := config.Load(cfg.Path)
	if err != nil {
		t.Fatalf("reload saved config: %v", err)
	}
	if reloaded.BillingDSL == nil || len(reloaded.BillingDSL.Profiles) == 0 {
		t.Fatalf("saved config has no billing DSL: %s", raw)
	}
	// A round-trip through YAML must not lose the price strings.
	found := false
	for _, r := range reloaded.BillingDSL.Profiles[0].Rules {
		if strings.Contains(r.URL, "commandcode") {
			found = true
			if m, ok := r.Models["deepseek-v4.1-flash"]; !ok || m.Prompt != "0.15" {
				t.Errorf("price did not survive the YAML round-trip: %+v", m)
			}
		}
	}
	if !found {
		t.Errorf("the token rule is missing after reload:\n%s", raw)
	}

	// 2. reached the plugin. Prices live in a separate Lua field from `state`,
	// and the GET /state endpoint returns the published state together with
	// the prices sidecar — which is exactly what the plugin's own UI reads.
	state := g.core.Plugins().State("billing")
	if state == nil {
		t.Fatalf("billing plugin published no state")
	}
	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())
	}
	var envelope struct {
		Prices map[string]interface{} `json:"prices"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &envelope); err != nil {
		t.Fatalf("decode state envelope: %v", err)
	}
	prices := envelope.Prices
	if prices == nil {
		t.Fatalf("state payload carries no prices")
	}
	sources, ok := prices["sources"].(map[string]interface{})
	if !ok {
		t.Fatalf("prices has no sources map: %#v", prices)
	}
	cc, ok := sources["commandcode"]
	if !ok {
		t.Fatalf("the rule was saved but never injected: prices has %v, "+
			"no commandcode entry. Every request on it stays unpriced.", keysOf(sources))
	}
	if cc == nil {
		t.Error("commandcode price entry is nil")
	}
	// The free rule must price localzen, and the two must not be conflated.
	if _, ok := sources["localzen"]; !ok {
		t.Errorf("the free rule did not reach the plugin: %v", keysOf(sources))
	}
	// The injected price must be per-token, i.e. 0.15 USD/M => 1.5e-7.
	// Asserting the exact magnitude catches a unit error (per-million vs
	// per-token), which would be off by a factor of a million and still look
	// like "a number".
	models, ok := cc.(map[string]interface{})["models"].(map[string]interface{})
	if ok {
		if m, ok := models["deepseek-v4.1-flash"].(map[string]interface{}); ok {
			p, _ := m["prompt"].(float64)
			if p <= 0 || p > 1e-6 {
				t.Errorf("injected prompt price = %v, want per-token (~1.5e-7); "+
					"a per-million value here would over-bill by 1e6x", p)
			}
		}
	}
}

func keysOf(m map[string]interface{}) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	sort.Strings(out)
	return out
}

// TestBillingRulesRequireAdmin keeps price rewriting behind the admin role,
// like every other plugin write: prices are configuration, not user data.
func TestBillingRulesRequireAdmin(t *testing.T) {
	g := rulesGateway(t)
	rec, err := g.core.CreateKey("viewer", "user", nil, "")
	if err != nil {
		t.Fatal(err)
	}
	userKey := rec.Key
	for _, m := range []string{http.MethodPut, http.MethodPost, http.MethodDelete} {
		req, _ := http.NewRequest(m, "/api/plugins/billing/rules", strings.NewReader(`{}`))
		req.Header.Set("Authorization", "Bearer "+userKey)
		req.Header.Set("Content-Type", "application/json")
		if rr := newRecorderFor(t, g, req); rr.Code != http.StatusForbidden {
			t.Errorf("user %s rules = %d, want 403", m, rr.Code)
		}
	}
}