llmgw/test/e2e/accounting_test.go

package e2e

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"testing"
	"time"

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

// paidConfig turns on auth with a price attached to the test model, so spend
// is a number rather than always zero.
func paidConfig(path string) func(*config.Server) {
	return func(c *config.Server) {
		c.Auth = config.Auth{
			Enabled: true, Store: path,
			CacheTTL:      50 * time.Millisecond,
			FlushInterval: 20 * time.Millisecond,
			BatchSize:     8,
			Pricing: map[string]config.Price{
				testModel: {PromptPer1k: 1_000_000, CompletionPer1k: 1_000_000},
			},
		}
	}
}

func usageRows(t *testing.T, path string) []store.Usage {
	t.Helper()
	s, err := store.OpenSQLite(context.Background(), path)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()

	rows, err := s.SpendSince(context.Background(), time.Now().Add(-time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	out := make([]store.Usage, 0, len(rows))
	for key, micros := range rows {
		out = append(out, store.Usage{KeyID: key, CostMicros: micros})
	}
	return out
}

func TestUsageReachesTheStoreWithoutBlockingRequests(t *testing.T) {
	path, mint := provision(t)
	key := mint(t, "spender")
	h := start(t, 2, paidConfig(path))

	for range 12 {
		if code, msg := h.call(t, key, testModel); code != http.StatusOK {
			t.Fatalf("%d %s", code, msg)
		}
	}

	// The writer is asynchronous, so this is the one place a wait is correct.
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if rows := usageRows(t, path); len(rows) > 0 && rows[0].CostMicros > 0 {
			return
		}
		time.Sleep(25 * time.Millisecond)
	}
	t.Fatal("no usage reached the store")
}

func TestRequestRateLimitsBindThroughTheGateway(t *testing.T) {
	path, mint := provision(t)
	key := mint(t, "burst")
	withStore(t, path, func(s *store.DB) {
		rpm := 5
		if err := s.PutLimits(context.Background(), store.Limits{
			Scope: store.ScopeKey, ScopeID: "key-burst", RPM: &rpm,
		}); err != nil {
			t.Fatal(err)
		}
	})

	h := start(t, 2, paidConfig(path))

	ok, limited := 0, 0
	for range 20 {
		switch code, _ := h.call(t, key, testModel); code {
		case http.StatusOK:
			ok++
		case http.StatusTooManyRequests:
			limited++
		default:
			t.Fatalf("unexpected status %d", code)
		}
	}
	if ok != 5 {
		t.Errorf("a limit of 5 let %d requests through", ok)
	}
	if limited != 15 {
		t.Errorf("%d requests were refused, expected 15", limited)
	}
}

func TestASpentBudgetRefusesAndDoesNotSuggestRetrying(t *testing.T) {
	path, mint := provision(t)
	key := mint(t, "broke")
	withStore(t, path, func(s *store.DB) {
		// One micro-unit: the first request exhausts it.
		budget := int64(1)
		if err := s.PutLimits(context.Background(), store.Limits{
			Scope: store.ScopeKey, ScopeID: "key-broke",
			Budget: &budget, BudgetPeriod: store.PeriodMonth,
		}); err != nil {
			t.Fatal(err)
		}
	})

	h := start(t, 2, paidConfig(path))

	if code, msg := h.call(t, key, testModel); code != http.StatusOK {
		t.Fatalf("the first request should go through: %d %s", code, msg)
	}

	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if code, _ := h.call(t, key, testModel); code == http.StatusTooManyRequests {
			return
		}
		time.Sleep(25 * time.Millisecond)
	}
	t.Fatal("a key kept spending past its budget")
}

func TestBudgetSurvivesAGatewayRestart(t *testing.T) {
	// The counters live in memory. If a restart reset them, a budget would be a
	// suggestion that anyone could clear by bouncing the process.
	path, mint := provision(t)
	key := mint(t, "persistent")

	h := start(t, 2, paidConfig(path))
	for range 3 {
		if code, msg := h.call(t, key, testModel); code != http.StatusOK {
			t.Fatalf("%d %s", code, msg)
		}
	}
	// Let the batch writer land before the process goes away.
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if rows := usageRows(t, path); len(rows) > 0 && rows[0].CostMicros > 0 {
			break
		}
		time.Sleep(25 * time.Millisecond)
	}

	spent := usageRows(t, path)
	if len(spent) == 0 || spent[0].CostMicros == 0 {
		t.Fatal("nothing was recorded to restore from")
	}

	// A second gateway over the same store, with a budget just under what was
	// already spent: it must refuse immediately rather than starting fresh.
	withStore(t, path, func(s *store.DB) {
		budget := spent[0].CostMicros
		if err := s.PutLimits(context.Background(), store.Limits{
			Scope: store.ScopeKey, ScopeID: "key-persistent",
			Budget: &budget, BudgetPeriod: store.PeriodMonth,
		}); err != nil {
			t.Fatal(err)
		}
	})

	h2 := start(t, 2, paidConfig(path))
	if code, _ := h2.call(t, key, testModel); code != http.StatusTooManyRequests {
		t.Errorf("the restarted gateway gave the key a fresh budget: %d", code)
	}
}

