llmgw/internal/router/pricing_test.go

package router

import (
	"crypto/sha256"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus/testutil"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// pricedRouter builds a gateway with authentication, one model under an alias,
// and a price table that charges that model and almost nothing for anything
// else — so a request priced by the wrong name is visible in the figure rather
// than merely plausible.
func pricedRouter(t *testing.T, engineURL string, served ...string) (*Router, string) {
	t.Helper()

	cfg := baseConfig(config.Endpoint{
		ID: "a", URL: engineURL,
		Models: append([]string{"qwen2.5-14b-instruct"}, served...),
	})
	cfg.Models = []catalog.Model{{
		Name: "qwen2.5-14b-instruct", Aliases: []string{"q14"},
	}}
	cfg.Auth.Enabled = true
	cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db")
	cfg.Auth.Pricing = map[string]config.Price{
		"qwen2.5-14b-instruct": {PromptPer1k: 14000, CompletionPer1k: 42000},
	}
	// Deliberately not zero: a fallback of nothing would make a miss look the
	// same as a model that is free.
	cfg.Auth.DefaultPrice = config.Price{PromptPer1k: 1, CompletionPer1k: 1}
	cfg.Auth.FlushInterval = 20 * time.Millisecond
	cfg.Auth.BatchSize = 1

	rt := newTestRouter(t, cfg)

	org, err := rt.store.CreateOrg(t.Context(), "acme")
	if err != nil {
		t.Fatal(err)
	}
	team, err := rt.store.CreateTeam(t.Context(), org, "platform")
	if err != nil {
		t.Fatal(err)
	}
	project, err := rt.store.CreateProject(t.Context(), team, "gateway")
	if err != nil {
		t.Fatal(err)
	}

	// A real key shape: authentication rejects anything else before it looks.
	secret := store.KeyPrefix + strings.Repeat("a", 43)
	sum := sha256.Sum256([]byte(secret))
	if err := rt.store.CreateKey(t.Context(), store.Key{
		ID: "key-1", Hash: sum[:], Prefix: "aaaa", Name: "test",
		ProjectID: project, CreatedAt: time.Now(),
	}); err != nil {
		t.Fatal(err)
	}
	return rt, secret
}

func TestAnAliasIsPricedAsTheModelItResolvesTo(t *testing.T) {
	// The price table is written against the catalog, and a caller asks for
	// whichever name they were given. Charging by the string they typed meant
	// every alias missed the table and fell to the default — so a budget could
	// be spent at the wrong rate, or not at all, depending on which name went
	// into the request.
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)

	rec := post(t, rt.Handler(),
		map[string]any{"model": "q14", "messages": []any{
			map[string]any{"role": "user", "content": "hello"}}},
		map[string]string{"Authorization": "Bearer " + secret})
	if rec.Code != 200 {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}

	u := waitForUsage(t, rt)
	want := charge(u.Prompt, 14000) + charge(u.Completion, 42000)
	fallback := charge(u.Prompt, 1) + charge(u.Completion, 1)
	if u.CostMicros != want {
		t.Errorf("charged %d micros for %d+%d tokens, expected %d; the "+
			"fallback rate would have charged %d",
			u.CostMicros, u.Prompt, u.Completion, want, fallback)
	}
}

// charge is Price.Cost for one half of a request, rounding up the way the
// accountant does — spelled out here so the expected figure is arithmetic the
// test states rather than the code under test restated.
func charge(tokens, per1k int64) int64 {
	if tokens <= 0 {
		return 0
	}
	return (tokens*per1k + 999) / 1000
}

func TestAModelWithNoPriceFallsBackRatherThanCostingNothing(t *testing.T) {
	e := newFakeEngine(t, "unpriced")
	rt, secret := pricedRouter(t, e.srv.URL, "unpriced")

	// Not in the catalog and not in the price table. It must still cost the
	// fallback, because a model nobody priced is the one most likely to be new.
	rec := post(t, rt.Handler(),
		map[string]any{"model": "unpriced", "messages": []any{
			map[string]any{"role": "user", "content": "hello"}}},
		map[string]string{"Authorization": "Bearer " + secret})
	if rec.Code != 200 {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}

	u := waitForUsage(t, rt)
	want := charge(u.Prompt, 1) + charge(u.Completion, 1)
	if u.CostMicros != want {
		t.Errorf("charged %d micros for %d+%d tokens, expected the fallback's %d",
			u.CostMicros, u.Prompt, u.Completion, want)
	}
	if u.CostMicros == 0 {
		t.Error("an unpriced model cost nothing; the fallback did not apply")
	}
}

