llmgw/internal/accounting/accounting_test.go

package accounting

import (
	"context"
	"errors"
	"log/slog"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
	"gitea.mynodes.xyz/ilker/llmgw/internal/notify/notifytest"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

func intp(v int) *int     { return &v }
func i64p(v int64) *int64 { return &v }

func principal(id string, tune func(*store.Effective)) auth.Principal {
	e := store.Effective{KeyID: id, ProjectID: "proj"}
	if tune != nil {
		tune(&e)
	}
	return auth.Principal{KeyID: id, ProjectID: "proj", Limits: e}
}

func pricing(t *testing.T) *Pricing {
	t.Helper()
	return NewPricing(map[string]Price{
		"paid": {PromptPer1k: 1000, CompletionPer1k: 3000},
		"free": {},
	}, Price{})
}

func TestCostRoundsUpSoSmallRequestsAreNotFree(t *testing.T) {
	// Rounding down would make a stream of tiny requests cost nothing, which is
	// exactly the traffic shape a quota exists to bound.
	p := Price{PromptPer1k: 1000, CompletionPer1k: 3000}
	if got := p.Cost(1, 0, 0, 0, 0); got != 1 {
		t.Errorf("one prompt token cost %d", got)
	}
	if got := p.Cost(1000, 0, 1000, 0, 0); got != 4000 {
		t.Errorf("a thousand of each cost %d, want 4000", got)
	}
	if got := p.Cost(0, 0, 0, 0, 0); got != 0 {
		t.Errorf("an empty request cost %d", got)
	}
}

func TestACachedPromptCostsTheCachedRate(t *testing.T) {
	p := Price{PromptPer1k: 1000, CachedPer1k: 100, CompletionPer1k: 3000}
	// 10,000 prompt tokens of which 9,000 were read from the cache: 1,000 at
	// the prompt rate and 9,000 at the cached one.
	if got, want := p.Cost(10000, 9000, 0, 0, 0), int64(1000+900); got != want {
		t.Errorf("a 90%% cached prompt cost %d, want %d", got, want)
	}
	// More cached than prompt is an engine's miscount, not a refund.
	if got := p.Cost(100, 500, 0, 0, 0); got != 10 {
		t.Errorf("cached beyond the prompt cost %d, want 10", got)
	}
	// No cached rate: cached tokens cost what any prompt token does.
	if got := (Price{PromptPer1k: 1000}).Cost(10000, 9000, 0, 0, 0); got != 10000 {
		t.Errorf("without a cached rate a cached prompt cost %d, want 10000", got)
	}
}

func TestAPictureCostsItsPriceOnTopOfItsPrompt(t *testing.T) {
	p := Price{PromptPer1k: 1000, PerImage: 50000}
	if got := p.Cost(20, 0, 0, 3, 0); got != 20+150000 {
		t.Errorf("three pictures from a 20-token prompt cost %d", got)
	}
	if (Price{PerImage: 1}).Zero() {
		t.Error("a model priced per picture reads as free, and would pass a spent budget")
	}
}

func TestZeroPricedModelsStayAvailableWhenTheBudgetIsSpent(t *testing.T) {
	// This is what makes budget fallback work: a key with nothing left can
	// still be served something that costs nothing.
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	p := principal("k1", func(e *store.Effective) {
		e.Budget = i64p(100)
		e.BudgetPeriod = store.PeriodMonth
	})

	res, err := a.Check(p, "paid", 10)
	if err != nil {
		t.Fatal(err)
	}
	a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 1000, CompletionTokens: 0})

	if _, err := a.Check(p, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
		t.Errorf("a spent budget still served a paid model: %v", err)
	}
	if _, err := a.Check(p, "free", 10); err != nil {
		t.Errorf("a spent budget refused a free model: %v", err)
	}
}

