netbird/management/internals/modules/agentnetwork/policyselect.go

package agentnetwork

import (
	"context"
	"fmt"
	"math"
	"sort"
	"strings"
	"time"

	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
	"github.com/netbirdio/netbird/management/server/store"
	"github.com/netbirdio/netbird/shared/management/status"
)

// validateUsageDeltas rejects negative or non-finite usage counters before they
// reach the consumption store, so a bad delta can't decrement or poison totals.
// The store batch method enforces the same invariant; this is the manager-level
// guard so direct callers fail fast with a clear error.
func validateUsageDeltas(tokensIn, tokensOut int64, costUSD float64) error {
	if tokensIn < 0 || tokensOut < 0 || costUSD < 0 || math.IsNaN(costUSD) || math.IsInf(costUSD, 0) {
		return status.Errorf(status.InvalidArgument, "usage deltas must be non-negative and finite")
	}
	return nil
}

// Deny codes the proxy surfaces back to the caller when every
// applicable policy is exhausted. The proxy converts these into
// upstream-shaped error responses.
const (
	//nolint:gosec // policy deny code label, not a credential
	denyCodeTokenCapExceeded = "llm_policy.token_cap_exceeded"
	//nolint:gosec // policy deny code label, not a credential
	denyCodeBudgetCapExceeded = "llm_policy.budget_cap_exceeded"
	//nolint:gosec // account deny code label, not a credential
	denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded"
	//nolint:gosec // account deny code label, not a credential
	denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded"
	// denyCodeModelBlocked is returned when policies govern the request's
	// (provider, caller-groups) but none permits the model. Matches the proxy
	// guardrail's code so both layers surface the same label.
	denyCodeModelBlocked = "llm_policy.model_blocked"
)

// consumptionCache holds the consumption counters prefetched for one
// policy-selection request, keyed by ConsumptionKey. A miss returns a zero
// counter — the same contract the store's single-row getter uses for absent
// rows — so the eval logic is identical whether a counter exists yet or not.
type consumptionCache map[types.ConsumptionKey]*types.Consumption

func (c consumptionCache) get(accountID string, kind types.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time) *types.Consumption {
	key := types.ConsumptionKey{Kind: kind, DimID: dimID, WindowSeconds: windowSeconds, WindowStartUTC: windowStart.UTC()}
	if row, ok := c[key]; ok && row != nil {
		return row
	}
	return &types.Consumption{
		AccountID:      accountID,
		DimensionKind:  kind,
		DimensionID:    dimID,
		WindowSeconds:  windowSeconds,
		WindowStartUTC: windowStart.UTC(),
	}
}

// addLimitKeys records the user/group consumption keys a single enabled (token
// or budget) limit window reads for the given attribution group, into a dedup
// set. attrGroup may be empty (no group dimension applies).
func addLimitKeys(set map[types.ConsumptionKey]struct{}, userID, attrGroup string, windowSeconds int64, now time.Time) {
	if windowSeconds <= 0 {
		return
	}
	ws := types.WindowStart(now, windowSeconds)
	if userID != "" {
		set[types.ConsumptionKey{Kind: types.DimensionUser, DimID: userID, WindowSeconds: windowSeconds, WindowStartUTC: ws}] = struct{}{}
	}
	if attrGroup != "" {
		set[types.ConsumptionKey{Kind: types.DimensionGroup, DimID: attrGroup, WindowSeconds: windowSeconds, WindowStartUTC: ws}] = struct{}{}
	}
}