func waitForUsage(t *testing.T, rt *Router) store.KeyUsage {
	t.Helper()
	since := time.Now().Add(-time.Hour)
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		rows, err := rt.store.UsageSince(t.Context(), since)
		if err == nil && len(rows) > 0 {
			return rows[0]
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("no usage was recorded")
	return store.KeyUsage{}
}

func TestARequestCanBeFollowedByItsID(t *testing.T) {
	// "Why was my request slow?" had no answer: the queue time went out in a
	// header and nowhere else, and nothing tied a request to its replica.
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)

	rec := post(t, rt.Handler(),
		map[string]any{"model": "q14", "messages": []any{
			map[string]any{"role": "user", "content": "hello"}}},
		map[string]string{"Authorization": "Bearer " + secret, "X-Request-Id": "client-abc.1"})
	if rec.Code != 200 {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	if got := rec.Header().Get("X-Request-Id"); got != "client-abc.1" {
		t.Errorf("the client's id came back as %q", got)
	}

	deadline := time.Now().Add(5 * time.Second)
	for {
		u, ok, err := rt.store.UsageByRequest(t.Context(), "client-abc.1")
		if err != nil {
			t.Fatal(err)
		}
		if ok {
			if u.Replica == "" || u.Status != 200 {
				t.Errorf("recorded %+v; expected the replica that answered", u)
			}
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("the request's usage cannot be found by its id")
		}
		time.Sleep(20 * time.Millisecond)
	}
}

func TestAnIDThatIsNotOneIsReplaced(t *testing.T) {
	for _, bad := range []string{"", "has space", "new\nline", strings.Repeat("x", 200)} {
		got := requestID(bad)
		if got == bad || !strings.HasPrefix(got, "req_") {
			t.Errorf("%q became %q", bad, got)
		}
	}
	if requestID("req_1") != "req_1" {
		t.Error("a plain id was not kept")
	}
}

func TestARequestALimitRefusesIsCountedAgainstItsOwner(t *testing.T) {
	// Refused requests were counted only in Prometheus, by model: a team
	// running into its limit looked, on every usage page, like a team that
	// had gone quiet.
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	rpm := 1
	if err := rt.store.PutLimits(t.Context(), store.Limits{Scope: store.ScopeKey, ScopeID: "key-1", RPM: &rpm}); err != nil {
		t.Fatal(err)
	}
	body := map[string]any{"model": "q14", "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
	auth := map[string]string{"Authorization": "Bearer " + secret}
	codes := []int{post(t, rt.Handler(), body, auth).Code, post(t, rt.Handler(), body, auth).Code}
	if codes[1] != 429 {
		t.Fatalf("the second request in the minute answered %v, want 429", codes)
	}
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		stats, err := rt.store.StatsBy(t.Context(), store.GroupByProject, time.Now().Add(-time.Hour))
		if err != nil {
			t.Fatal(err)
		}
		for _, s := range stats {
			if s.Limited == 1 && s.Requests == 2 {
				spend, _ := rt.store.SpendSince(t.Context(), time.Now().Add(-time.Hour))
				if spend["key-1"] == 0 {
					t.Error("the request that was served cost nothing")
				}
				return
			}
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("the refused request was never counted against its project")
}

func TestTenantsAreCountedAndTheirBudgetsGauged(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	ctx := t.Context()
	orgs, _ := rt.store.AllOrgs(ctx)
	teams, _ := rt.store.AllTeams(ctx)
	budget := int64(1_000_000)
	if err := rt.store.PutLimits(ctx, store.Limits{Scope: store.ScopeTeam, ScopeID: teams[0].ID,
		Budget: &budget, BudgetPeriod: store.PeriodMonth}); err != nil {
		t.Fatal(err)
	}
	if err := rt.m.tenant.refresh(ctx, rt.store, time.Now()); err != nil {
		t.Fatal(err)
	}
	rec := post(t, rt.Handler(), map[string]any{"model": "q14", "messages": []any{
		map[string]any{"role": "user", "content": "hello"}}},
		map[string]string{"Authorization": "Bearer " + secret})
	if rec.Code != 200 {
		t.Fatalf("status %d", rec.Code)
	}
	tenant := orgs[0].Name + "/" + teams[0].Name
	if got := testutil.ToFloat64(rt.m.tenant.requests.WithLabelValues(tenant, "ok")); got != 1 {
		t.Errorf("%s has %v ok requests, want 1", tenant, got)
	}

	u := waitForUsage(t, rt)
	if err := rt.m.tenant.refresh(ctx, rt.store, time.Now()); err != nil {
		t.Fatal(err)
	}
	want := float64(u.CostMicros) / float64(budget)
	got := testutil.ToFloat64(rt.m.tenant.budget.WithLabelValues("team", teams[0].ID, tenant, "month"))
	if want == 0 || got != want {
		t.Errorf("the team's budget gauge reads %v, want %v", got, want)
	}
}