func TestBudgetRefusalDoesNotSuggestRetrying(t *testing.T) {
	// A budget does not refill until the period rolls over, so a Retry-After of
	// a second would be a lie that turns one client into a retry loop.
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	p := principal("k1", func(e *store.Effective) {
		e.Budget = i64p(1)
		e.BudgetPeriod = store.PeriodDay
	})
	res, _ := a.Check(p, "paid", 10)
	a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 5000})

	_, err := a.Check(p, "paid", 10)
	var le *LimitError
	if !errors.As(err, &le) {
		t.Fatalf("expected a limit error, got %v", err)
	}
	if le.RetryAfter != 0 {
		t.Errorf("a spent budget suggested retrying in %v", le.RetryAfter)
	}
}

func TestRateLimitsBind(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	p := principal("k1", func(e *store.Effective) { e.RPM = intp(3) })

	for i := range 3 {
		if _, err := a.Check(p, "free", 1); err != nil {
			t.Fatalf("request %d was refused: %v", i, err)
		}
	}
	if _, err := a.Check(p, "free", 1); !errors.Is(err, ErrRateLimited) {
		t.Errorf("the fourth request was allowed: %v", err)
	}
}

func TestTokenLimitsUseTheEstimateAndThenSettle(t *testing.T) {
	// The real token count only exists once the request has run, so the check
	// has to work on an estimate. A caller that consistently underestimates
	// must still end up bounded.
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	p := principal("k1", func(e *store.Effective) { e.TPM = intp(1000) })

	// Claims 10 tokens, actually uses 900.
	res, err := a.Check(p, "free", 10)
	if err != nil {
		t.Fatal(err)
	}
	a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 400, CompletionTokens: 500})

	// The settled figure is what the next check sees, so a second request of
	// the same shape cannot fit.
	if _, err := a.Check(p, "free", 200); !errors.Is(err, ErrTokenLimited) {
		t.Errorf("the underestimate was never settled: %v", err)
	}
}

func TestNoLimitIsNotALimitOfZero(t *testing.T) {
	// The distinction the store keeps with a nil pointer. Getting it wrong here
	// would silently forbid every request for every key with no limits set.
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	p := principal("k1", nil)

	for range 50 {
		if _, err := a.Check(p, "paid", 10_000); err != nil {
			t.Fatalf("a key with no limits was refused: %v", err)
		}
	}
}

func TestAnUnauthenticatedRequestIsNotAccounted(t *testing.T) {
	// With auth off there is no key to charge, and Check must not refuse.
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	res, err := a.Check(auth.Principal{}, "paid", 1000)
	if err != nil {
		t.Fatal(err)
	}
	a.Record(res, store.Usage{PromptTokens: 1000})
	if spent, _, _ := a.Status(auth.Principal{}); spent != 0 {
		t.Errorf("an anonymous request was charged %d", spent)
	}
}

func TestBudgetSurvivesARestart(t *testing.T) {
	// The counters live in memory; the store is what makes them survive.
	ctx := context.Background()
	path := filepath.Join(t.TempDir(), "acct.db")
	s, err := store.OpenSQLite(ctx, path)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()

	must(t, s.PutOrg(ctx, store.Org{ID: "o", Name: "o"}))
	must(t, s.PutTeam(ctx, store.Team{ID: "t", OrgID: "o", Name: "t"}))
	must(t, s.PutProject(ctx, store.Project{ID: "proj", TeamID: "t", Name: "p"}))
	_, k, _ := store.NewKey("k1", "n", "proj")
	k.CreatedAt = time.Now()
	must(t, s.CreateKey(ctx, k))

	must(t, s.WriteUsage(ctx, []store.Usage{
		{At: time.Now(), KeyID: "k1", ProjectID: "proj", Model: "paid",
			ServedModel: "paid", CostMicros: 5000, Status: 200},
	}))

	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	must(t, a.RestoreBudgets(ctx, s, store.PeriodMonth))

	p := principal("k1", func(e *store.Effective) {
		e.Budget = i64p(5000)
		e.BudgetPeriod = store.PeriodMonth
	})
	if _, err := a.Check(p, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
		t.Errorf("the budget was reset by the restart: %v", err)
	}
}