// prefetchConsumption loads, in one store round-trip, every consumption counter
// that the account-budget ceiling and the candidate policies will read while
// scoring this request. This replaces the per-cap point reads the selector
// previously issued one at a time (the N+1 on the hot path).
func (m *managerImpl) prefetchConsumption(ctx context.Context, in PolicySelectionInput, rules []*types.AccountBudgetRule, candidates []*types.Policy, now time.Time) (consumptionCache, error) {
	set := make(map[types.ConsumptionKey]struct{})
	for _, p := range candidates {
		attr := lowestIntersect(p.SourceGroups, in.GroupIDs)
		if p.Limits.TokenLimit.Enabled {
			addLimitKeys(set, in.UserID, attr, p.Limits.TokenLimit.WindowSeconds, now)
		}
		if p.Limits.BudgetLimit.Enabled {
			addLimitKeys(set, in.UserID, attr, p.Limits.BudgetLimit.WindowSeconds, now)
		}
	}
	for _, r := range rules {
		if r == nil || !r.Enabled || !budgetRuleApplies(r, in) {
			continue
		}
		attr := lowestIntersect(r.TargetGroups, in.GroupIDs)
		if r.Limits.TokenLimit.Enabled {
			addLimitKeys(set, in.UserID, attr, r.Limits.TokenLimit.WindowSeconds, now)
		}
		if r.Limits.BudgetLimit.Enabled {
			addLimitKeys(set, in.UserID, attr, r.Limits.BudgetLimit.WindowSeconds, now)
		}
	}
	if len(set) == 0 {
		return consumptionCache{}, nil
	}
	keys := make([]types.ConsumptionKey, 0, len(set))
	for k := range set {
		keys = append(keys, k)
	}
	rows, err := m.store.GetAgentNetworkConsumptionBatch(ctx, store.LockingStrengthNone, in.AccountID, keys)
	if err != nil {
		return nil, fmt.Errorf("batch read consumption: %w", err)
	}
	return consumptionCache(rows), nil
}

// SelectPolicyForRequest picks the policy that "pays" for the
// incoming request. The chosen policy is the one with the largest
// pool that still has headroom — drain the bigger bucket first,
// fall through to the next-biggest only when the current one's
// group cap or shared per-user cap is exhausted. This matches
// operator intuition for layered tiers ("privileged group has the
// 10k budget, regular group has 1k as the safety net") and avoids
// the load-balancer flapping that fraction-based scoring produces
// once any cap has been touched.
//
// Ordering across non-exhausted candidates:
//  1. Policies with NO enabled caps (catch-all-allow) win over any
//     capped policy — operators who configure unlimited access
//     expect requests to attribute there until they explicitly add
//     caps.
//  2. Larger group token cap wins.
//  3. Larger group budget USD cap wins.
//  4. Larger user token cap wins.
//  5. Larger user budget USD cap wins.
//  6. Older created_at wins (deterministic final tiebreak so
//     multi-node selection converges).
//
// Returns Allow=true with empty SelectedPolicyID when no policy in
// the account targets the (provider, caller-groups) combination —
// llm_router is the gate that owns "no policy authorises this
// request" semantics; this function trusts that authorisation has
// already happened upstream and only does the limit-aware
// attribution.
func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelectionInput) (*PolicySelectionResult, error) {
	if in.AccountID == "" {
		return nil, status.Errorf(status.InvalidArgument, "account_id is required")
	}

	now := time.Now().UTC()

	rules, err := m.store.GetAccountAgentNetworkBudgetRules(ctx, store.LockingStrengthNone, in.AccountID)
	if err != nil {
		return nil, fmt.Errorf("list account budget rules: %w", err)
	}
	policies, err := m.store.GetAccountAgentNetworkPolicies(ctx, store.LockingStrengthNone, in.AccountID)
	if err != nil {
		return nil, fmt.Errorf("list account policies: %w", err)
	}
	candidates := filterApplicablePolicies(policies, in)

	candidates, denied, err := m.applyModelGate(ctx, in, candidates)
	if err != nil {
		return nil, err
	}
	if denied != nil {
		return denied, nil
	}

	// Prefetch every consumption counter the ceiling + candidate policies will
	// read, in a single store round-trip, then score against the cache.
	cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now)
	if err != nil {
		return nil, err
	}

	// Account-level budget rules are an always-on ceiling, evaluated
	// independently of policy selection (they bind even for catch-all-allow
	// policies or requests that match no policy). All applicable rules must
	// pass — this is where min-wins lives.
	if deny, code, reason := checkAccountBudget(in, rules, cache, now); deny {
		return &PolicySelectionResult{Allow: false, DenyCode: code, DenyReason: reason}, nil
	}

	if len(candidates) == 0 {
		return &PolicySelectionResult{Allow: true}, nil
	}
	scored, lastDenyCode, lastDenyReason := scoreCandidates(in, candidates, cache, now)
	if len(scored) == 0 {
		return &PolicySelectionResult{
			Allow:      false,
			DenyCode:   lastDenyCode,
			DenyReason: lastDenyReason,
		}, nil
	}

	sort.SliceStable(scored, func(i, j int) bool {
		// Catch-all-allow (no caps configured) wins outright over
		// any capped policy.
		iNoCap := isUncapped(scored[i].policy)
		jNoCap := isUncapped(scored[j].policy)
		if iNoCap != jNoCap {
			return iNoCap
		}
		// Bigger pool drains first. Group caps dominate (shared
		// across the group) before individual caps.
		if a, b := groupCapTokens(scored[i].policy), groupCapTokens(scored[j].policy); a != b {
			return a > b
		}
		if a, b := groupCapBudgetUsd(scored[i].policy), groupCapBudgetUsd(scored[j].policy); a != b {
			return a > b
		}
		if a, b := userCapTokens(scored[i].policy), userCapTokens(scored[j].policy); a != b {
			return a > b
		}
		if a, b := userCapBudgetUsd(scored[i].policy), userCapBudgetUsd(scored[j].policy); a != b {
			return a > b
		}
		return scored[i].policy.CreatedAt.Before(scored[j].policy.CreatedAt)
	})

	winner := scored[0]
	return &PolicySelectionResult{
		Allow:              true,
		SelectedPolicyID:   winner.policy.ID,
		AttributionGroupID: winner.attributionGroup,
		WindowSeconds:      winner.windowSeconds,
	}, nil
}