func TestLimitsAreNotAppliedWithoutAuthentication(t *testing.T) {
	// Accounting only exists where there is an identity to charge.
	h := start(t, 2)
	for range 30 {
		if code, msg := h.call(t, "", testModel); code != http.StatusOK {
			t.Fatalf("an unauthenticated gateway rate-limited a request: %d %s", code, msg)
		}
	}
}

const cheapModel = "cheap-model"

// budgetFallbackConfig prices the main model and leaves the fallback free,
// which is what makes it affordable to a key with nothing left.
func budgetFallbackConfig(path string) func(*config.Server) {
	return func(c *config.Server) {
		paidConfig(path)(c)
		// The fallback has no price, so it costs nothing and the budget check
		// exempts it. That exemption has existed all along with a comment
		// saying it is what lets a spent key still be served something cheap;
		// nothing ever tried, so it could never fire.
		c.Models = budgetFallbackModels()
		// The same engines answer for both names, so this is about accounting
		// rather than about placement.
		for i := range c.Endpoints {
			c.Endpoints[i].Models = []string{testModel, cheapModel}
		}
	}
}

// callWithHeaders is call, keeping the response headers.
func (h *harness) callWithHeaders(t *testing.T, key, model string) (int, http.Header) {
	t.Helper()
	body, _ := json.Marshal(map[string]any{
		"model":      model,
		"messages":   []map[string]string{{"role": "user", "content": "hi"}},
		"max_tokens": 4,
	})
	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
		h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+key)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, resp.Body)
	return resp.StatusCode, resp.Header
}

func TestASpentBudgetIsServedAnAffordableModel(t *testing.T) {
	// A team trying models will exhaust a quota, and being dropped to a smaller
	// model is a far better answer than a refusal: it makes a budget a
	// guardrail rather than a wall.
	path, mint := provision(t)
	key := mint(t, "broke")
	withStore(t, path, func(s *store.DB) {
		budget := int64(1) // one micro-unit: the first request exhausts it
		if err := s.PutLimits(context.Background(), store.Limits{
			Scope: store.ScopeKey, ScopeID: "key-broke",
			Budget: &budget, BudgetPeriod: store.PeriodMonth,
		}); err != nil {
			t.Fatal(err)
		}
	})

	h := start(t, 2, budgetFallbackConfig(path))

	if code, msg := h.call(t, key, testModel); code != http.StatusOK {
		t.Fatalf("the first request should go through: %d %s", code, msg)
	}

	// Once the spend lands, the priced model is unaffordable and the free one
	// answers instead — with the response saying which, rather than leaving
	// the caller to infer it from the reply.
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		code, hdr := h.callWithHeaders(t, key, testModel)
		if code == http.StatusOK && hdr.Get("x-llmgw-served-model") == cheapModel {
			return
		}
		if code == http.StatusTooManyRequests {
			t.Fatal("a spent budget refused the request instead of falling back")
		}
		time.Sleep(25 * time.Millisecond)
	}
	t.Fatal("the budget never ran out, or the fallback never served")
}

func TestASpentBudgetWithNoFallbackStillRefuses(t *testing.T) {
	// The limit has to remain a limit. A model with nowhere cheaper to go is
	// refused exactly as before, or a budget would stop meaning anything.
	path, mint := provision(t)
	key := mint(t, "broke")
	withStore(t, path, func(s *store.DB) {
		budget := int64(1)
		if err := s.PutLimits(context.Background(), store.Limits{
			Scope: store.ScopeKey, ScopeID: "key-broke",
			Budget: &budget, BudgetPeriod: store.PeriodMonth,
		}); err != nil {
			t.Fatal(err)
		}
	})

	h := start(t, 2, paidConfig(path)) // no catalog, so no fallbacks

	if code, msg := h.call(t, key, testModel); code != http.StatusOK {
		t.Fatalf("the first request should go through: %d %s", code, msg)
	}
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if code, _ := h.call(t, key, testModel); code == http.StatusTooManyRequests {
			return
		}
		time.Sleep(25 * time.Millisecond)
	}
	t.Fatal("a key with no fallback kept spending past its budget")
}

// budgetFallbackModels is a priced model with a free one behind it.
func budgetFallbackModels() []catalog.Model {
	return []catalog.Model{
		{
			Name: testModel, Source: "/models/test",
			Fallbacks: []string{cheapModel},
			Scaling:   catalog.Scaling{Min: 1, Max: 1},
			Profiles:  []catalog.Profile{{Name: "p", Engine: "llamacpp", Weight: 1}},
		},
		{
			Name: cheapModel, Source: "/models/cheap",
			Scaling:  catalog.Scaling{Min: 1, Max: 1},
			Profiles: []catalog.Profile{{Name: "p", Engine: "llamacpp", Weight: 1}},
		},
	}
}