func TestCatchingUpRaisesButNeverLowers(t *testing.T) {
	// Another router's spend arrives through the store; this one's own, not
	// written yet, must not be wiped by a store that is behind.
	b := NewBudgets()
	b.Add("mine", store.PeriodMonth, 900)
	b.Add("shared", store.PeriodMonth, 100)

	b.Raise(store.PeriodMonth, map[string]int64{"mine": 400, "shared": 700, "theirs": 50})

	for key, want := range map[string]int64{"mine": 900, "shared": 700, "theirs": 50} {
		if got := b.Spent(key, store.PeriodMonth); got != want {
			t.Errorf("%s: spent %d, want %d", key, got, want)
		}
	}
}

func TestBudgetsResetAtThePeriodBoundary(t *testing.T) {
	b := NewBudgets()
	now := time.Date(2026, 1, 31, 23, 59, 0, 0, time.UTC)
	b.now = func() time.Time { return now }

	b.Add("k1", store.PeriodMonth, 900)
	if got := b.Spent("k1", store.PeriodMonth); got != 900 {
		t.Fatalf("spent %d", got)
	}

	now = now.Add(2 * time.Minute) // into February
	if got := b.Spent("k1", store.PeriodMonth); got != 0 {
		t.Errorf("the counter carried %d across the month boundary", got)
	}
}

func TestDailyAndMonthlyBoundariesDiffer(t *testing.T) {
	b := NewBudgets()
	now := time.Date(2026, 3, 15, 12, 0, 0, 0, time.UTC)
	b.now = func() time.Time { return now }

	b.Add("daily", store.PeriodDay, 100)
	b.Add("monthly", store.PeriodMonth, 100)

	now = now.Add(24 * time.Hour)
	if got := b.Spent("daily", store.PeriodDay); got != 0 {
		t.Errorf("the daily counter carried %d across a day", got)
	}
	if got := b.Spent("monthly", store.PeriodMonth); got != 100 {
		t.Errorf("the monthly counter lost %d within a month", 100-got)
	}
}

func TestTheWindowDoesNotAllowTwiceTheLimit(t *testing.T) {
	// The reason for a sliding window. With a fixed one, a caller sends its
	// whole allowance at the end of a window and again at the start of the
	// next, so the observed rate is double the limit exactly when the system is
	// least able to take it.
	l := NewLimiter()
	base := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC)
	now := base.Add(59 * time.Second)
	l.now = func() time.Time { return now }

	const limit = 60
	sent := 0
	for range limit {
		if l.Allow("k", limit, 1) {
			sent++
		}
	}
	if sent != limit {
		t.Fatalf("the first burst only managed %d", sent)
	}

	// One second later, in the next fixed window.
	now = base.Add(61 * time.Second)
	second := 0
	for range limit {
		if l.Allow("k", limit, 1) {
			second++
		}
	}
	// Some allowance returns as the old window slides out of view, but nowhere
	// near the whole of it.
	if second > limit/2 {
		t.Errorf("a second burst of %d passed immediately after %d; the window is not sliding",
			second, sent)
	}
}

func TestLimiterIsSafeUnderConcurrency(t *testing.T) {
	// Checking and charging have to be one operation, or two concurrent
	// requests both see room and both take it.
	l := NewLimiter()
	const limit = 100

	var wg sync.WaitGroup
	var mu sync.Mutex
	allowed := 0
	for range 500 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if l.Allow("k", limit, 1) {
				mu.Lock()
				allowed++
				mu.Unlock()
			}
		}()
	}
	wg.Wait()
	if allowed != limit {
		t.Errorf("%d requests passed a limit of %d", allowed, limit)
	}
}

func must(t *testing.T, err error) {
	t.Helper()
	if err != nil {
		t.Fatal(err)
	}
}

// teamMember is a key whose budget and rate limit were set on its team, which
// is the arrangement an enterprise actually asks for: a group gets a quota and
// the people in it share it.
func teamMember(keyID, teamID string, tune func(*store.Effective)) auth.Principal {
	e := store.Effective{KeyID: keyID, ProjectID: "proj", TeamID: teamID}
	owner := store.LimitOwner{Scope: store.ScopeTeam, ID: teamID}
	e.BudgetOwner, e.RPMOwner, e.TPMOwner = owner, owner, owner
	if tune != nil {
		tune(&e)
	}
	return auth.Principal{KeyID: keyID, ProjectID: "proj", Limits: e}
}