// filterApplicablePolicies returns the enabled policies that target
// the requested provider and have at least one of the caller's groups
// in their source_groups. Caller's group set is matched
// case-sensitively against policy.SourceGroups.
func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput) []*types.Policy {
	if len(policies) == 0 {
		return nil
	}
	groupSet := make(map[string]struct{}, len(in.GroupIDs))
	for _, g := range in.GroupIDs {
		if g != "" {
			groupSet[g] = struct{}{}
		}
	}
	out := make([]*types.Policy, 0, len(policies))
	for _, p := range policies {
		if p == nil || !p.Enabled {
			continue
		}
		if !sliceContains(p.DestinationProviderIDs, in.ProviderID) {
			continue
		}
		if !anyGroupMatches(p.SourceGroups, groupSet) {
			continue
		}
		out = append(out, p)
	}
	return out
}

// anyPolicyHasGuardrails reports whether any policy references at least one
// guardrail, so the selector can skip loading guardrails when none do.
func anyPolicyHasGuardrails(policies []*types.Policy) bool {
	for _, p := range policies {
		if p != nil && len(p.GuardrailIDs) > 0 {
			return true
		}
	}
	return false
}

// applyModelGate is the model-allowlist gate scoped to the matched policies:
// it keeps the candidates whose guardrails permit the model (none enabled =
// unrestricted) and returns a deny result when policies apply but none
// permits it. The guardrail load is skipped when no candidate references a
// guardrail, and the provider's catalog id — which picks the model-id
// normalizer — is resolved only when a candidate actually restricts models:
// with no enabled allowlist every candidate is unrestricted, and a
// provider-store failure must not fail a request the gate would have waved
// through.
func (m *managerImpl) applyModelGate(ctx context.Context, in PolicySelectionInput, candidates []*types.Policy) ([]*types.Policy, *PolicySelectionResult, error) {
	if len(candidates) == 0 || !anyPolicyHasGuardrails(candidates) {
		return candidates, nil, nil
	}
	guardrailsByID, err := m.loadGuardrailsByID(ctx, in.AccountID)
	if err != nil {
		return nil, nil, err
	}
	if !anyEnabledModelAllowlist(candidates, guardrailsByID) {
		return candidates, nil, nil
	}
	catalogID, err := m.providerCatalogID(ctx, in.AccountID, in.ProviderID)
	if err != nil {
		return nil, nil, err
	}
	permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model, catalogID)
	if len(permitted) == 0 {
		return nil, &PolicySelectionResult{
			Allow:      false,
			DenyCode:   denyCodeModelBlocked,
			DenyReason: modelBlockedReason(in.Model),
		}, nil
	}
	return permitted, nil, nil
}