func TestATeamBudgetIsOnePoolRatherThanOnePerMember(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	budget := func(e *store.Effective) {
		e.Budget = i64p(3000)
		e.BudgetPeriod = store.PeriodMonth
	}
	first := teamMember("k1", "t1", budget)
	second := teamMember("k2", "t1", budget)

	// One member spends the lot.
	res, err := a.Check(first, "paid", 1000)
	if err != nil {
		t.Fatal(err)
	}
	a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 1000, CompletionTokens: 1000})

	// The second must now be refused. Counted per key it would have its own
	// untouched 3000 — a team budget of 3000 with ten keys would permit thirty
	// thousand, which is the bug this change exists for.
	if _, err := a.Check(second, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
		t.Errorf("the team's budget should be shared; the second member got %v", err)
	}
}

func TestATeamRateLimitIsSharedToo(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	rpm := func(e *store.Effective) { e.RPM = intp(1) }

	if _, err := a.Check(teamMember("k1", "t1", rpm), "free", 1); err != nil {
		t.Fatal(err)
	}
	if _, err := a.Check(teamMember("k2", "t1", rpm), "free", 1); !errors.Is(err, ErrRateLimited) {
		t.Error("a rate limit set on the team should be shared by its keys")
	}
}

func TestTwoTeamsDoNotShareACounter(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	rpm := func(e *store.Effective) { e.RPM = intp(1) }

	if _, err := a.Check(teamMember("k1", "t1", rpm), "free", 1); err != nil {
		t.Fatal(err)
	}
	// Pooling is the point; pooling the wrong things is the danger. One team
	// exhausting another's quota would be worse than the bug being fixed.
	if _, err := a.Check(teamMember("k2", "t2", rpm), "free", 1); err != nil {
		t.Errorf("a second team should have its own counter, got %v", err)
	}
}

func TestAKeyLimitStaysPerKey(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	perKey := func(id string) auth.Principal {
		return principal(id, func(e *store.Effective) {
			e.RPM = intp(1)
			e.RPMOwner = store.LimitOwner{Scope: store.ScopeKey, ID: id}
		})
	}

	if _, err := a.Check(perKey("k1"), "free", 1); err != nil {
		t.Fatal(err)
	}
	// The tree still has to be able to say "each of you may do this much".
	// Sharing a group pool must not become the only thing expressible.
	if _, err := a.Check(perKey("k2"), "free", 1); err != nil {
		t.Errorf("a limit owned by the key should not be shared, got %v", err)
	}
}

func TestTheDifferentLimitsCanBeOwnedByDifferentScopes(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	// A team budget with a per-key rate limit is the ordinary arrangement: the
	// group's money, the individual's throughput.
	mk := func(id string) auth.Principal {
		e := store.Effective{KeyID: id, TeamID: "t1"}
		e.Budget, e.BudgetPeriod = i64p(3000), store.PeriodMonth
		e.BudgetOwner = store.LimitOwner{Scope: store.ScopeTeam, ID: "t1"}
		e.RPM = intp(1)
		e.RPMOwner = store.LimitOwner{Scope: store.ScopeKey, ID: id}
		return auth.Principal{KeyID: id, Limits: e}
	}

	res, err := a.Check(mk("k1"), "paid", 1000)
	if err != nil {
		t.Fatal(err)
	}
	a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 1000, CompletionTokens: 1000})

	// k2's rate limit is its own and untouched; the budget is the team's and
	// gone. The budget is checked first, so that is what comes back.
	if _, err := a.Check(mk("k2"), "paid", 10); !errors.Is(err, ErrBudgetSpent) {
		t.Errorf("expected the shared budget to bite, got %v", err)
	}
}