// anyEnabledModelAllowlist reports whether any policy references a guardrail
// whose model allowlist is enabled — the only case the model gate restricts
// anything. Disabled allowlists, stale guardrail references, and guardrails
// carrying only other checks all leave every candidate unrestricted.
func anyEnabledModelAllowlist(policies []*types.Policy, byID map[string]*types.Guardrail) bool {
	for _, p := range policies {
		if p == nil {
			continue
		}
		for _, gID := range p.GuardrailIDs {
			if g, ok := byID[gID]; ok && g != nil && g.Checks.ModelAllowlist.Enabled {
				return true
			}
		}
	}
	return false
}

// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
// model-allowlist gate to resolve each candidate policy's attached guardrails.
func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
	guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID)
	if err != nil {
		return nil, fmt.Errorf("list account guardrails: %w", err)
	}
	byID := make(map[string]*types.Guardrail, len(guardrails))
	for _, g := range guardrails {
		if g != nil {
			byID[g.ID] = g
		}
	}
	return byID, nil
}

// providerCatalogID resolves a provider record id to its catalog provider
// id, the key the model-id normalizers are picked by. A missing provider
// resolves to the empty catalog id — the compare then runs verbatim-only,
// which can never widen an allowlist — while a store failure propagates
// rather than degrading a security decision.
func (m *managerImpl) providerCatalogID(ctx context.Context, accountID, providerID string) (string, error) {
	if providerID == "" {
		return "", nil
	}
	provider, err := m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthNone, accountID, providerID)
	switch {
	case err == nil:
		return provider.ProviderID, nil
	case isNotFound(err):
		return "", nil
	default:
		return "", fmt.Errorf("get provider: %w", err)
	}
}

// filterModelPermittedPolicies returns the subset of policies whose guardrails
// permit the model on the provider with the given catalog id. Order is
// preserved so downstream scoring is unaffected.
func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model, catalogProviderID string) []*types.Policy {
	out := make([]*types.Policy, 0, len(policies))
	for _, p := range policies {
		if policyPermitsModel(p, byID, model, catalogProviderID) {
			out = append(out, p)
		}
	}
	return out
}

// policyPermitsModel reports whether a policy permits the model. No
// allowlist-enabled guardrail = unrestricted (permits any, incl. empty);
// otherwise the model must be in the union of its allowlists, so an
// empty/undetermined model fails closed. An entry matches on its own
// normalised form or, for a path-style provider, its canonical form: the
// parser emits the canonical id for path-routed requests, while an
// allowlist may hold the raw declared id the dashboard's picker copies
// from the provider. The catalog id picks the normalizer, so a plain
// provider's entries always compare verbatim.
func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model, catalogProviderID string) bool {
	if p == nil {
		return false
	}
	wanted := normaliseModelID(model)
	restricted := false
	for _, gID := range p.GuardrailIDs {
		g, ok := byID[gID]
		if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
			continue
		}
		restricted = true
		if wanted == "" {
			continue
		}
		for _, allowed := range g.Checks.ModelAllowlist.Models {
			if normaliseModelID(allowed) == wanted || canonicalModelKey(catalogProviderID, allowed) == wanted {
				return true
			}
		}
	}
	return !restricted
}

// normaliseModelID lowercases and trims a model identifier so the allowlist
// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's
// normaliseModel so both layers agree on what "same model" means.
func normaliseModelID(model string) string {
	return strings.ToLower(strings.TrimSpace(model))
}

// modelBlockedReason builds the human-readable deny reason for a model-allowlist
// rejection. The model is quoted when known; an undetermined model is reported
// as such so the access log distinguishes "wrong model" from "no model".
func modelBlockedReason(model string) string {
	if normaliseModelID(model) == "" {
		return "request model could not be determined for the policy allowlist"
	}
	return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model)
}

// candidate is the per-policy intermediate the selector ranks. A
// policy that's been exhausted on any enabled cap never makes it
// into this slice; the selector's deny envelope carries the latest
// exhaustion's reason out separately.
type candidate struct {
	policy           *types.Policy
	attributionGroup string
	windowSeconds    int64
}

// scoreCandidates evaluates every applicable policy against the
// caller's current consumption. Exhausted policies are filtered out
// of the returned slice; the most recent exhaustion's deny code +
// human reason is returned alongside so the caller can surface it
// when no candidate survives.
func scoreCandidates(
	in PolicySelectionInput,
	candidates []*types.Policy,
	cache consumptionCache,
	now time.Time,
) ([]candidate, string, string) {
	out := make([]candidate, 0, len(candidates))
	var lastDenyCode, lastDenyReason string

	for _, p := range candidates {
		c, exhausted, denyCode, denyReason := scoreOne(in, p, cache, now)
		if exhausted {
			lastDenyCode = denyCode
			lastDenyReason = denyReason
			continue
		}
		out = append(out, c)
	}
	return out, lastDenyCode, lastDenyReason
}

// scoreOne checks a single policy for cap exhaustion. Returns the
// candidate envelope when the policy still has headroom on every
// enabled cap; reports exhausted=true with a deny code naming the
// offending cap kind otherwise.
func scoreOne(
	in PolicySelectionInput,
	p *types.Policy,
	cache consumptionCache,
	now time.Time,
) (candidate, bool, string, string) {
	attrGroup := lowestIntersect(p.SourceGroups, in.GroupIDs)
	c := candidate{
		policy:           p,
		attributionGroup: attrGroup,
		windowSeconds:    effectiveWindowSeconds(p),
	}

	if p.Limits.TokenLimit.Enabled && p.Limits.TokenLimit.WindowSeconds > 0 {
		if exhausted, reason := evalTokenCap(cache, in.AccountID, in.UserID, attrGroup, p.Limits.TokenLimit, now, "policy "+p.ID); exhausted {
			return candidate{}, true, denyCodeTokenCapExceeded, reason
		}
	}

	if p.Limits.BudgetLimit.Enabled && p.Limits.BudgetLimit.WindowSeconds > 0 {
		if exhausted, reason := evalBudgetCap(cache, in.AccountID, in.UserID, attrGroup, p.Limits.BudgetLimit, now, "policy "+p.ID); exhausted {
			return candidate{}, true, denyCodeBudgetCapExceeded, reason
		}
	}

	return c, false, "", ""
}

// evalTokenCap reports whether the token limit is already exhausted for the
// caller in its own window. attrGroup may be empty (no group dimension applies).
// label identifies the cap source ("policy <id>" or "account rule <id>") for the
// deny reason. It is the shared primitive behind both policy and account-rule
// enforcement.
func evalTokenCap(
	cache consumptionCache,
	accountID, userID, attrGroup string,
	tl types.PolicyTokenLimit,
	now time.Time,
	label string,
) (bool, string) {
	windowStart := types.WindowStart(now, tl.WindowSeconds)

	if tl.UserCap > 0 && userID != "" {
		row := cache.get(accountID, types.DimensionUser, userID, tl.WindowSeconds, windowStart)
		used := row.TokensInput + row.TokensOutput
		if used >= tl.UserCap {
			return true, fmt.Sprintf("user token cap exhausted on %s (used %d of %d)", label, used, tl.UserCap)
		}
	}

	if tl.GroupCap > 0 && attrGroup != "" {
		row := cache.get(accountID, types.DimensionGroup, attrGroup, tl.WindowSeconds, windowStart)
		used := row.TokensInput + row.TokensOutput
		if used >= tl.GroupCap {
			return true, fmt.Sprintf("group token cap exhausted on %s (used %d of %d)", label, used, tl.GroupCap)
		}
	}

	return false, ""
}