func TestStatusReportsTheCounterThatWillRefuseYou(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	budget := func(e *store.Effective) {
		e.Budget = i64p(9000)
		e.BudgetPeriod = store.PeriodMonth
	}
	res, err := a.Check(teamMember("k1", "t1", budget), "paid", 1000)
	if err != nil {
		t.Fatal(err)
	}
	a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 1000, CompletionTokens: 1000})

	// A member who has spent nothing themselves still sees the team's spend,
	// because that is the number that will stop them. A header disagreeing with
	// the refusal it precedes is worse than no header.
	spent, _, _ := a.Status(teamMember("k2", "t1", budget))
	if spent == 0 {
		t.Error("a member should see the pool's spend, not their own")
	}
}

func TestATeamBudgetSurvivesARestart(t *testing.T) {
	// The per-key restore could not do this: it rebuilt a counter named after
	// the key, while the check reads one named after the team. The control
	// plane came back believing the team had spent nothing — a budget that
	// resets on every restart, for exactly the scope where it matters most.
	ctx := context.Background()
	s, err := store.OpenSQLite(ctx, filepath.Join(t.TempDir(), "team.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()

	must(t, s.PutOrg(ctx, store.Org{ID: "o", Name: "o"}))
	must(t, s.PutTeam(ctx, store.Team{ID: "t1", OrgID: "o", Name: "t"}))
	must(t, s.PutProject(ctx, store.Project{ID: "proj", TeamID: "t1", Name: "p"}))
	for _, id := range []string{"k1", "k2"} {
		_, k, _ := store.NewKey(id, "n", "proj")
		k.CreatedAt = time.Now()
		must(t, s.CreateKey(ctx, k))
	}

	// Two members of one team, 2500 each: neither alone exceeds a 4000 budget,
	// together they do.
	must(t, s.WriteUsage(ctx, []store.Usage{
		{At: time.Now(), KeyID: "k1", ProjectID: "proj", Model: "paid",
			ServedModel: "paid", CostMicros: 2500, Status: 200},
		{At: time.Now(), KeyID: "k2", ProjectID: "proj", Model: "paid",
			ServedModel: "paid", CostMicros: 2500, Status: 200},
	}))

	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	must(t, a.RestoreBudgets(ctx, s, store.PeriodMonth))

	p := teamMember("k1", "t1", func(e *store.Effective) {
		e.Budget = i64p(4000)
		e.BudgetPeriod = store.PeriodMonth
	})
	if _, err := a.Check(p, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
		t.Errorf("the team's spend should have been restored and pooled: %v", err)
	}
}

// inTree is a key under org/team/project, with whatever each level sets.
func inTree(keyID string, chain ...store.Limits) auth.Principal {
	e := store.Effective{KeyID: keyID, Chain: chain}
	// The collapsed values still exist for headers and the profile page; they
	// are not what enforcement reads any more.
	for _, l := range chain {
		if l.Budget != nil && (e.Budget == nil || *l.Budget < *e.Budget) {
			e.Budget, e.BudgetPeriod = l.Budget, l.BudgetPeriod
			e.BudgetOwner = store.LimitOwner{Scope: l.Scope, ID: l.ScopeID}
		}
	}
	return auth.Principal{KeyID: keyID, Limits: e}
}

func orgBudget(id string, v int64) store.Limits {
	return store.Limits{Scope: store.ScopeOrg, ScopeID: id,
		Budget: i64p(v), BudgetPeriod: store.PeriodMonth}
}

func teamBudget(id string, v int64) store.Limits {
	return store.Limits{Scope: store.ScopeTeam, ScopeID: id,
		Budget: i64p(v), BudgetPeriod: store.PeriodMonth}
}

func TestTheSumOfTeamsCannotExceedTheOrganization(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))

	// Each request costs 800. The organization allows 2000 and each of its four
	// teams allows 2000, so no team comes close to its own ceiling — but the
	// organization's fills after three. Taking only the most restrictive value,
	// the organization's number was never looked at and all four went through.
	//
	// A budget stops spending once it is reached rather than refusing the
	// request that would cross it, so the third is served and the fourth is not.
	spend := func(team string) error {
		p := inTree(team+"-k", orgBudget("o", 2000), teamBudget(team, 2000))
		res, err := a.Check(p, "paid", 1000)
		if err != nil {
			return err
		}
		a.Record(res, store.Usage{KeyID: team + "-k", PromptTokens: 200, CompletionTokens: 200})
		return nil
	}

	total := 0
	for _, team := range []string{"t1", "t2", "t3", "t4"} {
		if err := spend(team); err == nil {
			total++
		}
	}
	if total != 3 {
		t.Errorf("expected the organization to stop the fourth team, %d got through", total)
	}
}

func TestATeamStillStopsBeforeItsOrganizationDoes(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	p := func() auth.Principal {
		return inTree("k1", orgBudget("o", 100_000), teamBudget("t1", 800))
	}

	res, err := a.Check(p(), "paid", 1000)
	if err != nil {
		t.Fatal(err)
	}
	a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 200, CompletionTokens: 200})

	// The narrower pool binds first, which is what it is for. Enforcing every
	// level must not stop the tighter one working.
	if _, err := a.Check(p(), "paid", 10); !errors.Is(err, ErrBudgetSpent) {
		t.Errorf("the team's own budget should still bite, got %v", err)
	}
}

func TestSpendingFillsEveryPoolAboveIt(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	p := inTree("k1", orgBudget("o", 5000), teamBudget("t1", 5000))

	res, err := a.Check(p, "paid", 1000)
	if err != nil {
		t.Fatal(err)
	}
	a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 200, CompletionTokens: 200})

	// Settled onto both, or the organization's pool would only ever fill from
	// requests that named it directly — which is none of them.
	if got := a.budgets.Spent("org:o", store.PeriodMonth); got == 0 {
		t.Error("the organization's pool did not fill")
	}
	if got := a.budgets.Spent("team:t1", store.PeriodMonth); got == 0 {
		t.Error("the team's pool did not fill")
	}
}

func TestARefusalHandsBackWhatTheNarrowerLevelsTook(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	// The team allows plenty; the organization allows one request a minute.
	chain := []store.Limits{
		{Scope: store.ScopeOrg, ScopeID: "o", RPM: intp(1)},
		{Scope: store.ScopeTeam, ScopeID: "t1", RPM: intp(100)},
	}
	if _, err := a.Check(inTree("k1", chain...), "free", 1); err != nil {
		t.Fatal(err)
	}
	before := a.requests.Rate("team:t1")

	if _, err := a.Check(inTree("k2", chain...), "free", 1); !errors.Is(err, ErrRateLimited) {
		t.Fatalf("the organization should have refused the second, got %v", err)
	}

	// A rate limiter charges as it checks, so a refusal further along has
	// already taken from the levels before it. Left there, a team at its
	// organization's ceiling would quietly eat its own rate limit with requests
	// that were never served.
	if after := a.requests.Rate("team:t1"); after != before {
		t.Errorf("the refused request stayed charged to the team: %v -> %v", before, after)
	}
}

func TestAChainWithNoLimitsAnywhereAllowsEverything(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	for i := 0; i < 5; i++ {
		if _, err := a.Check(inTree("k1"), "paid", 1000); err != nil {
			t.Fatalf("nothing is limited, but request %d was refused: %v", i, err)
		}
	}
}

func TestStatusReportsTheFullestPoolNotTheSmallestNumber(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	chain := func(team string) auth.Principal {
		return inTree(team+"-k", orgBudget("o", 9000), teamBudget(team, 9000))
	}

	// Another team spends; this one has spent nothing of its own.
	res, err := a.Check(chain("t1"), "paid", 1000)
	if err != nil {
		t.Fatal(err)
	}
	a.Record(res, store.Usage{KeyID: "t1-k", PromptTokens: 200, CompletionTokens: 200})

	spent, _, _ := a.Status(chain("t2"))
	// The organization's pool is the one nearer its ceiling, so it is the one
	// that will refuse them. Reporting the team's nil would disagree with the
	// refusal that eventually follows it.
	if spent == 0 {
		t.Error("a member should see the fullest pool that bounds them")
	}
}