// evalBudgetCap is the budget (USD) counterpart of evalTokenCap.
func evalBudgetCap(
	cache consumptionCache,
	accountID, userID, attrGroup string,
	bl types.PolicyBudgetLimit,
	now time.Time,
	label string,
) (bool, string) {
	windowStart := types.WindowStart(now, bl.WindowSeconds)

	if bl.UserCapUsd > 0 && userID != "" {
		row := cache.get(accountID, types.DimensionUser, userID, bl.WindowSeconds, windowStart)
		if row.CostUSD >= bl.UserCapUsd {
			return true, fmt.Sprintf("user budget cap exhausted on %s (used $%.4f of $%.4f)", label, row.CostUSD, bl.UserCapUsd)
		}
	}

	if bl.GroupCapUsd > 0 && attrGroup != "" {
		row := cache.get(accountID, types.DimensionGroup, attrGroup, bl.WindowSeconds, windowStart)
		if row.CostUSD >= bl.GroupCapUsd {
			return true, fmt.Sprintf("group budget cap exhausted on %s (used $%.4f of $%.4f)", label, row.CostUSD, bl.GroupCapUsd)
		}
	}

	return false, ""
}

// checkAccountBudget evaluates every applicable account-level budget rule as an
// all-must-pass ceiling. A rule applies when the caller is in its TargetUsers,
// one of its TargetGroups, or it has no targets at all (account-wide). Returns
// deny=true with an llm_account.* code on the first exhausted rule. Group caps
// attribute to the lowest intersecting group (the same model policies use), so
// multi-group behavior is unchanged.
func checkAccountBudget(in PolicySelectionInput, rules []*types.AccountBudgetRule, cache consumptionCache, now time.Time) (bool, string, string) {
	for _, r := range rules {
		if r == nil || !r.Enabled || !budgetRuleApplies(r, in) {
			continue
		}
		attrGroup := lowestIntersect(r.TargetGroups, in.GroupIDs)
		label := "account rule " + r.ID

		if r.Limits.TokenLimit.Enabled && r.Limits.TokenLimit.WindowSeconds > 0 {
			if exhausted, reason := evalTokenCap(cache, in.AccountID, in.UserID, attrGroup, r.Limits.TokenLimit, now, label); exhausted {
				return true, denyCodeAccountTokenCapExceeded, reason
			}
		}

		if r.Limits.BudgetLimit.Enabled && r.Limits.BudgetLimit.WindowSeconds > 0 {
			if exhausted, reason := evalBudgetCap(cache, in.AccountID, in.UserID, attrGroup, r.Limits.BudgetLimit, now, label); exhausted {
				return true, denyCodeAccountBudgetCapExceeded, reason
			}
		}
	}

	return false, "", ""
}

// budgetRuleApplies reports whether an account budget rule binds the caller:
// a direct user match, a group intersection, or an untargeted (account-wide)
// rule.
func budgetRuleApplies(r *types.AccountBudgetRule, in PolicySelectionInput) bool {
	if len(r.TargetUsers) == 0 && len(r.TargetGroups) == 0 {
		return true
	}
	if in.UserID != "" && sliceContains(r.TargetUsers, in.UserID) {
		return true
	}
	groupSet := make(map[string]struct{}, len(in.GroupIDs))
	for _, g := range in.GroupIDs {
		if g != "" {
			groupSet[g] = struct{}{}
		}
	}
	return anyGroupMatches(r.TargetGroups, groupSet)
}

// RecordAccountBudgetUsage fans the served request's usage out to every
// applicable account budget rule's own (dimension, window) counter. The user
// dimension is always booked when a rule has a user-applicable cap; the group
// dimension books against the rule's lowest intersecting group. This runs
// alongside the policy-window record so account ceilings accumulate in their own
// windows (commonly monthly) independently of the per-policy window.
func (m *managerImpl) RecordAccountBudgetUsage(ctx context.Context, accountID, userID string, groupIDs []string, tokensIn, tokensOut int64, costUSD float64) error {
	if accountID == "" {
		return status.Errorf(status.InvalidArgument, "account_id is required")
	}
	if err := validateUsageDeltas(tokensIn, tokensOut, costUSD); err != nil {
		return err
	}
	rules, err := m.store.GetAccountAgentNetworkBudgetRules(ctx, store.LockingStrengthNone, accountID)
	if err != nil {
		return fmt.Errorf("list account budget rules: %w", err)
	}
	set := make(map[types.ConsumptionKey]struct{})
	addAccountBudgetKeys(set, PolicySelectionInput{AccountID: accountID, UserID: userID, GroupIDs: groupIDs}, rules, time.Now().UTC())
	if len(set) == 0 {
		return nil
	}
	return m.store.IncrementAgentNetworkConsumptionBatch(ctx, accountID, keysSlice(set), tokensIn, tokensOut, costUSD)
}