func TestAnEmptyPriceTableMeansNoBudgetCanEverBind(t *testing.T) {
	// The ordinary case for self-hosted models, and the reason a budget on a
	// team can read as a limit while binding on nothing.
	if NewPricing(nil, Price{}).AnythingPriced() {
		t.Error("nothing priced should say so")
	}
	if NewPricing(map[string]Price{"m": {}}, Price{}).AnythingPriced() {
		t.Error("a model priced at zero is not priced")
	}
	if !NewPricing(map[string]Price{"m": {PromptPer1k: 1}}, Price{}).AnythingPriced() {
		t.Error("one priced model is enough for a budget to move")
	}
	if !NewPricing(nil, Price{CompletionPer1k: 1}).AnythingPriced() {
		t.Error("a fallback price prices everything")
	}
	var nilPricing *Pricing
	if nilPricing.AnythingPriced() {
		t.Error("no price table at all is not priced")
	}
}

func TestHeadroomIsTheTightestLimitAfterThisRequest(t *testing.T) {
	// A project allowing 100 a minute under a team allowing 3: the team is
	// what will refuse, so the team is what a client is told about.
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	p := principal("k1", func(e *store.Effective) {
		e.Chain = []store.Limits{
			{Scope: store.ScopeProject, ScopeID: "proj", RPM: intp(100), TPM: intp(50000)},
			{Scope: store.ScopeTeam, ScopeID: "team", RPM: intp(3)},
		}
	})
	if _, err := a.Check(p, "m", 10); err != nil {
		t.Fatal(err)
	}
	h := a.Headroom(p)
	if h.RequestLimit != 3 || h.RequestsLeft != 2 {
		t.Errorf("requests %d left of %d; expected 2 of the team's 3", h.RequestsLeft, h.RequestLimit)
	}
	if h.TokenLimit != 50000 || h.TokensLeft != 50000-10 {
		t.Errorf("tokens %d left of %d", h.TokensLeft, h.TokenLimit)
	}
	if h.Reset <= 0 || h.Reset > time.Minute {
		t.Errorf("reset %v", h.Reset)
	}
	if none := a.Headroom(principal("k2", nil)); none.RequestLimit != 0 || none.TokenLimit != 0 {
		t.Errorf("a key with no limits advertised some: %+v", none)
	}
}

func TestAFallbackIsChargedWhatTheModelThatAnsweredCosts(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	p := inTree("fb-k", teamBudget("fb", 1_000_000))
	res, err := a.Check(p, "paid", 400)
	if err != nil {
		t.Fatal(err)
	}
	// Asked for "paid", answered by "free": nothing is owed.
	a.Record(res, store.Usage{KeyID: "fb-k", Model: "paid", ServedModel: "free", PromptTokens: 200, CompletionTokens: 200})
	if spent, _, _ := a.Status(p); spent != 0 {
		t.Errorf("charged %d for a request the free model answered", spent)
	}
	res, _ = a.Check(p, "paid", 400)
	a.Record(res, store.Usage{KeyID: "fb-k", Model: "paid", ServedModel: "paid", PromptTokens: 200, CompletionTokens: 200})
	if spent, _, _ := a.Status(p); spent == 0 {
		t.Error("the paid model answering cost nothing")
	}
}

func TestASpentBudgetTellsSomeoneOnce(t *testing.T) {
	a := New(pricing(t), nil, slog.New(slog.DiscardHandler))
	n, sent := notifytest.Catch(t)
	a.SetNotifier(n)
	p := principal("k1", func(e *store.Effective) {
		e.Budget = i64p(100)
		e.BudgetPeriod = store.PeriodMonth
	})
	res, err := a.Check(p, "paid", 10)
	if err != nil {
		t.Fatal(err)
	}
	a.Record(res, store.Usage{KeyID: "k1", PromptTokens: 1000})
	for range 3 {
		if _, err := a.Check(p, "paid", 10); !errors.Is(err, ErrBudgetSpent) {
			t.Fatalf("not refused: %v", err)
		}
	}
	got := sent()
	if len(got) != 1 || got[0].Kind != notify.BudgetSpent || !strings.Contains(got[0].Message, "month") {
		t.Fatalf("sent %+v, want one budget_spent", got)
	}
}