// RecordUsageInput carries everything RecordUsage books for one served request.
type RecordUsageInput struct {
	AccountID          string
	UserID             string
	AttributionGroupID string // selected policy's attribution group (policy window)
	GroupIDs           []string
	WindowSeconds      int64 // selected policy's window; 0 means no policy cap
	TokensIn           int64
	TokensOut          int64
	CostUSD            float64
}

// RecordUsage books a served request's usage against every counter it touches —
// the selected policy's per-(user, group) window plus every applicable account
// budget rule's own window — deduplicated and written in a single transaction.
// Two counters that collapse to the same (dimension, window) tuple are booked
// once, so a single request can never double-count against one cap.
func (m *managerImpl) RecordUsage(ctx context.Context, in RecordUsageInput) error {
	if in.AccountID == "" {
		return status.Errorf(status.InvalidArgument, "account_id is required")
	}
	if err := validateUsageDeltas(in.TokensIn, in.TokensOut, in.CostUSD); err != nil {
		return err
	}
	now := time.Now().UTC()
	set := make(map[types.ConsumptionKey]struct{})

	// Policy-window dimensions are booked only when a policy cap bound this
	// request (window > 0). A zero window means catch-all-allow / no policy cap;
	// the account fan-out below still books against the budget rules' windows.
	if in.WindowSeconds > 0 {
		addLimitKeys(set, in.UserID, in.AttributionGroupID, in.WindowSeconds, now)
	}

	rules, err := m.store.GetAccountAgentNetworkBudgetRules(ctx, store.LockingStrengthNone, in.AccountID)
	if err != nil {
		return fmt.Errorf("list account budget rules: %w", err)
	}
	addAccountBudgetKeys(set, PolicySelectionInput{AccountID: in.AccountID, UserID: in.UserID, GroupIDs: in.GroupIDs}, rules, now)

	if len(set) == 0 {
		return nil
	}
	return m.store.IncrementAgentNetworkConsumptionBatch(ctx, in.AccountID, keysSlice(set), in.TokensIn, in.TokensOut, in.CostUSD)
}

// addAccountBudgetKeys adds the (dimension, window) keys a served request books
// against every applicable account budget rule into the dedup set.
func addAccountBudgetKeys(set map[types.ConsumptionKey]struct{}, in PolicySelectionInput, rules []*types.AccountBudgetRule, now time.Time) {
	for _, r := range rules {
		if r == nil || !r.Enabled || !budgetRuleApplies(r, in) {
			continue
		}
		attrGroup := lowestIntersect(r.TargetGroups, in.GroupIDs)
		for _, window := range ruleWindows(r) {
			addLimitKeys(set, in.UserID, attrGroup, window, now)
		}
	}
}

// keysSlice flattens a ConsumptionKey set into a slice.
func keysSlice(set map[types.ConsumptionKey]struct{}) []types.ConsumptionKey {
	keys := make([]types.ConsumptionKey, 0, len(set))
	for k := range set {
		keys = append(keys, k)
	}
	return keys
}

// ruleWindows returns the distinct enabled window lengths a budget rule books
// against (token window and/or budget window, deduplicated).
func ruleWindows(r *types.AccountBudgetRule) []int64 {
	var windows []int64
	if r.Limits.TokenLimit.Enabled && r.Limits.TokenLimit.WindowSeconds > 0 {
		windows = append(windows, r.Limits.TokenLimit.WindowSeconds)
	}
	if r.Limits.BudgetLimit.Enabled && r.Limits.BudgetLimit.WindowSeconds > 0 {
		bw := r.Limits.BudgetLimit.WindowSeconds
		if len(windows) == 0 || windows[0] != bw {
			windows = append(windows, bw)
		}
	}
	return windows
}

// effectiveWindowSeconds returns the window length the proxy should
// hand back to RecordLLMUsage. When both halves are enabled with
// different windows, token_limit wins (the more common config); when
// only one is enabled that one wins; when neither is enabled the
// returned value is 0 — RecordLLMUsage treats 0 as "no limit
// tracking" and skips the increment, which is the right pass-through
// for catch-all-allow policies with no caps configured.
func effectiveWindowSeconds(p *types.Policy) int64 {
	if p.Limits.TokenLimit.Enabled && p.Limits.TokenLimit.WindowSeconds > 0 {
		return p.Limits.TokenLimit.WindowSeconds
	}
	if p.Limits.BudgetLimit.Enabled && p.Limits.BudgetLimit.WindowSeconds > 0 {
		return p.Limits.BudgetLimit.WindowSeconds
	}
	return 0
}

// lowestIntersect returns the lowest-by-string-sort element of
// callerGroups ∩ sourceGroups. Empty when the intersection is empty.
// Lowest is deterministic so multi-node selection converges.
func lowestIntersect(sourceGroups, callerGroups []string) string {
	if len(sourceGroups) == 0 || len(callerGroups) == 0 {
		return ""
	}
	srcSet := make(map[string]struct{}, len(sourceGroups))
	for _, g := range sourceGroups {
		srcSet[g] = struct{}{}
	}
	var best string
	for _, g := range callerGroups {
		if _, ok := srcSet[g]; !ok {
			continue
		}
		if best == "" || g < best {
			best = g
		}
	}
	return best
}

func anyGroupMatches(sourceGroups []string, callerSet map[string]struct{}) bool {
	for _, g := range sourceGroups {
		if _, ok := callerSet[g]; ok {
			return true
		}
	}
	return false
}

// isUncapped reports whether a policy has any enabled cap with a
// positive limit value. Mirrors the eval functions' guards: a policy
// with token_limit.enabled=true but every cap value at 0 still
// counts as uncapped because the eval would query nothing and bind
// nothing.
func isUncapped(p *types.Policy) bool {
	tl := p.Limits.TokenLimit
	if tl.Enabled && tl.WindowSeconds > 0 && (tl.GroupCap > 0 || tl.UserCap > 0) {
		return false
	}
	bl := p.Limits.BudgetLimit
	if bl.Enabled && bl.WindowSeconds > 0 && (bl.GroupCapUsd > 0 || bl.UserCapUsd > 0) {
		return false
	}
	return true
}

// groupCapTokens returns the policy's group-token cap when the token
// limit is enabled, zero otherwise. Drives the primary "bigger pool
// first" sort.
func groupCapTokens(p *types.Policy) int64 {
	if p.Limits.TokenLimit.Enabled {
		return p.Limits.TokenLimit.GroupCap
	}
	return 0
}

// groupCapBudgetUsd returns the policy's group-budget cap in USD
// when the budget limit is enabled, zero otherwise. Secondary sort
// key after token group cap so budget-only policies still order
// predictably.
func groupCapBudgetUsd(p *types.Policy) float64 {
	if p.Limits.BudgetLimit.Enabled {
		return p.Limits.BudgetLimit.GroupCapUsd
	}
	return 0
}

// userCapTokens returns the policy's per-user token cap when the
// token limit is enabled, zero otherwise. Tertiary sort key, used
// when group caps tie or are absent.
func userCapTokens(p *types.Policy) int64 {
	if p.Limits.TokenLimit.Enabled {
		return p.Limits.TokenLimit.UserCap
	}
	return 0
}

// userCapBudgetUsd returns the policy's per-user budget cap in USD
// when the budget limit is enabled, zero otherwise. Quaternary sort
// key for budget-only policies whose group caps tie or are absent.
func userCapBudgetUsd(p *types.Policy) float64 {
	if p.Limits.BudgetLimit.Enabled {
		return p.Limits.BudgetLimit.UserCapUsd
	}
	return 0
}

func sliceContains(haystack []string, needle string) bool {
	for _, v := range haystack {
		if v == needle {
			return true
		}
	}
	return false
}

// mockManager fallback so tests that don't care about selection still
// compile.
func (*mockManager) SelectPolicyForRequest(_ context.Context, _ PolicySelectionInput) (*PolicySelectionResult, error) {
	return &PolicySelectionResult{Allow: true}, nil
}