From 496bf8817c76e33b8652a28e4ebceb140cb9a16b Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 07:25:13 +0000 Subject: [PATCH 01/17] =?UTF-8?q?migrations:=200008=20group=20scopes=20?= =?UTF-8?q?=E2=80=94=20billing=5Fevent.member=5Fgroup=5Fids=20+=20group=5F?= =?UTF-8?q?usage?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Membership-aware group quotas (ruled 2026-10-07, Q-T1(b)+Q-T2(b)). billing_event gains member_group_ids TEXT[] (membership evidence, never part of the money grain); the new group_usage table is the per-(group, hour) ATTRIBUTION rollup the rater writes while rating and the admission group spend check reads. Rollout: code and schema together, in the 0007 order. --- migrations/0008_group_scopes.down.sql | 8 ++++ migrations/0008_group_scopes.up.sql | 63 +++++++++++++++++++++++++++ migrations/README.md | 19 ++++++++ 3 files changed, 90 insertions(+) create mode 100644 migrations/0008_group_scopes.down.sql create mode 100644 migrations/0008_group_scopes.up.sql diff --git a/migrations/0008_group_scopes.down.sql b/migrations/0008_group_scopes.down.sql new file mode 100644 index 0000000..a034c65 --- /dev/null +++ b/migrations/0008_group_scopes.down.sql @@ -0,0 +1,8 @@ +-- Roll back 0008. group_usage is fully owned by this migration (drop the +-- table); member_group_ids is evidence only (drop the column). Rated windows +-- keep their rated_usage money rows either way -- this rollback never touches +-- money. Re-rating a window after a down/up cycle rebuilds the group rows. + +DROP TABLE IF EXISTS group_usage; + +ALTER TABLE billing_event DROP COLUMN IF EXISTS member_group_ids; diff --git a/migrations/0008_group_scopes.up.sql b/migrations/0008_group_scopes.up.sql new file mode 100644 index 0000000..5184e1b --- /dev/null +++ b/migrations/0008_group_scopes.up.sql @@ -0,0 +1,63 @@ +-- Membership-aware group quotas (ruled 2026-10-07, Q-T1(b) + Q-T2(b)). +-- +-- Two additions, both in service of group-level admission enforcement: +-- +-- 1. billing_event.member_group_ids — the groups the CALLER belongs to +-- (X-Saturn-Group-Scopes membership), captured at meter time as evidence. +-- EVIDENCE ONLY: it never enters the money grain. rated_usage stays keyed +-- (auth_id, owner_type, owner_id, resource_id, model_id, serving_mode, +-- window_start) and the money row is still written exactly once per window. +-- +-- 2. group_usage — the GROUP ATTRIBUTION rollup the monthly spend cap reads: +-- per (group_id, hour) token sums and cost, written by the rater WHILE it +-- rates a window into money. Attribution sources per event: the token's own +-- group (billing_event.group_id, the group-token case) PLUS every group in +-- member_group_ids (the membership case). A user in N groups contributes to +-- N group rows. This is ATTRIBUTION, not money: the same event cost appears +-- under every group it belongs to, and withheld (non-money) events +-- contribute nothing. +-- +-- ROLLOUT: the rater's statement references member_group_ids and group_usage, +-- so a pre-0008 rater cannot run against this schema and a post-0008 rater +-- cannot run against the old one (SQLSTATE 42703) — roll code and schema +-- together, in the order used for migration 0007 (migrate up, then deploy). + +-- --------------------------------------------------------------------------- +-- billing_event: membership evidence +-- --------------------------------------------------------------------------- + +-- The caller's group memberships, captured at meter time from the trusted +-- X-Saturn-Group-Scopes envelope. NULLABLE: an event whose caller belongs to no +-- group (or whose edge does not stamp the envelope) carries NULL, exactly like +-- every other identity column. NULL means "no membership evidence", never +-- "zero groups" as a chargeable fact. +ALTER TABLE billing_event ADD COLUMN member_group_ids TEXT[]; + +COMMENT ON COLUMN billing_event.member_group_ids IS + 'Groups the caller belongs to (X-Saturn-Group-Scopes membership), captured at meter time. Evidence only -- never part of the billing grain. NULL when the caller has no group memberships.'; + +-- --------------------------------------------------------------------------- +-- group_usage: the group attribution rollup +-- --------------------------------------------------------------------------- + +CREATE TABLE group_usage ( + group_id TEXT NOT NULL, + -- The UTC hour bucket, same bucketing expression the rater uses for + -- rated_usage.window_start, so a group's monthly spend sum reads + -- window_start >= date_trunc('month', now()) over aligned hours. + window_start TIMESTAMPTZ NOT NULL, + prompt_tokens BIGINT NOT NULL DEFAULT 0, + cached_tokens BIGINT NOT NULL DEFAULT 0, + completion_tokens BIGINT NOT NULL DEFAULT 0, + billable_prompt_tokens BIGINT NOT NULL DEFAULT 0, + -- NUMERIC(20,9), the rated_usage.cost unit, so the admission spend check + -- compares in the same unit caps are authored in. Attribution cost: the + -- summed per-event cost of the events attributed to this group in this + -- hour. The same event cost legitimately appears under several groups. + cost NUMERIC(20,9) NOT NULL DEFAULT 0, + event_count BIGINT NOT NULL DEFAULT 0, + CONSTRAINT pk_group_usage PRIMARY KEY (group_id, window_start) +); + +COMMENT ON TABLE group_usage IS + 'Group attribution rollup written by the rater while rating a window into money: per (group_id, hour) token sums and cost. Attribution, not money -- an event attributed to N groups contributes to N rows, and withheld (non-money) events contribute nothing. The admission group spend check reads SUM(cost) here; rated_usage remains the only money rollup.'; diff --git a/migrations/README.md b/migrations/README.md index cb9d0c9..8eb5e93 100644 --- a/migrations/README.md +++ b/migrations/README.md @@ -50,6 +50,7 @@ golang-migrate up/down pairs, applied in version order: | 0005 | `0005_invoice_grade_attempts.{up,down}.sql` | trusted/client request identity, attempt outcome and usage evidence, invalid-usage reconciliation, and the hourly reconciliation view at the rated natural key (exposing missing/conflicting org evidence) | | 0006 | `0006_rollup_grain.{up,down}.sql` | widens the `rated_usage` grain with `serving_mode`, the owner pair and `graph_k8s_name` evidence | | 0007 | `0007_serving_mode_explicit.{up,down}.sql` | `rated_usage.serving_mode` becomes `'shared'`/`'dedicated'` only (CHECK, no default; dedicated rows renamed from `''` with their ids recomputed) | +| 0008 | `0008_group_scopes.{up,down}.sql` | `billing_event.member_group_ids` (membership evidence) + `group_usage`, the group attribution rollup the admission group spend check reads | `embed.go` embeds these into the `migrations` package; `cmd/migrate` applies them. @@ -64,6 +65,24 @@ migrate down # roll back one step migrate version # print the current applied version ``` +### Rollout order for migration 0008 + +Migration 0008 adds `billing_event.member_group_ids` and the `group_usage` +table, both of which the rater's single rating statement reads/writes. A +pre-0008 rater fails against this schema and a post-0008 rater fails against +the old one (SQLSTATE 42703), so roll code and schema together, in the +migration-0007 order: run `cmd/migrate up`, then deploy the new drainer and +rater (the drainer's INSERT gains the column; an old drainer against schema +0008 simply writes NULL membership). The interceptor's group scope enforcement +is envelope-driven, so it is safe across the cutover: no stamped envelope, no +group check. + +`group_usage` rows for already-rated hours exist only after a re-rate (the +rollup is written by the rater, not backfilled). Until a window is re-rated, +the admission spend check sums an empty rollup for that group's hours — group +rate limits still enforce from the envelope; spend caps enforce once rating +has covered the month. + Rolling back 0007 is not exact. The 0007 down migration maps every `billing_event.serving_mode = 'dedicated'` back to NULL. It also renames the `rated_usage` `'dedicated'` rows back to `''` and recomputes their ids, so a From 1a1776490bfd0c34c58c08c27fabbf03a143cfe9 Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 07:25:23 +0000 Subject: [PATCH 02/17] identity: X-Saturn-Group-Scopes joins the pinned trusted set; membership rides Identity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The whole group scope set travels one allowlisted header (grammar frozen with the design). FromRequest reads it only through the R3 trusted-header registry (pinned fallback 13 → 14, in lockstep with the chart) and extracts the caller's group list leniently — the strict, fail-closed parse runs at the proxy, and a malformed envelope never reaches metering. --- internal/identity/group_scopes_test.go | 89 +++++++++++++++++++++ internal/identity/identity.go | 95 ++++++++++++++++++++++- internal/identity/trusted_headers.go | 9 ++- internal/identity/trusted_headers_test.go | 17 ++-- 4 files changed, 195 insertions(+), 15 deletions(-) create mode 100644 internal/identity/group_scopes_test.go diff --git a/internal/identity/group_scopes_test.go b/internal/identity/group_scopes_test.go new file mode 100644 index 0000000..8fe3d16 --- /dev/null +++ b/internal/identity/group_scopes_test.go @@ -0,0 +1,89 @@ +package identity + +import ( + "net/http/httptest" + "strings" + "testing" +) + +// TestFromRequestMemberGroupIDs pins the metering-evidence side of the group +// quota envelope: the trusted X-Saturn-Group-Scopes header parses to the +// caller's group list (envelope order), and an untrusted or absent header +// leaves the list empty (the drainer stores NULL). +func TestFromRequestMemberGroupIDs(t *testing.T) { + const gidA = "a1b2c3d4e5f60718293a4b5c6d7e8f90" + const gidB = "00112233445566778899aabbccddeeff" + + cases := []struct { + name string + header string + want []string + }{ + {"absent", "", nil}, + {"single", "v1;" + gidA + ":30,1000,500,2000,100.000000000", []string{gidA}}, + {"multi in envelope order", "v1;" + gidB + ":,,,,;" + gidA + ":1,,,", []string{gidB, gidA}}, + {"all-empty entry still lists the group", "v1;" + gidA + ":,,,,", []string{gidA}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + if tc.header != "" { + r.Header.Set(HeaderGroupScopes, tc.header) + } + id := FromRequest(r) + if len(id.MemberGroupIDs) != len(tc.want) { + t.Fatalf("MemberGroupIDs = %v, want %v", id.MemberGroupIDs, tc.want) + } + for i, want := range tc.want { + if id.MemberGroupIDs[i] != want { + t.Fatalf("MemberGroupIDs[%d] = %q, want %q", i, id.MemberGroupIDs[i], want) + } + } + // The raw envelope rides the Identity for the strict proxy parse. + if tc.header != "" && id.GroupScopes != tc.header { + t.Fatalf("GroupScopes = %q, want the raw header verbatim", id.GroupScopes) + } + }) + } +} + +// TestFromRequestMemberGroupIDsLenient pins the deliberate split of +// responsibility: the identity parse extracts group ids LENIENTLY (only the +// v1 prefix and the 32-hex shape), because the strict fail-closed parse runs +// at the proxy and a malformed envelope never reaches metering. A header with +// a bad entry still yields the well-formed groups; a non-v1 header yields +// none. +func TestFromRequestMemberGroupIDsLenient(t *testing.T) { + const gidA = "a1b2c3d4e5f60718293a4b5c6d7e8f90" + + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + r.Header.Set(HeaderGroupScopes, "v1;NOT-A-GROUP:,,,,"+";"+gidA+":1,,,") + if id := FromRequest(r); len(id.MemberGroupIDs) != 1 || id.MemberGroupIDs[0] != gidA { + t.Fatalf("MemberGroupIDs = %v, want only the well-formed group %s", id.MemberGroupIDs, gidA) + } + + r = httptest.NewRequest("POST", "/v1/chat/completions", nil) + r.Header.Set(HeaderGroupScopes, "v9;"+gidA+":,,,,") + if id := FromRequest(r); id.MemberGroupIDs != nil { + t.Fatalf("MemberGroupIDs = %v for a non-v1 header, want nil", id.MemberGroupIDs) + } +} + +// TestFromRequestGroupScopesUntrustedHeaderAbsent pins the R3 gate for the new +// envelope: when the header is NOT in the active trusted set, it reads as +// ABSENT — no group list, no raw value, no matter what the client sends. +func TestFromRequestGroupScopesUntrustedHeaderAbsent(t *testing.T) { + const gidA = "a1b2c3d4e5f60718293a4b5c6d7e8f90" + // The active set WITHOUT X-Saturn-Group-Scopes, as a not-yet-upgraded + // chart would render it. + withTrustedHeadersEnv(t, strings.Join([]string{ + HeaderGateway, HeaderOrgID, HeaderOwnerID, HeaderServingMode, HeaderServedModel, + }, ","), true) + + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + r.Header.Set(HeaderGroupScopes, "v1;"+gidA+":,,,,") + id := FromRequest(r) + if id.GroupScopes != "" || id.MemberGroupIDs != nil { + t.Fatalf("untrusted group scopes read as present: GroupScopes=%q MemberGroupIDs=%v, want both absent", id.GroupScopes, id.MemberGroupIDs) + } +} diff --git a/internal/identity/identity.go b/internal/identity/identity.go index ef9f715..7b67e73 100644 --- a/internal/identity/identity.go +++ b/internal/identity/identity.go @@ -4,7 +4,10 @@ // headers, exactly as auth-server emits them. package identity -import "net/http" +import ( + "net/http" + "strings" +) // Header names injected by atlas-auth. Kept identical to auth-server's // constants so the contract between the two services stays in one shape. @@ -147,6 +150,23 @@ const ( HeaderOwnerRateLimitUncachedPromptTokens = "X-Saturn-Owner-Rate-Limit-Uncached-Prompt-Tokens" HeaderOwnerRateLimitGeneratedTokens = "X-Saturn-Owner-Rate-Limit-Generated-Tokens" + // HeaderGroupScopes carries the membership-aware GROUP quota envelope + // (ruled 2026-10-07, Q-T1(b)+Q-T2(b)): the whole group scope set rides ONE + // allowlisted header because Traefik's authResponseHeaders allowlist is a + // fixed name list and a member can belong to several groups. Grammar + // (frozen envelope, parsed strictly at the proxy/admission layer): + // + // v1;:,,,,;:... + // + // is the 32-hex group id; each rate field is empty (unlimited), 0 + // (zero cap), or a positive per-minute value (R4 sentinels); is + // empty (no cap) or a plain decimal NUMERIC(20,9) cap, 0 = zero cap. At + // most 16 entries, under 4 KiB; malformed or oversize fails closed (503) + // at the proxy. Like every scoped limit header it resolves through the R3 + // trusted-header registry: outside the active set it reads as ABSENT, so + // no group quota is ever read from a client-supplied value. + HeaderGroupScopes = "X-Saturn-Group-Scopes" + // HeaderUpstream carries the EXACT backend the request must be forwarded to — // `host:port` (e.g. pd-abcde-mymodel-r123.main-namespace.svc.cluster.local:8000). // @@ -228,6 +248,23 @@ type Identity struct { OwnerRateLimitTotalPromptTokens string OwnerRateLimitUncachedPromptTokens string OwnerRateLimitGeneratedTokens string + // GroupScopes is the raw, trusted X-Saturn-Group-Scopes envelope value + // (see HeaderGroupScopes), read through the R3 trusted-header registry. + // Empty when the header is absent or untrusted. The strict parse — per-group + // rate limits, monthly spend caps, bounds — happens at the proxy/admission + // layer (parseTrustedGroupScopes); Identity carries the envelope verbatim, + // exactly like the eight scoped rate-limit fields above. + GroupScopes string + // MemberGroupIDs is the caller's group list extracted from GroupScopes — + // the of every well-formed entry, in envelope order. It rides the + // metering event as billing_event.member_group_ids so the rater can + // attribute usage to each group the caller belongs to (membership-aware + // group quotas, ruled 2026-10-07). Empty when the envelope is absent or + // carries no group. Extraction is deliberately LENIENT (see + // memberGroupIDs): the strict, fail-closed parse runs at the proxy before + // enforcement, and a malformed envelope never reaches metering because the + // request is refused first. + MemberGroupIDs []string } // The two serving modes, spelled exactly as they are stored in billing_event and @@ -248,8 +285,9 @@ func ValidServingMode(s string) bool { // validation beyond reading the values; authorization happened at the edge. // // The R3 envelope reads (gateway mark, org, owner, serving mode, served -// model, and every rate-limit policy header — the pinned 13) resolve through -// the trusted-header registry: a header outside the active set is treated +// model, every rate-limit policy header, and the group scope envelope — +// the pinned 14) resolve through the trusted-header registry: a header +// outside the active set is treated // as ABSENT, never read for a trust decision. The remaining identity headers // are read directly (ratified edge contract, outside the R3 gate). The proxy // calls FromRequest BEFORE StripSaturnHeaders removes every X-Saturn-* header @@ -262,7 +300,7 @@ func ValidServingMode(s string) bool { // request's serving mode is overwritten later by gateway resolution from the // served-model registry. func FromRequest(r *http.Request) Identity { - return Identity{ + id := Identity{ AuthID: r.Header.Get(HeaderAuthID), UserID: r.Header.Get(HeaderUserID), GroupID: r.Header.Get(HeaderGroupID), @@ -284,7 +322,56 @@ func FromRequest(r *http.Request) Identity { OwnerRateLimitTotalPromptTokens: trustedHeaderValue(r, HeaderOwnerRateLimitTotalPromptTokens), OwnerRateLimitUncachedPromptTokens: trustedHeaderValue(r, HeaderOwnerRateLimitUncachedPromptTokens), OwnerRateLimitGeneratedTokens: trustedHeaderValue(r, HeaderOwnerRateLimitGeneratedTokens), + GroupScopes: trustedHeaderValue(r, HeaderGroupScopes), + } + id.MemberGroupIDs = memberGroupIDs(id.GroupScopes) + return id +} + +// memberGroupIDs extracts the group list from a raw X-Saturn-Group-Scopes +// envelope value: the of every well-formed v1 entry, in envelope order. +// The extraction is LENIENT by design — it validates only the pieces metering +// needs (the v1 prefix and each entry's 32-hex group id) and skips anything +// else. The full envelope — rate fields, spend cap, the 16-entry and 4 KiB +// bounds, duplicate gids — is parsed STRICTLY at the proxy before enforcement, +// and a malformed or oversize envelope fails closed there (503), so a request +// carrying one never reaches metering; this extraction feeding the evidence +// column therefore only ever sees values the strict parse already accepted. +// Absent/empty/untrusted header → nil, stored as NULL by the drainer like +// every other absent identity field. +func memberGroupIDs(raw string) []string { + if raw == "" { + return nil + } + parts := strings.Split(raw, ";") + if len(parts) < 2 || parts[0] != "v1" { + return nil + } + var out []string + for _, entry := range parts[1:] { + fields := strings.SplitN(entry, ":", 2) + if len(fields) != 2 || !isGroupID(fields[0]) { + continue + } + out = append(out, fields[0]) + } + return out +} + +// isGroupID reports whether s is a 32-char lowercase hex group id (uuid4 hex, +// as Atlas mints). The frozen envelope pins this grammar; uppercase or any +// other spelling is not a group id. +func isGroupID(s string) bool { + if len(s) != 32 { + return false + } + for i := 0; i < len(s); i++ { + c := s[i] + if !('0' <= c && c <= '9' || 'a' <= c && c <= 'f') { + return false + } } + return true } // OrgScopedRateLimitHeaders and OwnerScopedRateLimitHeaders are the scoped diff --git a/internal/identity/trusted_headers.go b/internal/identity/trusted_headers.go index 9dc0df3..aa28682 100644 --- a/internal/identity/trusted_headers.go +++ b/internal/identity/trusted_headers.go @@ -9,10 +9,10 @@ package identity // PHOEBE_TRUSTED_HEADERS, comma-separated, rendered from the saturn-k8s // phoebe chart's ConfigMap); when the config is empty, unset, or malformed // (nothing left after trimming and dropping empties), a hard-coded fallback -// equal to the pinned 13 engages so phoebe fails closed, and a loud warning +// equal to the pinned 14 engages so phoebe fails closed, and a loud warning // flags the misrendered chart for the operator. // -// The set covers exactly the R3 envelope headers listed in the pinned 13. +// The set covers exactly the R3 envelope headers listed in the pinned 14. // The remaining identity headers (AuthID / UserID / GroupID / ResourceID / // ResourceType / BaseModel / Adapter / Upstream) are read under the ratified // edge contract (ForwardAuth authResponseHeaders allowlist + Atlas @@ -62,6 +62,7 @@ var pinnedTrustedHeaders = []string{ HeaderOwnerRateLimitTotalPromptTokens, HeaderOwnerRateLimitUncachedPromptTokens, HeaderOwnerRateLimitGeneratedTokens, + HeaderGroupScopes, } // requiredTrustedHeaders are the envelope headers request handling cannot work @@ -91,7 +92,7 @@ type trustedHeaderSet map[string]struct{} var activeTrustedHeaders atomic.Pointer[trustedHeaderSet] func init() { - // Before config load the pinned 13 are active: phoebe never reads an + // Before config load the pinned 14 are active: phoebe never reads an // envelope header outside the list, before OR after the runtime config // is loaded, and the pre-load default must preserve request handling. fallback := newTrustedHeaderSet(pinnedTrustedHeaders) @@ -129,7 +130,7 @@ func trustedHeaderValue(r *http.Request, name string) string { // active (entries are canonicalized; case is ignored per HTTP // convention); // - unset, empty/whitespace-only, or nothing left after splitting and -// trimming -> the pinned 13 engage AND a loud warning flags the +// trimming -> the pinned 14 engage AND a loud warning flags the // misrendered chart. Fallback, never a startup failure: phoebe must keep // serving with the pinned set rather than crash-loop behind a broken // ConfigMap render. diff --git a/internal/identity/trusted_headers_test.go b/internal/identity/trusted_headers_test.go index 9f96166..6a9180e 100644 --- a/internal/identity/trusted_headers_test.go +++ b/internal/identity/trusted_headers_test.go @@ -48,12 +48,13 @@ func pinnedSet() map[string]struct{} { return set } -// TestPinnedFallbackIsExactlyThe13 pins the fallback to the exact +// TestPinnedFallbackIsExactlyThe14 pins the fallback to the exact // pinned names, in the exact pinned order — the set the chart's ConfigMap // is supposed to render. Ruling R8 removed the five legacy single-scope // quota headers (X-Saturn-Service-Tier, X-Saturn-Rate-Limit-*), taking the -// R3 set from 18 to 13. -func TestPinnedFallbackIsExactlyThe13(t *testing.T) { +// R3 set from 18 to 13; the membership-aware group scope envelope +// (X-Saturn-Group-Scopes, ruled 2026-10-07) takes it from 13 to 14. +func TestPinnedFallbackIsExactlyThe14(t *testing.T) { want := []string{ "X-Saturn-Gateway", "X-Saturn-Org-Id", @@ -68,14 +69,15 @@ func TestPinnedFallbackIsExactlyThe13(t *testing.T) { "X-Saturn-Owner-Rate-Limit-Total-Prompt-Tokens", "X-Saturn-Owner-Rate-Limit-Uncached-Prompt-Tokens", "X-Saturn-Owner-Rate-Limit-Generated-Tokens", + "X-Saturn-Group-Scopes", } if !reflect.DeepEqual(pinnedTrustedHeaders, want) { - t.Fatalf("pinnedTrustedHeaders = %v, want the pinned 13 in order %v", pinnedTrustedHeaders, want) + t.Fatalf("pinnedTrustedHeaders = %v, want the pinned 14 in order %v", pinnedTrustedHeaders, want) } } // TestLoadTrustedHeadersFallback: unset, empty, whitespace-only, or -// all-empty-after-split config engages the pinned 13 — never an empty set. +// all-empty-after-split config engages the pinned 14 — never an empty set. func TestLoadTrustedHeadersFallback(t *testing.T) { cases := []struct { name string @@ -92,7 +94,7 @@ func TestLoadTrustedHeadersFallback(t *testing.T) { t.Run(tc.name, func(t *testing.T) { withTrustedHeadersEnv(t, tc.value, tc.set) if got := ActiveTrustedHeaders(); !reflect.DeepEqual(got, pinnedSet()) { - t.Fatalf("active set = %v, want the pinned 13 %v", got, pinnedSet()) + t.Fatalf("active set = %v, want the pinned 14 %v", got, pinnedSet()) } }) } @@ -308,7 +310,7 @@ func TestLoadTrustedHeadersErrorsWhenOwnerOrOrgAnchorMissing(t *testing.T) { } } -// TestLoadTrustedHeadersNoErrorForPinnedList: the full pinned 13, configured +// TestLoadTrustedHeadersNoErrorForPinnedList: the full pinned 14, configured // explicitly, satisfies every required header and logs nothing at ERROR. func TestLoadTrustedHeadersNoErrorForPinnedList(t *testing.T) { out := loadTrustedHeadersCapturingErrors(t, strings.Join(pinnedTrustedHeaders, ",")) @@ -434,6 +436,7 @@ var candidateSaturnHeaders = []string{ HeaderOwnerRateLimitTotalPromptTokens, HeaderOwnerRateLimitUncachedPromptTokens, HeaderOwnerRateLimitGeneratedTokens, + HeaderGroupScopes, } // packageSaturnHeaderConstants parses the package's non-test Go files and From 911ff22e1fe0c8a550b045365c614251242eacbe Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 07:25:33 +0000 Subject: [PATCH 03/17] metering/drain: carry member_group_ids to billing_event (NULL when empty) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The event gains the caller's group list; the drainer binds it like every other identity column (nil for none — no empty arrays), making the column 24-wide. Recovery's duplicate-evidence compare switches to canonical JSON: Event now carries a slice and is no longer comparable with ==. --- internal/drain/drainer_test.go | 2 +- internal/drain/store.go | 17 ++++++++++- internal/drain/store_test.go | 49 +++++++++++++++++++----------- internal/metering/metering.go | 8 +++++ internal/recovery/recovery.go | 11 ++++++- internal/recovery/recovery_test.go | 7 +++-- 6 files changed, 71 insertions(+), 23 deletions(-) diff --git a/internal/drain/drainer_test.go b/internal/drain/drainer_test.go index e7f5125..5a740ff 100644 --- a/internal/drain/drainer_test.go +++ b/internal/drain/drainer_test.go @@ -615,7 +615,7 @@ func TestEnsureGroup_Idempotent(t *testing.T) { // explicit "" (a post-cutover producer bug) is stored as the empty string so // the rater withholds it, and "shared" and "dedicated" are stored as sent. func TestDecodeEvent_OnlyAbsentServingModeStoredAsDedicated(t *testing.T) { - const servingModeIdx = 11 + const servingModeIdx = 12 cases := []struct { name string json string diff --git a/internal/drain/store.go b/internal/drain/store.go index 051ed90..0e33428 100644 --- a/internal/drain/store.go +++ b/internal/drain/store.go @@ -91,6 +91,7 @@ var upsertColumns = []string{ "group_id", "resource_id", "resource_type", + "member_group_ids", "org_id", "model", "base_model", @@ -109,7 +110,7 @@ var upsertColumns = []string{ "event_ts", } -const colsPerRow = 23 // len(upsertColumns); created_at is DB-defaulted. +const colsPerRow = 24 // len(upsertColumns); created_at is DB-defaulted. // Upsert writes a batch of events in a single transaction with a multi-row // INSERT ... ON CONFLICT (request_id) DO NOTHING. @@ -186,6 +187,10 @@ func eventArgs(e metering.Event) []any { nullStr(e.GroupID), nullStr(e.ResourceID), nullStr(e.ResourceType), + // MemberGroupIDs is the caller's group list (X-Saturn-Group-Scopes + // membership). nil for len==0 so the column stores NULL, not an empty + // array — like every other identity column, absent must read as absent. + nullGroupIDs(e.MemberGroupIDs), // OrgID is "" when Atlas isn't injecting X-Saturn-Org-Id yet (producer-rollout // gap). nullStr so it stores NULL, not '' — the rater/push fail-closed predicate // is `org_id IS NULL` (held + screamed at push), and a stored '' would dodge it @@ -233,6 +238,16 @@ func eventArgs(e metering.Event) []any { } } +// nullGroupIDs returns a driver NULL for an empty slice and the []string +// otherwise (pgx encodes it as a text[] bind). Mirror of nullStr for the +// membership list: no memberships is NULL, not an empty array. +func nullGroupIDs(ids []string) any { + if len(ids) == 0 { + return nil + } + return ids +} + // nullStr returns a driver NULL for "" and the string otherwise. func nullStr(s string) any { if s == "" { diff --git a/internal/drain/store_test.go b/internal/drain/store_test.go index 035426c..7154559 100644 --- a/internal/drain/store_test.go +++ b/internal/drain/store_test.go @@ -3,6 +3,7 @@ package drain import ( "context" "errors" + "reflect" "regexp" "testing" "time" @@ -45,14 +46,14 @@ func TestPostgresStore_UpsertSQL(t *testing.T) { mock.ExpectBegin() mock.ExpectExec(regexp.QuoteMeta( - "INSERT INTO billing_event (request_id, client_request_id, auth_id, user_id, group_id, resource_id, resource_type, org_id, model, base_model, adapter, serving_mode, prompt_tokens, cached_tokens, completion_tokens, finish_reason, gpu_type, aborted, usage_found, status_code, streamed, graph_k8s_name, event_ts) VALUES", + "INSERT INTO billing_event (request_id, client_request_id, auth_id, user_id, group_id, resource_id, resource_type, member_group_ids, org_id, model, base_model, adapter, serving_mode, prompt_tokens, cached_tokens, completion_tokens, finish_reason, gpu_type, aborted, usage_found, status_code, streamed, graph_k8s_name, event_ts) VALUES", )). WithArgs( // row 1 (org_id + base_model NULL: a pre-cutover-shaped event with no org // header and no derived_from) - "req-1", "logical-1", "auth-1", nil, nil, nil, nil, nil, "m1", nil, nil, "dedicated", 5, 0, 7, nil, nil, false, false, nil, false, nil, time.UnixMilli(ts).UTC(), + "req-1", "logical-1", "auth-1", nil, nil, nil, nil, nil, nil, "m1", nil, nil, "dedicated", 5, 0, 7, nil, nil, false, false, nil, false, nil, time.UnixMilli(ts).UTC(), // row 2 (no identity, no timestamp → event_ts NULL) - "req-2", nil, nil, nil, nil, nil, nil, nil, "m2", nil, nil, "shared", 0, 0, 0, nil, nil, false, false, nil, false, nil, nil, + "req-2", nil, nil, nil, nil, nil, nil, nil, nil, "m2", nil, nil, "shared", 0, 0, 0, nil, nil, false, false, nil, false, nil, nil, ). WillReturnResult(sqlmock.NewResult(0, 2)) mock.ExpectCommit() @@ -124,7 +125,7 @@ func TestPostgresStore_EmptyModelStoredAsNull(t *testing.T) { mock.ExpectBegin() mock.ExpectExec("INSERT INTO billing_event"). WithArgs( - "req-no-model", nil, "auth-1", nil, nil, nil, nil, + "req-no-model", nil, "auth-1", nil, nil, nil, nil, nil, nil, // org_id: "" must bind NULL nil, // model: "" must bind NULL nil, // base_model: "" must bind NULL @@ -158,7 +159,7 @@ func TestPostgresStore_EmptyModelStoredAsNull(t *testing.T) { // and billing_event_status_code_ck admits NULL or 100..599 — so 0 must bind NULL // (a never-answered attempt), while a real status binds through verbatim. func TestEventArgs_ZeroStatusCodeBindsNullNoResponse(t *testing.T) { - const statusCodeIdx = 19 // request_id..usage_found is 19 columns; status_code is next. + const statusCodeIdx = 20 // request_id..usage_found is 20 columns; status_code is next. noResponse := eventArgs(metering.Event{RequestID: "r", Model: "m"}) if noResponse[statusCodeIdx] != nil { @@ -184,31 +185,43 @@ func TestEventArgs_NullsEmptyIdentities(t *testing.T) { if args[2] != nil { t.Fatalf("auth_id arg = %v, want nil for empty AuthID", args[2]) } - // org_id is index 6 (after resource_type) — must be nil for empty (no producer + // member_group_ids is index 7 (after resource_type) — must be nil for empty + // (no memberships), like every other absent identity field. + if args[7] != nil { + t.Fatalf("member_group_ids arg = %v, want nil for empty MemberGroupIDs", args[7]) + } + // org_id is index 8 (member_group_ids pushed it down by one) — must be nil for empty (no producer // header / rollout gap); a stored '' would dodge the rater/push `org_id IS NULL` // held-not-billed predicate. - if args[7] != nil { - t.Fatalf("org_id arg = %v, want nil for empty OrgID", args[7]) + if args[8] != nil { + t.Fatalf("org_id arg = %v, want nil for empty OrgID", args[8]) } - // base_model is index 8 (org_id pushed model/base_model down by one) — nil for empty. - if args[9] != nil { - t.Fatalf("base_model arg = %v, want nil for empty BaseModel", args[9]) + // base_model is index 10 (member_group_ids pushed org_id/model/base_model down by one) — nil for empty. + if args[10] != nil { + t.Fatalf("base_model arg = %v, want nil for empty BaseModel", args[10]) } - // prompt_tokens is index 11 (…base_model=8, adapter=9, serving_mode=10) — int, not nil. - if args[12] != 3 { - t.Fatalf("prompt_tokens arg = %v, want 3", args[12]) + // prompt_tokens is index 13 (…base_model=10, adapter=11, serving_mode=12) — int, not nil. + if args[13] != 3 { + t.Fatalf("prompt_tokens arg = %v, want 3", args[13]) } // event_ts is the last index — nil when TimestampUnixMs==0. if args[colsPerRow-1] != nil { t.Fatalf("event_ts arg = %v, want nil for zero timestamp", args[colsPerRow-1]) } - // POSITIVE bind: a non-empty OrgID must reach the org_id column verbatim (index 6), + // POSITIVE bind: a non-empty OrgID must reach the org_id column verbatim (index 8), // not just be NULLed when empty — the meter-time org capture is the point of the // change, so the happy path is pinned here (not only transitively via e2e). withOrg := eventArgs(metering.Event{RequestID: "r", Model: "m", OrgID: "org-xyz"}) - if withOrg[7] != "org-xyz" { - t.Fatalf("org_id arg = %v, want \"org-xyz\" (a non-empty OrgID must bind through)", withOrg[7]) + if withOrg[8] != "org-xyz" { + t.Fatalf("org_id arg = %v, want \"org-xyz\" (a non-empty OrgID must bind through)", withOrg[8]) + } + + // POSITIVE bind: a membership list binds as the slice (pgx encodes text[]), + // so the group attribution rollup can read it. + withGroups := eventArgs(metering.Event{RequestID: "r", Model: "m", MemberGroupIDs: []string{"group-a"}}) + if got := withGroups[7]; !reflect.DeepEqual(got, []string{"group-a"}) { + t.Fatalf("member_group_ids arg = %#v, want [group-a] (a membership list must bind through)", got) } } @@ -218,7 +231,7 @@ func TestEventArgs_NullsEmptyIdentities(t *testing.T) { // as ” (not NULL, not 'dedicated') so the rater withholds it as an invalid // serving mode, and every other value binds verbatim. func TestEventArgs_ServingModeStoredAsDecoded(t *testing.T) { - const servingModeIdx = 11 // request_id..adapter is 11 columns; serving_mode is next. + const servingModeIdx = 12 // request_id..adapter is 12 columns; serving_mode is next. cases := []struct { in string diff --git a/internal/metering/metering.go b/internal/metering/metering.go index 52a8401..24fd624 100644 --- a/internal/metering/metering.go +++ b/internal/metering/metering.go @@ -62,6 +62,14 @@ type Event struct { GroupID string `json:"group_id,omitempty"` // present on group tokens ResourceID string `json:"resource_id,omitempty"` // model / deployment id ResourceType string `json:"resource_type,omitempty"` // e.g. workspace, deployment + // MemberGroupIDs is the caller's group list from the trusted + // X-Saturn-Group-Scopes envelope (membership-aware group quotas, ruled + // 2026-10-07): the groups whose quotas this request counts against even + // though the token is the caller's own (a user token). EVIDENCE ONLY — it + // never enters the money grain; the rater attributes usage to each listed + // group (plus GroupID, the group-token case) in the group_usage rollup. + // Empty/nil when the caller belongs to no group (stored as NULL). + MemberGroupIDs []string `json:"member_group_ids,omitempty"` // OrgID is the org that OWNS the served deployment (E2 customer attribution), // injected by Atlas as a per-deployment Traefik header (X-Saturn-Org-Id). Captured // verbatim at meter time so push reads org off the rollup instead of re-joining diff --git a/internal/recovery/recovery.go b/internal/recovery/recovery.go index 7fc5d7d..dee58e0 100644 --- a/internal/recovery/recovery.go +++ b/internal/recovery/recovery.go @@ -274,7 +274,16 @@ func validateAndDedupe(events []metering.Event) (Evidence, error) { return Evidence{}, fmt.Errorf("record %d: %w", i+1, err) } if prior, ok := seen[ev.RequestID]; ok { - if prior != ev { + // Two records with one request_id must be identical evidence; a + // differing field (including a differing membership list) is a + // conflict. Compare canonical JSON, not ==: Event carries a slice + // (MemberGroupIDs) and is no longer comparable with the equality + // operator, and the marshaled form also normalizes a nil vs an + // empty membership list (both omit under omitempty — "no + // memberships" either way). + priorJSON, _ := json.Marshal(prior) + evJSON, _ := json.Marshal(ev) + if !bytes.Equal(priorJSON, evJSON) { return Evidence{}, fmt.Errorf("record %d: conflicting duplicate request_id %q", i+1, ev.RequestID) } result.Duplicates++ diff --git a/internal/recovery/recovery_test.go b/internal/recovery/recovery_test.go index 073bba1..db72330 100644 --- a/internal/recovery/recovery_test.go +++ b/internal/recovery/recovery_test.go @@ -1,6 +1,7 @@ package recovery import ( + "bytes" "context" "encoding/json" "fmt" @@ -175,8 +176,10 @@ func TestReplayUsesEmitterStreamShapeAndPreservesIDs(t *testing.T) { if err := json.Unmarshal([]byte(raw), &got); err != nil { t.Fatal(err) } - if got != events[i] { - t.Fatalf("row %d mismatch: got %+v want %+v", i, got, events[i]) + gotJSON, _ := json.Marshal(got) + wantJSON, _ := json.Marshal(events[i]) + if !bytes.Equal(gotJSON, wantJSON) { + t.Fatalf("row %d mismatch: got %s want %s", i, gotJSON, wantJSON) } } } From 81e0ed5884fc5c6ecb5a26fc786238f93093cf52 Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 07:25:43 +0000 Subject: [PATCH 04/17] rating: upsert the group_usage attribution rollup while rating a window MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The single rate statement now writes both rollups from one snapshot: money into rated_usage (schema, grain, and write count untouched) and attribution into group_usage — each event attributed to its token's own group plus every member_group_ids group, deduped per (event, group), withheld (non-money) events contributing nothing. The same reconcile contract as rated_usage: in-window group rows the run does not reproduce are deleted. Integration harnesses apply 0008; the new conformance test pins token-group + membership + multi-group + withheld-excluded + re-rate convergence against live Postgres. --- internal/e2e/e2e_test.go | 3 + .../rating/group_rollup_integration_test.go | 218 ++++++++++++++++++ internal/rating/rater.go | 14 ++ internal/rating/store.go | 154 ++++++++++++- internal/rating/store_integration_test.go | 21 +- internal/rating/store_test.go | 25 +- 6 files changed, 427 insertions(+), 8 deletions(-) create mode 100644 internal/rating/group_rollup_integration_test.go diff --git a/internal/e2e/e2e_test.go b/internal/e2e/e2e_test.go index a4c0eab..5ae0c5a 100644 --- a/internal/e2e/e2e_test.go +++ b/internal/e2e/e2e_test.go @@ -155,6 +155,9 @@ func newHarness(t *testing.T, schema string) *harness { // 0007 makes rated_usage.serving_mode 'shared'/'dedicated' only (CHECK, no // default); the rater's upsert must satisfy it. mustExec(t, db, readMigration(t, "0007_serving_mode_explicit.up.sql")) + // 0008 adds billing_event.member_group_ids (in the drainer's INSERT and the + // rater's group attribution) and the group_usage table the rater upserts. + mustExec(t, db, readMigration(t, "0008_group_scopes.up.sql")) mr, err := miniredis.Run() if err != nil { diff --git a/internal/rating/group_rollup_integration_test.go b/internal/rating/group_rollup_integration_test.go new file mode 100644 index 0000000..1450499 --- /dev/null +++ b/internal/rating/group_rollup_integration_test.go @@ -0,0 +1,218 @@ +//go:build integration + +package rating + +import ( + "context" + "database/sql" + "errors" + "fmt" + "os" + "testing" + "time" +) + +// TestIntegration_GroupUsageAttribution runs the REAL rating SQL against a live +// Postgres and pins the group attribution rollup (group_usage, migration 0008) +// the admission group spend check reads: +// +// - an event is attributed to its token's own group (group-token case) PLUS +// every group in member_group_ids (membership case): a user in two groups +// contributes to both rows; +// - an event whose own group also appears in member_group_ids contributes +// ONCE to that group (dedupe per (event, group)); +// - WITHHELD (non-money) events contribute nothing to any group row; +// - the rollup's cost is the summed per-event cost at the applied rate +// (asserted against the Rate() oracle), and the same event cost +// legitimately appears under every attributed group (attribution, not +// money — rated_total must not double-count it); +// - re-running the window is idempotent (same sums, zero reconcile +// deletions), and deleting an event re-rates: its groups' rows shrink or +// are reconcile-deleted. +func TestIntegration_GroupUsageAttribution(t *testing.T) { + dsn := os.Getenv("PHOEBE_TEST_DATABASE_URL") + if dsn == "" { + t.Skip("PHOEBE_TEST_DATABASE_URL not set; skipping live-Postgres conformance") + } + ctx := context.Background() + db, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatalf("open: %v", err) + } + defer db.Close() + + const sch = "phoebe_rating_group_it" + exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") + exec(t, db, "CREATE SCHEMA "+sch) + exec(t, db, "SET search_path TO "+sch) + defer func() { exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") }() + exec(t, db, ratingSchemaDDL(t)) + + hour := mustTime("2026-06-08T10:00:00Z") + book := newTestBook( + map[string]Rate3{"b": rate3("0.000005", "0.0000005", "0.00002")}, + nil, PolicyIdentity, Dec{}, Dec{}, + ) + rate, err := book.Resolve("b") + if err != nil { + t.Fatalf("resolve: %v", err) + } + billed := rate.Quantized() + + const ( + g1 = "a1b2c3d4e5f60718293a4b5c6d7e8f90" + g2 = "00112233445566778899aabbccddeeff" + ) + + // groupEvent is a seeded billing_event with group attribution set. + type groupEvent struct { + auth string + group string // billing_event.group_id (the token's own group) + members []string // billing_event.member_group_ids + prompt int + cached int + completion int + usage bool + } + events := []groupEvent{ + // 0: a GROUP TOKEN's traffic attributes to the token's own group. + {auth: "a-grp", group: g1, prompt: 100, cached: 30, completion: 50, usage: true}, + // 1: a USER TOKEN belonging to two groups attributes to BOTH. + {auth: "a-user", members: []string{g1, g2}, prompt: 10, usage: true}, + // 2: WITHHELD (no authoritative usage): money excludes it, so must the + // attribution — even though it names g1. + {auth: "a-dead", members: []string{g1}, usage: false}, + // 3: an event whose own group ALSO appears in its membership list + // attributes to that group exactly ONCE. + {auth: "a-both", group: g1, members: []string{g1}, prompt: 5, usage: true}, + } + for i, e := range events { + _, err := db.ExecContext(ctx, + `INSERT INTO billing_event (request_id, auth_id, group_id, member_group_ids, resource_id, model, serving_mode, usage_found, prompt_tokens, cached_tokens, completion_tokens, event_ts) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12)`, + fmt.Sprintf("req-%d", i), nullableStr(e.auth), nullableStr(e.group), groupIDsOrNil(e.members), + "r", "b", "dedicated", e.usage, e.prompt, e.cached, e.completion, hour.Add(time.Duration(i)*time.Minute)) + if err != nil { + t.Fatalf("seed event %d: %v", i, err) + } + } + + // Oracle per-event costs (the same events the money rollup bills). + moneyEvents := []RatedEvent{ + {AuthID: events[0].auth, ResourceID: "r", ModelID: "b", ServingMode: "dedicated", PromptTokens: 100, CachedTokens: 30, CompletionTokens: 50}, + {AuthID: events[1].auth, ResourceID: "r", ModelID: "b", ServingMode: "dedicated", PromptTokens: 10}, + {AuthID: events[3].auth, ResourceID: "r", ModelID: "b", ServingMode: "dedicated", PromptTokens: 5}, + } + costs := make([]Dec, len(moneyEvents)) + for i, e := range moneyEvents { + costs[i] = Rate(e, billed) + } + + store := NewPostgresStore(db) + res, err := store.RateWindow(ctx, book, hour, hour.Add(time.Hour)) + if err != nil { + t.Fatalf("RateWindow: %v", err) + } + + // Money is written exactly once per window, untouched by the attribution: + // three priced events → three rated_usage rows, and the window total is the + // sum of the three per-event costs (no group double-count in the money). + if res.RollupsWritten != 3 { + t.Fatalf("money rollups = %d, want 3 (the withheld event never enters money)", res.RollupsWritten) + } + wantTotal := costs[0].Add(costs[1]).Add(costs[2]).String() + if MustDec(res.TotalCost).String() != wantTotal { + t.Fatalf("total cost = %s, want %s (attribution must not touch the money)", res.TotalCost, wantTotal) + } + if res.EventsRated != 3 || res.MissingUsageEvents != 1 { + t.Fatalf("rated/missing = %d/%d, want 3/1", res.EventsRated, res.MissingUsageEvents) + } + if res.GroupRollupsWritten != 2 { + t.Fatalf("group rollups = %d, want 2 (groups %s and %s)", res.GroupRollupsWritten, g1[:8], g2[:8]) + } + + // Attribution rows: g1 carries events 0, 1, 3 (event 2 is withheld, event 3 + // dedupes); g2 carries event 1 only. + type groupRow struct { + prompt, cached, completion, billable, count int64 + cost string + } + readRow := func(gid string) (groupRow, bool) { + var r groupRow + err := db.QueryRowContext(ctx, + `SELECT prompt_tokens, cached_tokens, completion_tokens, billable_prompt_tokens, event_count, cost::text + FROM group_usage WHERE group_id=$1 AND window_start=$2`, gid, hour). + Scan(&r.prompt, &r.cached, &r.completion, &r.billable, &r.count, &r.cost) + if errors.Is(err, sql.ErrNoRows) { + return groupRow{}, false + } + if err != nil { + t.Fatalf("read group_usage %s: %v", gid, err) + } + return r, true + } + + g1WantCost := costs[0].Add(costs[1]).Add(costs[2]).String() + g2WantCost := costs[1].String() + r1, ok := readRow(g1) + if !ok { + t.Fatalf("no group_usage row for %s", g1) + } + if r1.count != 3 || r1.prompt != 115 || r1.cached != 30 || r1.completion != 50 || r1.billable != 85 { + t.Fatalf("group %s row = %+v, want count 3 / prompt 115 / cached 30 / completion 50 / billable 85", g1, r1) + } + if MustDec(r1.cost).String() != g1WantCost { + t.Fatalf("group %s cost = %s, want %s (oracle)", g1, r1.cost, g1WantCost) + } + r2, ok := readRow(g2) + if !ok { + t.Fatalf("no group_usage row for %s", g2) + } + if r2.count != 1 || r2.prompt != 10 || MustDec(r2.cost).String() != g2WantCost { + t.Fatalf("group %s row = %+v, want count 1 / prompt 10 / cost %s (oracle)", g2, r2, g2WantCost) + } + + // Idempotent re-run: same rows, same sums, nothing reconciled away. + res2, err := store.RateWindow(ctx, book, hour, hour.Add(time.Hour)) + if err != nil { + t.Fatalf("re-run: %v", err) + } + if res2.GroupRollupsWritten != 2 || res2.GroupReconciledDeletions != 0 { + t.Fatalf("re-run group rollups/deletions = %d/%d, want 2/0 (idempotent)", res2.GroupRollupsWritten, res2.GroupReconciledDeletions) + } + if r1b, _ := readRow(g1); MustDec(r1b.cost).String() != g1WantCost || r1b.count != 3 { + t.Fatalf("re-run changed the %s row: %+v", g1, r1b) + } + + // Reconcile: the membership event (req-1) vanishes (data-loss drill). Its + // groups' rows shrink accordingly and g2's row is deleted outright. + exec(t, db, `DELETE FROM billing_event WHERE request_id='req-1'`) + res3, err := store.RateWindow(ctx, book, hour, hour.Add(time.Hour)) + if err != nil { + t.Fatalf("re-rate after deletion: %v", err) + } + if _, ok := readRow(g2); ok { + t.Fatalf("group %s row survived the reconcile, want deleted (its only event vanished)", g2) + } + if res3.GroupReconciledDeletions != 1 { + t.Fatalf("group reconciled deletions = %d, want 1 (the stale %s row)", res3.GroupReconciledDeletions, g2[:8]) + } + r1c, ok := readRow(g1) + if !ok { + t.Fatalf("group %s row vanished, want it retained (events 0 and 3 remain)", g1) + } + g1AfterCost := costs[0].Add(costs[2]).String() + if r1c.count != 2 || MustDec(r1c.cost).String() != g1AfterCost { + t.Fatalf("group %s row after reconcile = %+v, want count 2 / cost %s", g1, r1c, g1AfterCost) + } +} + +// groupIDsOrNil binds a membership list as a driver value (nil for none — +// mirrors the drainer's nullGroupIDs so the fixture writes what production +// writes). +func groupIDsOrNil(ids []string) any { + if len(ids) == 0 { + return nil + } + return ids +} diff --git a/internal/rating/rater.go b/internal/rating/rater.go index 6c4a43a..96ff5f9 100644 --- a/internal/rating/rater.go +++ b/internal/rating/rater.go @@ -98,6 +98,11 @@ type Result struct { // than a guessed one. Reported so lost attribution is visible. AmbiguousGraphRollups int64 RollupsWritten int64 // distinct grain rows upserted (see rated_usage_grain_uq) + GroupRollupsWritten int64 // group_usage attribution rows upserted (NOT money — see RateResult.GroupRollupsWritten) + // GroupReconciledDeletions is stale in-window group_usage rows this run + // deleted (same reconcile contract as ReconciledDeletions, at the + // attribution grain). 0 on a first run or a clean identical re-run. + GroupReconciledDeletions int64 // UnratedHours counts hours in the window that were SKIPPED — prices could not // be obtained, or rating them failed. Each hour is independent and the upsert is // idempotent, so a skipped hour is rated by a later run whose trailing window @@ -256,6 +261,8 @@ func (r *Rater) Run(ctx context.Context, windowStart, windowEnd time.Time, windo } res.EventsRated = rr.EventsRated res.RollupsWritten = rr.RollupsWritten + res.GroupRollupsWritten = rr.GroupRollupsWritten + res.GroupReconciledDeletions = rr.GroupReconciledDeletions res.ReconciledDeletions = rr.ReconciledDeletions res.TotalCost = rr.TotalCost res.UnpricedEvents = rr.UnpricedEvents @@ -348,6 +355,11 @@ func (r *Rater) Run(ctx context.Context, windowStart, windowEnd time.Time, windo windowStart.Format(time.RFC3339), windowEnd.Format(time.RFC3339), res.EventsRated, res.RollupsWritten, res.TotalCost) } + if res.GroupRollupsWritten > 0 || res.GroupReconciledDeletions > 0 { + r.log.Info.Printf("rating: window [%s,%s) attributed usage to %d group_usage rollup(s) (membership-aware group quotas — attribution, NOT money; the same event cost legitimately appears under every group it belongs to); %d stale group rollup(s) reconcile-deleted", + windowStart.Format(time.RFC3339), windowEnd.Format(time.RFC3339), + res.GroupRollupsWritten, res.GroupReconciledDeletions) + } return res, nil } @@ -474,6 +486,8 @@ func (r *Rater) RunWindow(ctx context.Context, windowStart, windowEnd time.Time, func (r *Result) accumulate(hour Result) { r.EventsRated += hour.EventsRated r.RollupsWritten += hour.RollupsWritten + r.GroupRollupsWritten += hour.GroupRollupsWritten + r.GroupReconciledDeletions += hour.GroupReconciledDeletions r.ReconciledDeletions += hour.ReconciledDeletions r.UnpricedEvents += hour.UnpricedEvents r.UnattributableEvents += hour.UnattributableEvents diff --git a/internal/rating/store.go b/internal/rating/store.go index 82300e6..1de136f 100644 --- a/internal/rating/store.go +++ b/internal/rating/store.go @@ -41,7 +41,18 @@ type RateResult struct { // wide window can exceed 2^31. Widened with the ::bigint SQL casts to avoid a // silent 32-bit overflow. RollupsWritten int64 - EventsRated int64 + // GroupRollupsWritten is how many group_usage attribution rows this run + // upserted (membership-aware group quotas, ruled 2026-10-07). ATTRIBUTION, + // not money: the same event cost legitimately appears under every group it + // is attributed to, so this count (and the group cost sums) must never be + // read as billing volume. rated_usage's money write count is unchanged. + GroupRollupsWritten int64 + // GroupReconciledDeletions is how many in-window group_usage rows this run + // DELETED as stale — the same "what the latest run says is what bills" + // reconcile rated_usage performs, at the attribution grain. 0 on a first + // run or a clean identical re-run. + GroupReconciledDeletions int64 + EventsRated int64 // ReconciledDeletions is how many stale rated_usage rows this run DELETED because // they billed in a prior run but fell out of the current priced set (re-rate // convergence — "what the latest run says is what bills"). 0 on a first run or a @@ -214,7 +225,12 @@ CREATE TEMP TABLE rating_derived ( ) ON COMMIT DROP` // rateWindowSQL resolves, sums, upserts, and counts in ONE statement over the -// transient rating_price table (populated from the YAML PriceBook for this run). +// transient rating_price table (populated from the YAML PriceBook for this +// run). It writes TWO rollups from the same snapshot: the money rollup +// (rated_usage — written exactly once per window, schema and grain unchanged) +// and the group attribution rollup (group_usage, migration 0008 — the same +// events attributed to each group that counts against them; attribution, not +// money). // // RESOLUTION (the C4 ladder, precedence a > b > c > d; the Go mirror is // PriceBook.ResolveEvent): vLLM serves Token Factory endpoints under the ENDPOINT @@ -360,6 +376,18 @@ WITH ev AS ( -- owner_conflict: both a user AND a group on one event. Carried so the grouped -- gate below can withhold the whole rollup rather than bill a guessed owner. (COALESCE(user_id, '') <> '' AND COALESCE(group_id, '') <> '') AS owner_conflict, + -- GROUP ATTRIBUTION (membership-aware group quotas, ruled 2026-10-07): + -- the token's own group (group-token case) plus the caller's memberships + -- (X-Saturn-Group-Scopes). Both feed the group_usage rollup below; + -- neither enters the money grain — rated_usage is unchanged. + group_id, + member_group_ids, + -- request_id: the event id, carried ONLY so the attribution dedupe can + -- key on it (DISTINCT ON (request_id, gid): one event whose own group + -- also appears in member_group_ids must attribute once, and without a + -- unique event key two genuinely different events with identical + -- counts and cost would collapse into one). + request_id, -- graph_k8s_name: the DynamoGraphDeployment that served the request -- the cost -- centre. Carried as EVIDENCE onto the rollup, never a grain key: a shared graph -- serves many orgs and has no database row, so without it shared-mode cost is @@ -400,6 +428,9 @@ resolved AS ( ev.owner_type, ev.owner_id, ev.owner_conflict, + ev.group_id, + ev.member_group_ids, + ev.request_id, ev.graph_k8s_name, ev.ev_ts, ev.usage_found, @@ -688,8 +719,7 @@ upserted AS ( -- the cross-rater hazard is unreachable and no delete-lock-ordering machinery is -- added here. See cmd/rater's package doc for the single-flight contract. ORDER BY auth_id, owner_type, owner_id, resource_id, model_id, serving_mode, window_start - ON CONFLICT (auth_id, owner_type, owner_id, resource_id, model_id, serving_mode, window_start) DO UPDATE SET - -- Refresh org_id on re-rate, but NEVER erase a known org: COALESCE prefers the + ON CONFLICT (auth_id, owner_type, owner_id, resource_id, model_id, serving_mode, window_start) DO UPDATE SET -- Refresh org_id on re-rate, but NEVER erase a known org: COALESCE prefers the -- new snapshot's org and FALLS BACK to the existing row's org when the new one is -- NULL. So a rollup first written with a NULL org (header not yet wired) picks up -- the real org on a later re-rate (NULL -> real, convergence), but a re-rate over @@ -734,10 +764,124 @@ upserted AS ( event_count = EXCLUDED.event_count, rated_at = now() RETURNING event_count, cost +), +-- --------------------------------------------------------------------------- +-- GROUP ATTRIBUTION (group_usage; membership-aware group quotas, ruled +-- 2026-10-07): while the same statement writes the window's money into +-- rated_usage, it ALSO upserts the per-(group, hour) attribution rollup the +-- admission group spend check reads. Attribution, not money: an event +-- attributed to N groups contributes to N group rows (the same event cost +-- appears under each), and WITHHELD (non-money) events contribute nothing — +-- an event reaches attributed only if it also entered a priced rollup, so the +-- rollup's rows and the money always agree on what was served. rated_usage's +-- schema, grain, and write count are untouched: money is still written +-- exactly once. +-- --------------------------------------------------------------------------- +attributed AS ( + -- Per (event, group) rows for the events that entered money. priced + -- already excludes the rollup-level ambiguity gates, so joining resolved + -- back onto priced on the FULL grain re-attaches exactly the events of + -- money rollups. The per-event money gates are re-applied on the resolved + -- side: a failing event of the same grain (an unpriced/invalid twin in an + -- otherwise priced rollup) must not sneak in with its rollup. DISTINCT ON + -- (request_id, gid) dedupes the one legitimate duplication — an event + -- whose own group ALSO appears in member_group_ids — while the event id + -- keeps two genuinely different but identical-looking events distinct. + SELECT DISTINCT ON (pe.request_id, g.gid) + g.gid AS group_id, + p.window_start AS window_start, + pe.prompt_tokens AS prompt_tokens, + pe.cached_tokens AS cached_tokens, + pe.completion_tokens AS completion_tokens, + pe.billable_prompt AS billable_prompt_tokens, + (pe.billable_prompt * pe.prompt_price + + pe.cached_tokens * pe.cached_price + + pe.completion_tokens * pe.completion_price) AS event_cost + FROM priced p + JOIN resolved pe + ON pe.auth_id = p.auth_id + AND pe.owner_type = p.owner_type + AND pe.owner_id = p.owner_id + AND pe.resource_id = p.resource_id + AND pe.model_id = p.model_id + AND pe.serving_mode = p.serving_mode + AND date_trunc('hour', pe.ev_ts AT TIME ZONE 'UTC') AT TIME ZONE 'UTC' = p.window_start + CROSS JOIN LATERAL ( + -- The token's own group (group-token case) PLUS each membership gid. + -- COALESCE both operands: NULL || anything is NULL in Postgres, and an + -- event with neither source simply attributes to no group. The empty- + -- string guard is belt-and-braces: the envelope parser never emits one. + SELECT gid + FROM unnest(COALESCE(pe.member_group_ids, '{}'::text[]) + || COALESCE(ARRAY[pe.group_id], '{}'::text[])) AS u(gid) + WHERE gid IS NOT NULL AND gid <> '' + ) g + WHERE pe.usage_found + AND pe.valid_usage + AND pe.prompt_price IS NOT NULL + AND pe.auth_id IS NOT NULL + AND pe.resource_id IS NOT NULL + AND pe.model_id IS NOT NULL + AND NOT pe.owner_conflict + AND pe.valid_serving_mode +), +group_rollup AS ( + SELECT + group_id, + window_start, + SUM(prompt_tokens)::bigint AS prompt_tokens, + SUM(cached_tokens)::bigint AS cached_tokens, + SUM(completion_tokens)::bigint AS completion_tokens, + SUM(billable_prompt_tokens)::bigint AS billable_prompt_tokens, + SUM(event_cost) AS cost, + COUNT(*)::bigint AS event_count + FROM attributed + GROUP BY group_id, window_start +), +group_deleted AS ( + -- Reconcile, the SAME contract as rated_usage's deleted CTE: an in-window + -- group_usage row this run does NOT reproduce is stale (its events became + -- withheld, or the window re-rated under different attribution), and a + -- clean re-run reproduces every row so nothing matches. Single-flight and + -- hour-aligned windows (cmd/rater's contract) make the delete safe. + DELETE FROM group_usage gu + WHERE gu.window_start >= $1 + AND gu.window_start < $2 + AND NOT EXISTS ( + SELECT 1 FROM group_rollup gr + WHERE gr.group_id = gu.group_id + AND gr.window_start = gu.window_start + ) + RETURNING gu.group_id +), +group_upserted AS ( + INSERT INTO group_usage ( + group_id, window_start, prompt_tokens, cached_tokens, completion_tokens, + billable_prompt_tokens, cost, event_count + ) + SELECT + group_id, window_start, prompt_tokens, cached_tokens, completion_tokens, + billable_prompt_tokens, cost, event_count + FROM group_rollup + -- Same deterministic-order discipline as the rated_usage upsert. + ORDER BY group_id, window_start + ON CONFLICT (group_id, window_start) DO UPDATE SET + prompt_tokens = EXCLUDED.prompt_tokens, + cached_tokens = EXCLUDED.cached_tokens, + completion_tokens = EXCLUDED.completion_tokens, + billable_prompt_tokens = EXCLUDED.billable_prompt_tokens, + cost = EXCLUDED.cost, + event_count = EXCLUDED.event_count + RETURNING group_id ) SELECT (SELECT COUNT(*)::bigint FROM upserted) AS rollups_written, (SELECT COALESCE(SUM(event_count), 0)::bigint FROM upserted) AS events_rated, + -- Group attribution rollup (group_usage), from the same snapshot: rows + -- upserted, and stale in-window rows the reconcile deleted. Attribution, + -- not money — group_rollups_total is NOT added to total_cost. + (SELECT COUNT(*)::bigint FROM group_upserted) AS group_rollups_written, + (SELECT COUNT(*)::bigint FROM group_deleted) AS group_reconciled_deletions, (SELECT COALESCE(SUM(cost), 0)::numeric FROM upserted) AS total_cost, -- Stale rollups DELETED by the reconcile (re-rate convergence). Rows that billed -- in a prior run but fell out of priced this run; surfaced so a re-rate that @@ -914,7 +1058,7 @@ func (s *PostgresStore) RateWindow(ctx context.Context, book *PriceBook, start, var res RateResult var total string err = tx.QueryRowContext(ctx, rateWindowSQL, start.UTC(), end.UTC(), ftLikePattern). - Scan(&res.RollupsWritten, &res.EventsRated, &total, &res.ReconciledDeletions, + Scan(&res.RollupsWritten, &res.EventsRated, &res.GroupRollupsWritten, &res.GroupReconciledDeletions, &total, &res.ReconciledDeletions, &res.UnpricedEvents, &res.UnattributableEvents, &res.MissingUsageEvents, &res.ExpectedMissingUsageEvents, &res.UnexplainedMissingUsageEvents, &res.InvalidUsageEvents, &res.AmbiguousBaseEvents, diff --git a/internal/rating/store_integration_test.go b/internal/rating/store_integration_test.go index 7e53e41..0b4dd01 100644 --- a/internal/rating/store_integration_test.go +++ b/internal/rating/store_integration_test.go @@ -59,6 +59,9 @@ func ratingSchemaDDL(t *testing.T) string { // 0007 makes rated_usage.serving_mode 'shared'/'dedicated' only (CHECK, no // default), which the rater's upsert must satisfy. "../../migrations/0007_serving_mode_explicit.up.sql", + // 0008 adds billing_event.member_group_ids (read by the rater's group + // attribution) and the group_usage rollup it writes. + "../../migrations/0008_group_scopes.up.sql", } { ddl, err := os.ReadFile(f) if err != nil { @@ -942,6 +945,10 @@ func TestIntegration_InvalidUsageEvidenceNeverEntersMoney(t *testing.T) { "0004_billing_event_serving_mode.up.sql", "0005_invoice_grade_attempts.up.sql", "0006_rollup_grain.up.sql", + // 0008 adds billing_event.member_group_ids (read by the rater's group + // attribution) and group_usage (written by it) — the current rater + // requires both, like 0006's graph_k8s_name before it. + "0008_group_scopes.up.sql", } { apply(name) } @@ -2106,6 +2113,7 @@ func TestIntegration_FreshInputTokensSurvivesInt32Overflow(t *testing.T) { "0004_billing_event_serving_mode.up.sql", "0005_invoice_grade_attempts.up.sql", "0006_rollup_grain.up.sql", + "0008_group_scopes.up.sql", } { ddl, readErr := os.ReadFile("../../migrations/" + name) if readErr != nil { @@ -2216,6 +2224,7 @@ func TestIntegration_MissingUsagePartitionedByCause(t *testing.T) { "0004_billing_event_serving_mode.up.sql", "0005_invoice_grade_attempts.up.sql", "0006_rollup_grain.up.sql", + "0008_group_scopes.up.sql", } { ddl, readErr := os.ReadFile("../../migrations/" + name) if readErr != nil { @@ -3235,6 +3244,10 @@ func TestIntegration_Migration0007ServingModeExplicit(t *testing.T) { exec(t, db, "UPDATE rated_usage SET id = "+idExpr("")+" WHERE auth_id = 'a-ded'") exec(t, db, readMigration(t, "0007_serving_mode_explicit.up.sql")) + // The current rater also reads member_group_ids / writes group_usage + // (0008): apply it before the first RateWindow below. 0008 is independent + // of everything 0007 asserts. + exec(t, db, readMigration(t, "0008_group_scopes.up.sql")) rows := map[string]string{} r, err := db.QueryContext(ctx, `SELECT auth_id || '/' || serving_mode, id FROM rated_usage`) @@ -3482,6 +3495,9 @@ func TestIntegration_Migration0007BackfillMakesPreCutoverReRateANoOp(t *testing. } exec(t, db, readMigration(t, "0007_serving_mode_explicit.up.sql")) + // The current rater reads member_group_ids / writes group_usage (0008) — + // apply it before the re-rate below (see TestIntegration_Migration0007ServingModeExplicit). + exec(t, db, readMigration(t, "0008_group_scopes.up.sql")) var leftover int if err := db.QueryRowContext(ctx, @@ -3577,7 +3593,7 @@ func TestIntegration_RolloutReRunBackfillLeavesExplicitEmptyServingModeWithheld( exec(t, db, "SET search_path TO "+sch) defer func() { exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") }() for _, f := range []string{"0001_billing_event", "0002_rating", "0004_billing_event_serving_mode", - "0005_invoice_grade_attempts", "0006_rollup_grain", "0007_serving_mode_explicit"} { + "0005_invoice_grade_attempts", "0006_rollup_grain", "0007_serving_mode_explicit", "0008_group_scopes"} { exec(t, db, readMigration(t, f+".up.sql")) } exec(t, db, "ALTER TABLE billing_event ALTER COLUMN usage_found SET DEFAULT TRUE") @@ -3680,6 +3696,9 @@ func TestIntegration_Migration0007TwinHourIsWholeAfterExplicitReRate(t *testing. } exec(t, db, readMigration(t, "0007_serving_mode_explicit.up.sql")) + // The current rater reads member_group_ids / writes group_usage (0008) — + // apply it before the re-rates below. + exec(t, db, readMigration(t, "0008_group_scopes.up.sql")) var partial int64 if err := db.QueryRowContext(ctx, `SELECT SUM(event_count) FROM rated_usage WHERE window_start = $1`, hour). diff --git a/internal/rating/store_test.go b/internal/rating/store_test.go index dd682d6..5afef54 100644 --- a/internal/rating/store_test.go +++ b/internal/rating/store_test.go @@ -46,8 +46,8 @@ func TestPostgresStore_RateWindowSQL(t *testing.T) { // missing_usage_events is the TOTAL; expected + unexplained partition it // (6 = 4 routine aborts/failures + 2 successes with no usage block), so the // fixture cannot pass while the SQL's partition is wrong. - rows := sqlmock.NewRows([]string{"rollups_written", "events_rated", "total_cost", "reconciled_deletions", "unpriced_events", "unattributable_events", "missing_usage_events", "expected_missing_usage_events", "unexplained_missing_usage_events", "invalid_usage_events", "ambiguous_base_events", "ambiguous_org_events", "owner_conflict_events", "ambiguous_graph_rollups", "invalid_serving_mode_events"}). - AddRow(2, 5, "0.001234500", 0, 3, 1, 6, 4, 2, 7, 4, 2, 3, 1, 5) + rows := sqlmock.NewRows([]string{"rollups_written", "events_rated", "group_rollups_written", "group_reconciled_deletions", "total_cost", "reconciled_deletions", "unpriced_events", "unattributable_events", "missing_usage_events", "expected_missing_usage_events", "unexplained_missing_usage_events", "invalid_usage_events", "ambiguous_base_events", "ambiguous_org_events", "owner_conflict_events", "ambiguous_graph_rollups", "invalid_serving_mode_events"}). + AddRow(2, 5, 3, 1, "0.001234500", 0, 3, 1, 6, 4, 2, 7, 4, 2, 3, 1, 5) // The statement binds $3 = the ft: LIKE pattern (single-sourced from fineTunePrefix). mock.ExpectQuery(`INSERT INTO rated_usage`). WithArgs(start.UTC(), end.UTC(), ftLikePattern). @@ -92,6 +92,9 @@ func TestPostgresStore_RateWindowSQL(t *testing.T) { if res.ReconciledDeletions != 0 { t.Fatalf("reconciled deletions = %d, want 0 (the projected count must scan into the result)", res.ReconciledDeletions) } + if res.GroupRollupsWritten != 3 || res.GroupReconciledDeletions != 1 { + t.Fatalf("group attribution = %d rollups/%d deletions, want 3/1 (the projected counts must scan into the result)", res.GroupRollupsWritten, res.GroupReconciledDeletions) + } if err := mock.ExpectationsWereMet(); err != nil { t.Fatalf("unmet: %v", err) } @@ -256,6 +259,24 @@ func TestRateWindowSQL_Shape(t *testing.T) { // a both-ambiguous rollup counts only as base) "WHERE ambiguous_org AND NOT ambiguous_base", "AS ambiguous_org_events", + // GROUP ATTRIBUTION (group_usage, migration 0008; membership-aware group + // quotas, ruled 2026-10-07): the money events re-attached to their groups + // (token group_id + member_group_ids), deduped per (event, group), summed + // per (group, hour), upserted with the SAME reconcile contract as + // rated_usage. Attribution rides the same statement and snapshot. + "SELECT DISTINCT ON (pe.request_id, g.gid)", + "COALESCE(pe.member_group_ids, '{}'::text[])", + "COALESCE(ARRAY[pe.group_id], '{}'::text[])", + "FROM unnest(", + "JOIN resolved pe", + "AND date_trunc('hour', pe.ev_ts AT TIME ZONE 'UTC') AT TIME ZONE 'UTC' = p.window_start", + "GROUP BY group_id, window_start", + "INSERT INTO group_usage (", + "ON CONFLICT (group_id, window_start) DO UPDATE SET", + "DELETE FROM group_usage gu", + "FROM group_rollup gr", + "AS group_rollups_written", + "AS group_reconciled_deletions", } // The price tables are GONE (prices are YAML now): no reference to model_price, // derivation_policy, effective-dating, or a derivation CASE may remain. From c67ecf02365504c2426352d002dc5a9dc797a3f6 Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 07:25:52 +0000 Subject: [PATCH 05/17] admission: group rate scopes + monthly spend caps (membership-aware quotas) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Request gains GroupScopes: per group, the four per-minute windows run as contract scopes through the existing Lua/counter machinery and settle via the lease path (R4 sentinels end to end), and the monthly spend cap compares SUM(cost) over group_usage in Postgres — cached ~60s per (group, cap), compared in SQL so money never becomes a Go number. Spend >= cap rejects contractually (a zero cap rejects all paid work); store errors fail open with a throttled loud log (2026-09-24 posture); a scope without its group id fails closed like every broken-envelope case. --- internal/admission/admission.go | 126 ++++++++++++ internal/admission/group_test.go | 323 +++++++++++++++++++++++++++++++ internal/admission/spend.go | 183 +++++++++++++++++ 3 files changed, 632 insertions(+) create mode 100644 internal/admission/group_test.go create mode 100644 internal/admission/spend.go diff --git a/internal/admission/admission.go b/internal/admission/admission.go index 97ad289..33d4093 100644 --- a/internal/admission/admission.go +++ b/internal/admission/admission.go @@ -19,6 +19,7 @@ import ( "github.com/redis/go-redis/v9" "github.com/saturncloud/phoebe/internal/config" + "github.com/saturncloud/phoebe/internal/logging" ) var ErrUnavailable = errors.New("distributed admission state unavailable") @@ -82,6 +83,26 @@ type Request struct { Adapter bool OrganizationLimits RateLimits OwnerLimits RateLimits + // GroupScopes is the parsed membership-aware group quota envelope + // (X-Saturn-Group-Scopes, ruled 2026-10-07): one entry per group the + // caller belongs to that carries limits. Each entry's Limits enforce the + // four per-minute rate windows as a contract scope (the same Lua/counter + // machinery as the contract_owner scope, settling through the same + // lease/CompleteUsage path); SpendCap, when set, enforces the group's + // monthly spend cap against the rater's group_usage rollup. Empty when the + // envelope carried no group (the common case). + GroupScopes []GroupScope +} + +// GroupScope is one group's quota contract within a Request: the group's id, +// its four per-minute rate limits (R4 sentinels: nil unlimited, 0 zero cap), +// and its monthly spend cap as a plain decimal NUMERIC(20,9) string straight +// from the trusted envelope ("" = no cap). The cap is compared in Postgres, +// never as a Go number. +type GroupScope struct { + GroupID string + Limits RateLimits + SpendCap string } // RateLimits is the authenticated customer contract. R4 sentinel semantics @@ -135,6 +156,28 @@ type RedisAdmitter struct { cfg config.AdmissionSettings counters, leases, expiries string windowExpiries string + // Group scope enforcement (membership-aware group quotas, ruled + // 2026-10-07). spendStore is nil on installs without a Postgres handle + // (serving-only spokes): group RATE limits still enforce from the Valkey + // store, spend caps are then unchecked — logged loudly at startup by + // cmd/interceptor, never silently assumed. + spendStore GroupSpendStore + spendCache *spendVerdictCache + spendLog *logging.Logger + spendFails *spendFailureLog +} + +// WithGroupSpend enables the monthly spend cap check against a GroupSpendStore +// and returns the admitter for chaining. log carries the loud-failure lines: +// the spend check FAILS OPEN (a quota may never take inference down), so a +// store error is always an ERROR line, throttled by spendFailureLog so a +// Postgres outage cannot emit one line per admitted request. +func (a *RedisAdmitter) WithGroupSpend(store GroupSpendStore, log *logging.Logger) *RedisAdmitter { + a.spendStore = store + a.spendLog = log + a.spendCache = newSpendVerdictCache() + a.spendFails = &spendFailureLog{} + return a } func New(client redis.Cmdable, cfg config.AdmissionSettings) *RedisAdmitter { @@ -199,9 +242,25 @@ func (a *RedisAdmitter) Admit(ctx context.Context, req Request) (*Lease, error) if req.OwnerLimits.Any() && req.Owner == "" { return nil, fmt.Errorf("%w: missing trusted owner identity", ErrInvalidIdentity) } + for _, gs := range req.GroupScopes { + // A group scope carrying any limit or spend cap must name its group: + // an anonymous quota is not a quota, it is a broken envelope (fail + // closed, like every other structural violation). + if gs.GroupID == "" && (gs.Limits.Any() || gs.SpendCap != "") { + return nil, fmt.Errorf("%w: group scope carries limits without a group id", ErrInvalidIdentity) + } + } if req.PromptBytes < 0 || req.EstimatedInputTokens <= 0 || req.ReservedOutputTokens <= 0 { return nil, &Rejected{Scope: "request", Dimension: "work estimate", RetryAfter: time.Second} } + // The monthly spend cap is a Go-side pre-check against Postgres (cached + // per (group, cap), compared in SQL), ahead of the Lua reservation: a group + // at its cap is rejected without reserving anything. It fails OPEN — a + // store error logs loud and admits, the 2026-09-24 posture (quotas are + // permissible; they never take inference down). + if err := a.checkGroupSpend(ctx, req.GroupScopes); err != nil { + return nil, err + } id, err := randomID() if err != nil { return nil, fmt.Errorf("%w: lease id: %v", ErrUnavailable, err) @@ -385,9 +444,76 @@ func (a *RedisAdmitter) scopes(r Request) []scope { if r.OwnerLimits.Any() { out = append(out, makeContractScope("contract_owner", r.Owner, r.OwnerLimits)) } + // Group rate scopes: one contract scope per group that carries any rate + // limit, same Lua/counter machinery and wireLimit sentinels as the org/ + // owner contract scopes, settling through the same lease/CompleteUsage + // path (the lease record holds this scope like any other, and finishScript + // charges its windows). The scope NAME carries the group id so a rejection + // names the team; the scope id digests it, so groups never collide. + for _, gs := range r.GroupScopes { + if gs.Limits.Any() { + out = append(out, makeContractScope("group:"+gs.GroupID, gs.GroupID, gs.Limits)) + } + } return out } +// checkGroupSpend enforces the monthly spend cap for every group scope that +// carries one: reject when the group's month-to-date attribution spend (the +// rater's group_usage rollup) has reached the cap. Reaching is the boundary — +// spend >= cap rejects, so a cap of 0 rejects all paid work for the group. +// A 429+Retry-After denial (Contractual); the retry hint is one minute, the +// cache TTL: the verdict can refresh no sooner anyway, and spend only relaxes +// when the month turns or a re-rate supersedes the rollup. +// +// FAIL OPEN, LOUD: a missing store (install without Postgres) or a store +// error admits the request and logs. Quotas are permissible; they never take +// inference down, and they are never silently enforced either — every bypass +// is an ERROR line (throttled so an outage cannot flood). +func (a *RedisAdmitter) checkGroupSpend(ctx context.Context, scopes []GroupScope) error { + if a.spendStore == nil { + for _, gs := range scopes { + if gs.SpendCap != "" { + a.spendBypass(gs.GroupID, errors.New("no spend store configured (group_usage is unreachable — is DATABASE_URL set?)")) + } + } + return nil + } + for _, gs := range scopes { + if gs.SpendCap == "" { + continue + } + if exhausted, ok := a.spendCache.lookup(gs.GroupID, gs.SpendCap); ok { + if exhausted { + return &Rejected{Scope: "group:" + gs.GroupID, Dimension: "monthly_spend", RetryAfter: groupSpendCacheTTL, Contractual: true} + } + continue + } + exhausted, err := a.spendStore.GroupSpendExhausted(ctx, gs.GroupID, gs.SpendCap) + if err != nil { + a.spendBypass(gs.GroupID, err) + continue + } + a.spendCache.store(gs.GroupID, gs.SpendCap, exhausted) + if exhausted { + return &Rejected{Scope: "group:" + gs.GroupID, Dimension: "monthly_spend", RetryAfter: groupSpendCacheTTL, Contractual: true} + } + } + return nil +} + +// spendBypass logs a fail-open spend-check bypass at ERROR, throttled through +// spendFailureLog so a sustained store outage stays visible at its onset +// without one line per admitted request. Nil-safe for admitters built without +// WithGroupSpend (a spendStore-less admitter can still see spend caps if a +// caller hands it group scopes directly). +func (a *RedisAdmitter) spendBypass(groupID string, err error) { + if a.spendLog == nil || a.spendFails == nil { + return + } + a.spendFails.logf(a.spendLog, "admission: group %s: monthly spend cap NOT enforced (fail open): %v", groupID, err) +} + // Lease owns all reservations for one accepted request. Its transition methods // are idempotent locally and atomically idempotent in Valkey. type Lease struct { diff --git a/internal/admission/group_test.go b/internal/admission/group_test.go new file mode 100644 index 0000000..7b897f8 --- /dev/null +++ b/internal/admission/group_test.go @@ -0,0 +1,323 @@ +package admission + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/saturncloud/phoebe/internal/config" + "github.com/saturncloud/phoebe/internal/logging" +) + +// fakeGroupSpendStore is the spend-check seam test double: verdicts and errors +// are scripted per call so a test can observe exactly how often Postgres would +// be read (the cache-behavior tests pin the once-per-TTL contract). +type fakeGroupSpendStore struct { + exhausted bool + err error + calls int + gotGID []string + gotCap []string +} + +func (f *fakeGroupSpendStore) GroupSpendExhausted(_ context.Context, groupID, cap string) (bool, error) { + f.calls++ + f.gotGID = append(f.gotGID, groupID) + f.gotCap = append(f.gotCap, cap) + return f.exhausted, f.err +} + +// spendTestAdmitter builds an admitter over miniredis with a fake spend store +// and a silent logger. +func spendTestAdmitter(t *testing.T) (*RedisAdmitter, *fakeGroupSpendStore) { + t.Helper() + a, _ := testAdmitter(t, config.AdmissionSettings{}) + fake := &fakeGroupSpendStore{} + a.WithGroupSpend(fake, logging.New(logging.ERROR)) + return a, fake +} + +const testGID = "a1b2c3d4e5f60718293a4b5c6d7e8f90" + +// groupScope builds one GroupScope. rate of nil leaves every rate unset +// (unlimited); a non-nil rate installs all four dimensions to the same value. +func groupScope(gid string, rate *int64, spend string) GroupScope { + gs := GroupScope{GroupID: gid, SpendCap: spend} + if rate != nil { + gs.Limits = RateLimits{Requests: rate, TotalPromptTokens: rate, + UncachedPromptTokens: rate, GeneratedTokens: rate} + } + return gs +} + +func groupRequest(org, model string, scopes ...GroupScope) Request { + r := request(org, model) + r.GroupScopes = scopes + return r +} + +// TestGroupRateLimitsEnforcedLikeContractScope: a group's per-minute windows +// are the same Lua/counter machinery as the contract scopes — a burst up to +// the cap admits, the next request is a contractual rejection (the proxy maps +// that to 429 + Retry-After). +func TestGroupRateLimitsEnforcedLikeContractScope(t *testing.T) { + a, _ := spendTestAdmitter(t) + // Only the requests window is capped, so the request's own estimates + // (the helper's 10-in/20-out shape) can never trip the unsatisfiable + // check: the third request is rejected by the window counter itself. + cap := int64(2) + scope := GroupScope{GroupID: testGID, Limits: RateLimits{Requests: &cap}} + + for i := 0; i < 2; i++ { + lease, err := a.Admit(context.Background(), groupRequest("org-g", "m", scope)) + if err != nil { + t.Fatalf("admit %d: %v, want admitted", i+1, err) + } + _ = lease.Complete(context.Background(), 0) + } + _, err := a.Admit(context.Background(), groupRequest("org-g", "m", scope)) + var rejected *Rejected + if !errors.As(err, &rejected) { + t.Fatalf("third admit: %v, want a contract rejection", err) + } + if !rejected.Contractual || rejected.RetryAfter <= 0 { + t.Fatalf("rejection = %+v, want Contractual with a Retry-After (429 mapping)", rejected) + } + if rejected.Scope != "group:"+testGID { + t.Fatalf("rejection scope = %q, want the group id named", rejected.Scope) + } +} + +// TestGroupRateZeroCapBlocksEveryRequest: an explicit 0 is a zero cap, not +// unlimited — the same R4 sentinel as every other contract scope. +func TestGroupRateZeroCapBlocksEveryRequest(t *testing.T) { + a, _ := spendTestAdmitter(t) + zero := int64(0) + _, err := a.Admit(context.Background(), groupRequest("org-z", "m", + groupScope(testGID, &zero, ""))) + var rejected *Rejected + if !errors.As(err, &rejected) || !rejected.Contractual { + t.Fatalf("zero-capped group admit: %v, want a contractual rejection", err) + } +} + +// TestGroupRateNilLimitIsUnlimited: a group scope whose rate fields are all +// nil contributes no contract scope (unlimited), so requests flow freely. +func TestGroupRateNilLimitIsUnlimited(t *testing.T) { + a, _ := spendTestAdmitter(t) + for i := 0; i < 3; i++ { + lease, err := a.Admit(context.Background(), groupRequest("org-u", "m", + groupScope(testGID, nil, ""))) + if err != nil { + t.Fatalf("admit %d: %v, want admitted (nil limits are unlimited)", i+1, err) + } + _ = lease.Complete(context.Background(), 0) + } +} + +// TestGroupSpendCapDeniesAtCap: spend >= cap rejects (reaching the cap IS the +// denial boundary), the denial is contractual (429 + Retry-After at the +// proxy), and the cap string reaches the store verbatim. +func TestGroupSpendCapDeniesAtCap(t *testing.T) { + a, store := spendTestAdmitter(t) + store.exhausted = true + + _, err := a.Admit(context.Background(), groupRequest("org-s", "m", + groupScope(testGID, nil, "100.000000000"))) + var rejected *Rejected + if !errors.As(err, &rejected) { + t.Fatalf("capped group admit: %v, want a rejection", err) + } + if !rejected.Contractual || rejected.Dimension != "monthly_spend" || rejected.RetryAfter <= 0 { + t.Fatalf("rejection = %+v, want Contractual monthly_spend with Retry-After (429 mapping)", rejected) + } + if store.gotGID[0] != testGID || store.gotCap[0] != "100.000000000" { + t.Fatalf("store bound (gid, cap) = (%q, %q), want the envelope values verbatim", store.gotGID[0], store.gotCap[0]) + } +} + +// TestGroupSpendZeroCapDeniesAllPaidWork: the rater-side predicate is +// spend >= cap, and a zero cap is reached by every request — zero spend is +// still >= 0. +func TestGroupSpendZeroCapDeniesAllPaidWork(t *testing.T) { + a, store := spendTestAdmitter(t) + store.exhausted = true // the SQL verdict at a zero cap: COALESCE(SUM,0) >= 0 + _, err := a.Admit(context.Background(), groupRequest("org-s", "m", + groupScope(testGID, nil, "0"))) + var rejected *Rejected + if !errors.As(err, &rejected) || !rejected.Contractual { + t.Fatalf("zero-cap group admit: %v, want a contractual rejection", err) + } +} + +// TestGroupSpendCapAdmitsBelowCap: under the cap the request admits, and the +// verdict — not the spend — is what the store returns. +func TestGroupSpendCapAdmitsBelowCap(t *testing.T) { + a, store := spendTestAdmitter(t) + lease, err := a.Admit(context.Background(), groupRequest("org-b", "m", + groupScope(testGID, nil, "100"))) + if err != nil { + t.Fatalf("under-cap admit: %v, want admitted", err) + } + _ = lease.Complete(context.Background(), 0) + if store.calls != 1 { + t.Fatalf("store calls = %d, want 1", store.calls) + } +} + +// TestGroupSpendVerdictCachedPerGroupAndCap: the verdict is cached ~60s per +// (group, cap) — a second request must not re-read Postgres, and a different +// cap for the same group is a different key. +func TestGroupSpendVerdictCachedPerGroupAndCap(t *testing.T) { + a, store := spendTestAdmitter(t) + scope := groupScope(testGID, nil, "100") + ctx := context.Background() + + if _, err := a.Admit(ctx, groupRequest("org-c", "m", scope)); err != nil { + t.Fatalf("first admit: %v", err) + } + if _, err := a.Admit(ctx, groupRequest("org-c", "m", scope)); err != nil { + t.Fatalf("second admit: %v", err) + } + if store.calls != 1 { + t.Fatalf("store calls = %d, want 1 (the verdict is cached per (group, cap))", store.calls) + } + + // A different cap is a different cache key: one more read, then cached. + scope.SpendCap = "50" + if _, err := a.Admit(ctx, groupRequest("org-c", "m", scope)); err != nil { + t.Fatalf("different-cap admit: %v", err) + } + if _, err := a.Admit(ctx, groupRequest("org-c", "m", scope)); err != nil { + t.Fatalf("different-cap admit 2: %v", err) + } + if store.calls != 2 { + t.Fatalf("store calls = %d, want 2 (a changed cap re-checks once)", store.calls) + } +} + +// TestGroupSpendCacheExpiry: past the TTL the verdict re-reads Postgres, so a +// group that hit its cap is not pinned to a stale verdict forever, and a cap +// RAISE takes effect within one TTL. +func TestGroupSpendCacheExpiry(t *testing.T) { + a, store := spendTestAdmitter(t) + now := time.Now() + a.spendCache.now = func() time.Time { return now } + scope := groupScope(testGID, nil, "100") + ctx := context.Background() + + if _, err := a.Admit(ctx, groupRequest("org-e", "m", scope)); err != nil { + t.Fatalf("first admit: %v", err) + } + if _, err := a.Admit(ctx, groupRequest("org-e", "m", scope)); err != nil { + t.Fatalf("cached admit: %v", err) + } + if store.calls != 1 { + t.Fatalf("store calls = %d, want 1 inside the TTL", store.calls) + } + + // Past the TTL: a fresh read. + now = now.Add(groupSpendCacheTTL + time.Second) + if _, err := a.Admit(ctx, groupRequest("org-e", "m", scope)); err != nil { + t.Fatalf("post-TTL admit: %v", err) + } + if store.calls != 2 { + t.Fatalf("store calls = %d, want 2 after the TTL expires", store.calls) + } +} + +// TestGroupSpendStoreErrorFailsOpen: a Postgres/store error must never deny — +// the check fails open with a loud log (the 2026-09-24 posture: quotas are +// permissible; they never take inference down). +func TestGroupSpendStoreErrorFailsOpen(t *testing.T) { + a, store := spendTestAdmitter(t) + store.err = errors.New("postgres unreachable") + + lease, err := a.Admit(context.Background(), groupRequest("org-f", "m", + groupScope(testGID, nil, "100"))) + if err != nil { + t.Fatalf("store-error admit: %v, want fail-open admission", err) + } + _ = lease.Complete(context.Background(), 0) + + // A second request hits the failing store again: errors are NOT cached — + // caching a fail-open verdict would silently extend the bypass past + // recovery, and caching nothing keeps the check self-healing. + lease2, err := a.Admit(context.Background(), groupRequest("org-f", "m", + groupScope(testGID, nil, "100"))) + if err != nil { + t.Fatalf("second store-error admit: %v, want fail-open admission", err) + } + _ = lease2.Complete(context.Background(), 0) + if store.calls != 2 { + t.Fatalf("store calls = %d, want 2 (errors never cache)", store.calls) + } +} + +// TestGroupSpendNoStoreFailsOpen: an admitter without a spend store (a +// serving-only install) admits requests carrying a spend cap — the bypass is +// the loud startup log's contract, not a per-request error. +func TestGroupSpendNoStoreFailsOpen(t *testing.T) { + a, _ := testAdmitter(t, config.AdmissionSettings{}) + lease, err := a.Admit(context.Background(), groupRequest("org-n", "m", + groupScope(testGID, nil, "100"))) + if err != nil { + t.Fatalf("no-store admit: %v, want admission (spend check bypassed)", err) + } + _ = lease.Complete(context.Background(), 0) +} + +// TestGroupScopeMissingGroupIDFailsClosed: a scope carrying limits without a +// group id is a broken envelope, and broken trusted identity fails closed — +// distinct from the store-outage bypass. +func TestGroupScopeMissingGroupIDFailsClosed(t *testing.T) { + a, _ := spendTestAdmitter(t) + cap := int64(10) + _, err := a.Admit(context.Background(), groupRequest("org-i", "m", + GroupScope{Limits: RateLimits{Requests: &cap}})) + if !errors.Is(err, ErrInvalidIdentity) { + t.Fatalf("anonymous group scope: %v, want ErrInvalidIdentity (fail closed)", err) + } + + // A scope with nothing to enforce and no id is inert, not an error. + lease, err := a.Admit(context.Background(), groupRequest("org-i", "m", GroupScope{})) + if err != nil { + t.Fatalf("inert anonymous scope: %v, want admission", err) + } + _ = lease.Complete(context.Background(), 0) +} + +// TestGroupScopesSettleThroughCompleteUsage: the group contract scope rides +// the lease like any other scope — CompleteUsage settles it through the same +// finishScript path, charging the window the engine-reported tokens (not the +// reservation). Generated cap 50: settle 30, then 20 — both admit; a third +// reservation that would push the window past 50 is rejected. +func TestGroupScopesSettleThroughCompleteUsage(t *testing.T) { + a, _ := spendTestAdmitter(t) + cap := int64(50) + r := groupRequest("org-l", "m", GroupScope{GroupID: testGID, + Limits: RateLimits{GeneratedTokens: &cap}}) + + lease, err := a.Admit(context.Background(), r) + if err != nil { + t.Fatalf("first admit: %v", err) + } + if err := lease.Complete(context.Background(), 30); err != nil { + t.Fatalf("first Complete: %v", err) + } + + lease2, err := a.Admit(context.Background(), r) + if err != nil { + t.Fatalf("second admit (window at 30/50): %v, want admitted", err) + } + if err := lease2.Complete(context.Background(), 20); err != nil { + t.Fatalf("second Complete: %v", err) + } + + _, err = a.Admit(context.Background(), r) + var rejected *Rejected + if !errors.As(err, &rejected) || !rejected.Contractual { + t.Fatalf("third admit (window at 50/50): %v, want a contractual rejection", err) + } +} diff --git a/internal/admission/spend.go b/internal/admission/spend.go new file mode 100644 index 0000000..2d1c031 --- /dev/null +++ b/internal/admission/spend.go @@ -0,0 +1,183 @@ +package admission + +import ( + "context" + "database/sql" + "fmt" + "sync" + "time" + + // pgx stdlib driver registers itself as "pgx" with database/sql — the same + // driver/DSN convention as the drainer and the rater. The spend check reads + // Postgres (group_usage, written by the rater), NOT the Valkey admission + // store, so the Valkey-outage ruling's store is not in this path. + _ "github.com/jackc/pgx/v5/stdlib" + + "github.com/saturncloud/phoebe/internal/logging" +) + +// groupSpendCacheTTL is how long a group's monthly-spend verdict is reused +// before re-reading Postgres. Ruling-scale: ~60 seconds per group (the spend +// sum moves hourly, when the rater writes a window, so a tighter TTL would +// only add load). A constant, not a config knob: the check is deliberately a +// coarse guardrail, and an operator-tunable TTL buys nothing the hourly +// rating cadence doesn't already bound. +const groupSpendCacheTTL = time.Minute + +// GroupSpendStore is the admission package's seam onto the group spend the +// monthly cap compares against. It is an interface so admission can be tested +// against a fake and the SQL tested in isolation via sqlmock, mirroring the +// drainer/rater store pattern. +type GroupSpendStore interface { + // GroupSpendExhausted reports whether the group's spend in the CURRENT + // CALENDAR MONTH has reached the cap: SUM(cost) over group_usage for + // window_start >= date_trunc('month', now()), compared in Postgres + // (NUMERIC, never a Go number). cap is a plain decimal NUMERIC(20,9) + // string straight from the trusted envelope. The verdict — not the spend + // sum — crosses this seam, so no money value becomes a Go number. + GroupSpendExhausted(ctx context.Context, groupID, cap string) (bool, error) +} + +// PostgresSpendStore answers the monthly spend check against group_usage +// (migration 0008), in phoebe's own Postgres — the same DSN convention +// (DATABASE_URL) as the drainer and the rater. +type PostgresSpendStore struct { + db *sql.DB +} + +// OpenPostgresSpendStore opens the spend-check pool. Failure modes are the +// caller's call: cmd/interceptor logs loudly and runs WITHOUT spend caps +// (group rate limits still enforce from the envelope) rather than refusing to +// serve — the 2026-09-24 posture: quotas are permissible, they never take +// inference down. +func OpenPostgresSpendStore(ctx context.Context, databaseURL string) (*PostgresSpendStore, error) { + if databaseURL == "" { + return nil, fmt.Errorf("admission: spend check: DATABASE_URL is empty (group_usage lives in phoebe's Postgres; the spend cap cannot be checked without it)") + } + db, err := sql.Open("pgx", databaseURL) + if err != nil { + return nil, fmt.Errorf("admission: spend check: open postgres: %w", err) + } + // The hot path is served from the in-process verdict cache: the pool sees + // at most one query per (group, cap) per TTL, so a handful of connections + // is generous. + db.SetMaxOpenConns(4) + db.SetMaxIdleConns(2) + db.SetConnMaxLifetime(30 * time.Minute) + s := &PostgresSpendStore{db: db} + if err := s.db.PingContext(ctx); err != nil { + _ = db.Close() + return nil, fmt.Errorf("admission: spend check: postgres ping: %w", err) + } + return s, nil +} + +// NewPostgresSpendStore wraps an existing *sql.DB (tests / caller-owned pools). +func NewPostgresSpendStore(db *sql.DB) *PostgresSpendStore { return &PostgresSpendStore{db: db} } + +func (s *PostgresSpendStore) Close() error { return s.db.Close() } + +// groupSpendExhaustedSQL compares the group's month-to-date attribution spend +// against the cap in ONE round-trip, in Postgres, so money never becomes a Go +// number. An empty group_usage month (no rating yet, or a group with no +// attributed usage) sums to NULL → COALESCE to 0 → the cap is unreached until +// rated spend says otherwise. +const groupSpendExhaustedSQL = ` +SELECT COALESCE(SUM(cost), 0) >= $2::numeric +FROM group_usage +WHERE group_id = $1 + AND window_start >= date_trunc('month', now())` + +func (s *PostgresSpendStore) GroupSpendExhausted(ctx context.Context, groupID, cap string) (bool, error) { + var exhausted bool + if err := s.db.QueryRowContext(ctx, groupSpendExhaustedSQL, groupID, cap).Scan(&exhausted); err != nil { + return false, fmt.Errorf("admission: spend check: group %s: %w", groupID, err) + } + return exhausted, nil +} + +// spendVerdictCache is the per-(group, cap) verdict cache: one entry per +// distinct cap a group has been checked against, refreshed after +// groupSpendCacheTTL. The TTL is measured from the verdict's capture, so a +// hot group re-reads Postgres at most once per TTL; a quiet group's entry is +// simply stale on its next request. Unbounded growth is not a hazard: the +// key set is bounded by the distinct (group, cap) pairs actually stamped on +// arriving envelopes (at most 16 groups per request). +type spendVerdictCache struct { + mu sync.Mutex + entries map[string]spendVerdictEntry + // now exists so a test can simulate the passage of time; nil selects + // time.Now (production). + now func() time.Time +} + +type spendVerdictEntry struct { + exhausted bool + at time.Time +} + +func newSpendVerdictCache() *spendVerdictCache { + return &spendVerdictCache{entries: map[string]spendVerdictEntry{}} +} + +// lookup returns the cached verdict for (groupID, cap) and whether it is +// fresh enough to trust. +func (c *spendVerdictCache) lookup(groupID, cap string) (bool, bool) { + c.mu.Lock() + defer c.mu.Unlock() + e, ok := c.entries[groupID+"\x00"+cap] + if !ok { + return false, false + } + now := time.Now + if c.now != nil { + now = c.now + } + if now().Sub(e.at) >= groupSpendCacheTTL { + return false, false + } + return e.exhausted, true +} + +// store records a fresh verdict for (groupID, cap). +func (c *spendVerdictCache) store(groupID, cap string, exhausted bool) { + c.mu.Lock() + defer c.mu.Unlock() + now := time.Now + if c.now != nil { + now = c.now + } + c.entries[groupID+"\x00"+cap] = spendVerdictEntry{exhausted: exhausted, at: now()} +} + +// spendFailureLog throttles the fail-open spend-check ERROR lines: the first +// bypass of an incident logs immediately, then at most one line per +// spendFailureLogInterval while the incident persists, carrying the count of +// suppressed occurrences. Without it a Postgres outage would emit one ERROR +// per admitted request carrying a group spend cap — flooding the log exactly +// when the check is down. +const spendFailureLogInterval = time.Minute + +type spendFailureLog struct { + mu sync.Mutex + last time.Time + suppressed int +} + +func (l *spendFailureLog) logf(log *logging.Logger, format string, args ...interface{}) { + l.mu.Lock() + defer l.mu.Unlock() + now := time.Now() + if now.Sub(l.last) < spendFailureLogInterval { + l.suppressed++ + return + } + l.last = now + msg := fmt.Sprintf(format, args...) + if l.suppressed > 0 { + log.Error.Printf("%s (+%d similar suppressed)", msg, l.suppressed) + l.suppressed = 0 + return + } + log.Error.Printf("%s", msg) +} From 48b9b62778fff8ccc9f40641af5014921aa51f44 Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 07:26:03 +0000 Subject: [PATCH 06/17] proxy: parse the group scope envelope and enforce it at admission MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit parseTrustedGroupScopes mirrors parseTrustedRateLimits: behind the admission-enabled gate, strict on the frozen grammar (v1 prefix, 32-hex gids, R4 rate sentinels, plain-decimal NUMERIC(20,9) spend caps, <=16 entries, <4 KiB) — malformed or oversize fails closed 503 with a loud log, the fail-visible default for multi-group callers. Denials map through the contractual class (429 + Retry-After); the undeclared-output reservation clamps to group generated-token limits like the org/owner ones; the parsed membership list rides the emitted event for the rater. --- internal/proxy/admission.go | 130 ++++++++- internal/proxy/group_scopes_test.go | 407 ++++++++++++++++++++++++++++ internal/proxy/proxy.go | 28 +- 3 files changed, 557 insertions(+), 8 deletions(-) create mode 100644 internal/proxy/group_scopes_test.go diff --git a/internal/proxy/admission.go b/internal/proxy/admission.go index 81dc316..196c206 100644 --- a/internal/proxy/admission.go +++ b/internal/proxy/admission.go @@ -7,7 +7,9 @@ import ( "fmt" "math" "net/http" + "regexp" "strconv" + "strings" "sync/atomic" "time" @@ -436,15 +438,17 @@ func anyScopedLimitPresent(id identity.Identity) bool { } // trustedPolicyEnvelopePresent reports whether any part of the trusted quota -// envelope — the owner-id anchor or any scoped limit header — reached this -// request. Any part of it engages admission's fail-closed checks even under -// admission.enabled=false: a missing organization identity, a malformed or -// partial envelope (a limit header without the owner-id anchor fails closed -// in parseTrustedRateLimits), and an underivable graph scope all answer 503. +// envelope — the owner-id anchor, any scoped limit header, or the group scope +// envelope — reached this request. Any part of it engages admission's +// fail-closed checks even under admission.enabled=false: a missing +// organization identity, a malformed or partial envelope (a limit header +// without the owner-id anchor fails closed in parseTrustedRateLimits; a +// malformed group scope envelope fails closed in parseTrustedGroupScopes), +// and an underivable graph scope all answer 503. // Under admission.enabled=false the admission store itself is only consulted // when a contract limit header is present (see the Admit call in proxy.go). func trustedPolicyEnvelopePresent(id identity.Identity) bool { - return id.OwnerID != "" || anyScopedLimitPresent(id) + return id.OwnerID != "" || anyScopedLimitPresent(id) || id.GroupScopes != "" } // undeclaredOutputReservation is the output-token reservation for a request @@ -467,6 +471,120 @@ func undeclaredOutputReservation(defaultOutput int64, scopes ...admission.RateLi // for pre-R8 producers (see legacyQuotaHeadersPresent). var errNoTrustedRateLimitPolicy = errors.New("incomplete trusted shared-inference rate-limit policy") +// The group scope envelope (X-Saturn-Group-Scopes) bounds, frozen with the +// grammar (see identity.HeaderGroupScopes): at most 16 entries and the whole +// header under 4 KiB. Exceeding either, like any malformed field, is a 503 at +// the proxy with a loud log — the fail-visible default for the multi-group +// case: usage counts against EVERY group the caller belongs to that carries +// limits, and the envelope must never silently drop one. +const ( + maxGroupScopeEntries = 16 + maxGroupScopeBytes = 4096 +) + +// groupIDRe pins a group id to 32 lowercase hex chars (uuid4 hex, as Atlas +// mints). The frozen envelope pins this grammar; anything else is malformed. +var groupIDRe = regexp.MustCompile(`^[0-9a-f]{32}$`) + +// spendCapRe pins a monthly spend cap to a plain fixed-point decimal in the +// NUMERIC(20,9) unit: digits, an optional fraction of at most 9 digits, no +// sign, no exponent (mirrors rating.decimalRe's no-float discipline — the cap +// is money). "" parses to no cap; "0" is the explicit zero cap. +var spendCapRe = regexp.MustCompile(`^[0-9]+(\.[0-9]{1,9})?$`) + +// parseTrustedGroupScopes parses the trusted group quota envelope +// (X-Saturn-Scopes grammar: v1;:,,,,;...). +// R4 sentinels exactly as parseTrustedRateLimits: an empty rate field is +// unlimited (nil), "0" is a zero cap, a positive value is the per-minute cap; +// an empty is no cap, "0" blocks all paid work for the group. The +// parse is STRICT — a malformed entry, an out-of-grammar gid or spend cap, a +// duplicate gid, an uncached-above-total relation, more than +// maxGroupScopeEntries entries, or a header at/over maxGroupScopeBytes is an +// error and the request fails closed (503), like every other envelope +// violation. An absent header parses to nil scopes (no group quotas on this +// request), and "v1" alone parses to zero groups. +// +// TRUST: the value is read by identity.FromRequest ONLY when the header is in +// the active trusted set, so a client can never stamp its own group quotas. +func parseTrustedGroupScopes(id identity.Identity) ([]admission.GroupScope, error) { + raw := id.GroupScopes + if raw == "" { + return nil, nil + } + if len(raw) >= maxGroupScopeBytes { + return nil, fmt.Errorf("trusted %s header exceeds the %d-byte envelope bound (%d bytes)", identity.HeaderGroupScopes, maxGroupScopeBytes, len(raw)) + } + parts := strings.Split(raw, ";") + if parts[0] != "v1" { + return nil, fmt.Errorf("trusted %s header missing the v1 prefix", identity.HeaderGroupScopes) + } + entries := parts[1:] + if len(entries) > maxGroupScopeEntries { + return nil, fmt.Errorf("trusted %s header carries %d group entries, over the %d-entry bound", identity.HeaderGroupScopes, len(entries), maxGroupScopeEntries) + } + scopes := make([]admission.GroupScope, 0, len(entries)) + seen := make(map[string]struct{}, len(entries)) + for _, entry := range entries { + fields := strings.Split(entry, ":") + if len(fields) != 2 { + return nil, fmt.Errorf("trusted %s entry %q is not :,,,,", identity.HeaderGroupScopes, entry) + } + gid := fields[0] + if !groupIDRe.MatchString(gid) { + return nil, fmt.Errorf("trusted %s entry has a malformed group id %q (want 32 lowercase hex)", identity.HeaderGroupScopes, gid) + } + if _, dup := seen[gid]; dup { + return nil, fmt.Errorf("trusted %s header lists group %s twice", identity.HeaderGroupScopes, gid) + } + seen[gid] = struct{}{} + rates := strings.Split(fields[1], ",") + if len(rates) != 5 { + return nil, fmt.Errorf("trusted %s entry for group %s has %d fields, want ,,,,", identity.HeaderGroupScopes, gid, len(rates)) + } + var scope admission.GroupScope + scope.GroupID = gid + var err error + if scope.Limits.Requests, err = parseGroupRate(gid, "requests", rates[0]); err != nil { + return nil, err + } + if scope.Limits.TotalPromptTokens, err = parseGroupRate(gid, "total prompt tokens", rates[1]); err != nil { + return nil, err + } + if scope.Limits.UncachedPromptTokens, err = parseGroupRate(gid, "uncached prompt tokens", rates[2]); err != nil { + return nil, err + } + if scope.Limits.GeneratedTokens, err = parseGroupRate(gid, "generated tokens", rates[3]); err != nil { + return nil, err + } + if scope.Limits.TotalPromptTokens != nil && scope.Limits.UncachedPromptTokens != nil && + *scope.Limits.UncachedPromptTokens > *scope.Limits.TotalPromptTokens { + return nil, fmt.Errorf("trusted %s group %s: uncached prompt limit exceeds total prompt limit", identity.HeaderGroupScopes, gid) + } + if spend := rates[4]; spend != "" { + if !spendCapRe.MatchString(spend) { + return nil, fmt.Errorf("trusted %s group %s has a malformed spend cap %q (want a plain decimal, at most 9 fraction digits)", identity.HeaderGroupScopes, gid, spend) + } + scope.SpendCap = spend + } + scopes = append(scopes, scope) + } + return scopes, nil +} + +// parseGroupRate parses one group envelope rate field: empty is unlimited +// (nil), "0" is a zero cap, a positive integer is the per-minute cap. A +// negative or non-integer value is malformed and fails closed. +func parseGroupRate(gid, dimension, value string) (*int64, error) { + if value == "" { + return nil, nil + } + limit, err := strconv.ParseInt(value, 10, 64) + if err != nil || limit < 0 { + return nil, fmt.Errorf("trusted %s group %s has a malformed %s limit %q", identity.HeaderGroupScopes, gid, dimension, value) + } + return &limit, nil +} + // legacyQuotaHeaderNames are the five single-scope quota headers that R8 // removed from the trusted envelope. Phoebe never reads them for a trust or // limit decision; they are listed here ONLY so the proxy can log that a diff --git a/internal/proxy/group_scopes_test.go b/internal/proxy/group_scopes_test.go new file mode 100644 index 0000000..b372d25 --- /dev/null +++ b/internal/proxy/group_scopes_test.go @@ -0,0 +1,407 @@ +package proxy + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "net/url" + "strconv" + "strings" + "testing" + "time" + + "github.com/redis/go-redis/v9" + + "github.com/saturncloud/phoebe/internal/admission" + "github.com/saturncloud/phoebe/internal/config" + "github.com/saturncloud/phoebe/internal/identity" + "github.com/saturncloud/phoebe/internal/logging" +) + +// These tests cover the membership-aware group quota envelope +// (X-Saturn-Group-Scopes, ruled 2026-10-07) at the proxy layer: the strict +// parse (fail-visible: malformed/oversize → 503) and the end-to-end +// enforcement of both scope classes (per-minute rates and the monthly spend +// cap, both answering 429 + Retry-After on denial). + +const ( + testGroupA = "a1b2c3d4e5f60718293a4b5c6d7e8f90" + testGroupB = "00112233445566778899aabbccddeeff" + + // groupScopesHeader is a fully-unlimited two-group envelope: every rate + // empty (unlimited), no spend cap. Tests override individual fields. + groupScopesHeader = "v1;" + testGroupA + ":,,,," + ";" + testGroupB + ":,,,," +) + +// groupScopesValue builds an envelope for one group with the given rate field +// string (verbatim between the colons) and spend field. +func groupScopesValue(gid, rates, spend string) string { + return "v1;" + gid + ":" + rates + "," + spend +} + +// TestParseTrustedGroupScopes exercises the frozen envelope grammar: +// v1;:,,,,;:... +func TestParseTrustedGroupScopes(t *testing.T) { + valid := []struct { + name string + in string + want []admission.GroupScope + }{ + { + "absent header", + "", + nil, + }, + { + "v1 alone is zero groups", + "v1", + nil, + }, + { + "single group all unlimited no spend cap", + groupScopesValue(testGroupA, ",,,", ""), + []admission.GroupScope{{GroupID: testGroupA}}, + }, + { + "all four rates set", + groupScopesValue(testGroupA, "30,1000,500,2000", ""), + []admission.GroupScope{{ + GroupID: testGroupA, + Limits: admission.RateLimits{ + Requests: ptr64(30), TotalPromptTokens: ptr64(1000), + UncachedPromptTokens: ptr64(500), GeneratedTokens: ptr64(2000), + }, + }}, + }, + { + "spend cap decimal", + groupScopesValue(testGroupA, ",,,", "100.000000001"), + []admission.GroupScope{{GroupID: testGroupA, SpendCap: "100.000000001"}}, + }, + { + "spend cap zero is a zero cap not absent", + groupScopesValue(testGroupA, ",,,", "0"), + []admission.GroupScope{{GroupID: testGroupA, SpendCap: "0"}}, + }, + { + "rate zero is a zero cap", + groupScopesValue(testGroupA, "0,,,", ""), + []admission.GroupScope{{ + GroupID: testGroupA, + Limits: admission.RateLimits{Requests: ptr64(0)}, + }}, + }, + { + "multi group in envelope order", + "v1;" + testGroupA + ":1,,,," + ";" + testGroupB + ":,2,,,100", + []admission.GroupScope{ + {GroupID: testGroupA, Limits: admission.RateLimits{Requests: ptr64(1)}}, + {GroupID: testGroupB, Limits: admission.RateLimits{TotalPromptTokens: ptr64(2)}, SpendCap: "100"}, + }, + }, + } + for _, tc := range valid { + t.Run(tc.name, func(t *testing.T) { + scopes, err := parseTrustedGroupScopes(identity.Identity{GroupScopes: tc.in}) + if err != nil { + t.Fatalf("parse(%q): %v, want success", tc.in, err) + } + if len(scopes) != len(tc.want) { + t.Fatalf("parse(%q) = %d scopes, want %d (%+v)", tc.in, len(scopes), len(tc.want), scopes) + } + for i, want := range tc.want { + got := scopes[i] + if got.GroupID != want.GroupID || got.SpendCap != want.SpendCap { + t.Fatalf("scope %d = %+v, want %+v", i, got, want) + } + if (got.Limits.Requests == nil) != (want.Limits.Requests == nil) || + (got.Limits.Requests != nil && *got.Limits.Requests != *want.Limits.Requests) { + t.Fatalf("scope %d requests = %v, want %v", i, got.Limits.Requests, want.Limits.Requests) + } + if (got.Limits.TotalPromptTokens == nil) != (want.Limits.TotalPromptTokens == nil) || + (got.Limits.TotalPromptTokens != nil && *got.Limits.TotalPromptTokens != *want.Limits.TotalPromptTokens) { + t.Fatalf("scope %d total = %v, want %v", i, got.Limits.TotalPromptTokens, want.Limits.TotalPromptTokens) + } + if (got.Limits.UncachedPromptTokens == nil) != (want.Limits.UncachedPromptTokens == nil) || + (got.Limits.UncachedPromptTokens != nil && *got.Limits.UncachedPromptTokens != *want.Limits.UncachedPromptTokens) { + t.Fatalf("scope %d uncached = %v, want %v", i, got.Limits.UncachedPromptTokens, want.Limits.UncachedPromptTokens) + } + if (got.Limits.GeneratedTokens == nil) != (want.Limits.GeneratedTokens == nil) || + (got.Limits.GeneratedTokens != nil && *got.Limits.GeneratedTokens != *want.Limits.GeneratedTokens) { + t.Fatalf("scope %d generated = %v, want %v", i, got.Limits.GeneratedTokens, want.Limits.GeneratedTokens) + } + } + }) + } + + malformed := []struct { + name string + in string + }{ + {"missing v1 prefix", testGroupA + ":,,,,"}, + {"wrong prefix version", "v2;" + testGroupA + ":,,,,"}, + {"empty entry from trailing semicolon", "v1;" + testGroupA + ":,,,,;"}, + {"empty entry in the middle", "v1;;" + testGroupA + ":,,,,"}, + {"gid uppercase hex", "v1;A1B2C3D4E5F60718293A4B5C6D7E8F90:,,,,"}, + {"gid too short", "v1;a1b2:,,,,"}, + {"gid not hex", "v1;g1b2c3d4e5f60718293a4b5c6d7e8f90:,,,,"}, + {"missing colon", "v1;" + testGroupA + ",,,,"}, + {"too few rate fields", groupScopesValue(testGroupA, ",,", "")}, + {"too many rate fields", "v1;" + testGroupA + ":1,2,3,4,5,6"}, + {"negative rate", groupScopesValue(testGroupA, "-1,,,", "")}, + {"non-numeric rate", groupScopesValue(testGroupA, "x,,,", "")}, + {"uncached above total", groupScopesValue(testGroupA, ",100,101,", "")}, + {"spend cap exponent", groupScopesValue(testGroupA, ",,,", "1e3")}, + {"spend cap sign", groupScopesValue(testGroupA, ",,,", "+100")}, + {"spend cap bare dot", groupScopesValue(testGroupA, ",,,", ".5")}, + {"spend cap trailing dot", groupScopesValue(testGroupA, ",,,", "5.")}, + {"spend cap too many fraction digits", groupScopesValue(testGroupA, ",,,", "0.0000000001")}, + {"duplicate gid", "v1;" + testGroupA + ":,,,," + ";" + testGroupA + ":1,,,"}, + } + for _, tc := range malformed { + t.Run(tc.name, func(t *testing.T) { + if _, err := parseTrustedGroupScopes(identity.Identity{GroupScopes: tc.in}); err == nil { + t.Fatalf("parse(%q) succeeded, want a fail-closed error", tc.in) + } + }) + } + + t.Run("over the entry bound", func(t *testing.T) { + entries := make([]string, 0, maxGroupScopeEntries+1) + for i := 0; i <= maxGroupScopeEntries; i++ { + entries = append(entries, groupIDFor(i)+":,,,,") + } + if _, err := parseTrustedGroupScopes(identity.Identity{GroupScopes: "v1;" + strings.Join(entries, ";")}); err == nil { + t.Fatalf("parse of %d entries succeeded, want the %d-entry bound enforced", len(entries), maxGroupScopeEntries) + } + }) + + t.Run("over the size bound", func(t *testing.T) { + // 16 maximal entries (32-hex gid + ":,,,,") plus the v1; prefix stays + // under 4 KiB, so pad one entry's (unused) spend field to cross it. + entries := make([]string, 0, maxGroupScopeEntries) + for i := 0; i < maxGroupScopeEntries; i++ { + entries = append(entries, groupIDFor(i)+":,,,,") + } + in := "v1;" + strings.Join(entries, ";") + in += strings.Repeat("0", maxGroupScopeBytes-len(in)) + if len(in) != maxGroupScopeBytes { + t.Fatalf("fixture size = %d, want exactly the %d-byte bound", len(in), maxGroupScopeBytes) + } + if _, err := parseTrustedGroupScopes(identity.Identity{GroupScopes: in}); err == nil { + t.Fatalf("parse of a %d-byte header succeeded, want the size bound enforced", len(in)) + } + }) +} + +// groupIDFor returns a distinct valid 32-hex group id for i. +func groupIDFor(i int) string { + s := strconv.FormatInt(int64(i), 16) + return strings.Repeat("0", 32-len(s)) + s +} + +// groupScopesRequest is a shared request carrying the group quota envelope. +func groupScopesRequest(t *testing.T, upstream *url.URL, envelope string) *http.Request { + t.Helper() + r := sharedRequest(upstream) + if envelope != "" { + r.Header.Set(identity.HeaderGroupScopes, envelope) + } + return r +} + +// proxyWithGroupScopes builds the admission-enabled proxy over a real +// (frozen-clock) miniredis, with an optional spend store. +func proxyWithGroupScopes(t *testing.T, spend admission.GroupSpendStore) (*Server, *recordingEmitter) { + t.Helper() + mr := frozenMiniredis(t) + cfg := proxyAdmissionConfig(10) + client := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = client.Close() }) + admitter := admission.New(client, cfg) + if spend != nil { + admitter.WithGroupSpend(spend, logging.New(logging.ERROR)) + } + em := &recordingEmitter{} + s := New(&config.Settings{Admission: cfg}, logging.New(logging.ERROR), em).WithAdmitter(admitter) + return s, em +} + +// TestSharedGroupRateLimitAnswers429: a group requests-per-minute cap on the +// envelope is enforced exactly like the contract scopes — the burst admits, +// the overrun answers 429 + Retry-After (the contractual mapping). +func TestSharedGroupRateLimitAnswers429(t *testing.T) { + up, hits := sharedGroupBackend(t) + s, _ := proxyWithGroupScopes(t, nil) + + r1 := groupScopesRequest(t, up, groupScopesValue(testGroupA, "1,,,", "")) + rr1 := httptest.NewRecorder() + s.Handler().ServeHTTP(rr1, r1) + if rr1.Code != http.StatusOK { + t.Fatalf("first request status=%d, want 200", rr1.Code) + } + + r2 := groupScopesRequest(t, up, groupScopesValue(testGroupA, "1,,,", "")) + rr2 := httptest.NewRecorder() + s.Handler().ServeHTTP(rr2, r2) + if rr2.Code != http.StatusTooManyRequests { + t.Fatalf("overrun status=%d, want 429", rr2.Code) + } + if rr2.Header().Get("Retry-After") == "" { + t.Fatal("overrun missing Retry-After") + } + if *hits != 1 { + t.Fatalf("upstream hits=%d, want 1 (the overrun never forwards)", *hits) + } +} + +// TestSharedGroupRateZeroCapAnswers429: an explicit 0 blocks every request +// for the group (R4 zero cap), whatever the other scopes say. +func TestSharedGroupRateZeroCapAnswers429(t *testing.T) { + up, _ := sharedGroupBackend(t) + s, _ := proxyWithGroupScopes(t, nil) + + rr := httptest.NewRecorder() + s.Handler().ServeHTTP(rr, groupScopesRequest(t, up, groupScopesValue(testGroupA, "0,,,", ""))) + if rr.Code != http.StatusTooManyRequests { + t.Fatalf("zero-cap status=%d, want 429", rr.Code) + } +} + +// staticSpendStore is a GroupSpendStore fake with a fixed verdict. +type staticSpendStore struct{ exhausted bool } + +func (s staticSpendStore) GroupSpendExhausted(context.Context, string, string) (bool, error) { + return s.exhausted, nil +} + +// TestSharedGroupSpendCapAnswers429: a group whose month-to-date spend +// reached its cap is refused with 429 + Retry-After — the contractual class, +// never 503 — even though its per-minute rates are unlimited. +func TestSharedGroupSpendCapAnswers429(t *testing.T) { + up, hits := sharedGroupBackend(t) + s, _ := proxyWithGroupScopes(t, staticSpendStore{exhausted: true}) + + r := groupScopesRequest(t, up, groupScopesValue(testGroupA, ",,,", "100")) + rr := httptest.NewRecorder() + s.Handler().ServeHTTP(rr, r) + if rr.Code != http.StatusTooManyRequests { + t.Fatalf("capped status=%d, want 429", rr.Code) + } + if rr.Header().Get("Retry-After") == "" { + t.Fatal("capped response missing Retry-After") + } + if *hits != 0 { + t.Fatalf("upstream hits=%d, want 0 (the capped request never forwards)", *hits) + } +} + +// TestSharedGroupSpendStoreErrorFailsOpen: a broken spend store must not deny +// traffic — the request is admitted (fail open) and independently metered. +func TestSharedGroupSpendStoreErrorFailsOpen(t *testing.T) { + up, hits := sharedGroupBackend(t) + s, em := proxyWithGroupScopes(t, errSpendStore{}) + + r := groupScopesRequest(t, up, groupScopesValue(testGroupA, ",,,", "100")) + rr := httptest.NewRecorder() + s.Handler().ServeHTTP(rr, r) + if rr.Code != http.StatusOK { + t.Fatalf("store-error status=%d, want 200 (fail open)", rr.Code) + } + if *hits != 1 { + t.Fatalf("upstream hits=%d, want 1", *hits) + } + if evs := em.waitForEvents(1, time.Second); len(evs) != 1 { + t.Fatalf("emitted events=%d, want 1 (metering is independent of admission)", len(evs)) + } +} + +type errSpendStore struct{} + +func (errSpendStore) GroupSpendExhausted(context.Context, string, string) (bool, error) { + return false, errors.New("postgres unreachable") +} + +// TestSharedGroupScopesMalformedHeaderFailsClosed: a malformed group envelope +// is a broken trusted policy, not a quota — 503 at the proxy, before +// admission, and nothing reaches the upstream or the emitter. +func TestSharedGroupScopesMalformedHeaderFailsClosed(t *testing.T) { + up, hits := sharedGroupBackend(t) + s, em := proxyWithGroupScopes(t, nil) + + r := groupScopesRequest(t, up, "v1;NOT-A-GROUP:,,,,") + rr := httptest.NewRecorder() + s.Handler().ServeHTTP(rr, r) + if rr.Code != http.StatusServiceUnavailable { + t.Fatalf("malformed envelope status=%d, want 503", rr.Code) + } + if *hits != 0 { + t.Fatalf("upstream hits=%d, want 0 (fail closed before forwarding)", *hits) + } + if evs := em.waitForEvents(1, 200*time.Millisecond); len(evs) != 0 { + t.Fatalf("emitted events=%d, want 0 (a refused request meters nothing)", len(evs)) + } +} + +// TestSharedGroupScopesUnlimitedHeaderAdmits: the fully-unlimited envelope +// (every field empty) parses to scopes with nothing to enforce — requests +// flow, and the membership list rides the emitted event for the rater. +func TestSharedGroupScopesUnlimitedHeaderAdmits(t *testing.T) { + up, hits := sharedGroupBackend(t) + s, em := proxyWithGroupScopes(t, nil) + + r := groupScopesRequest(t, up, groupScopesHeader) + rr := httptest.NewRecorder() + s.Handler().ServeHTTP(rr, r) + if rr.Code != http.StatusOK { + t.Fatalf("status=%d, want 200", rr.Code) + } + if *hits != 1 { + t.Fatalf("upstream hits=%d, want 1", *hits) + } + evs := em.waitForEvents(1, time.Second) + if len(evs) != 1 { + t.Fatalf("emitted events=%d, want 1", len(evs)) + } + got := evs[0].MemberGroupIDs + if len(got) != 2 || got[0] != testGroupA || got[1] != testGroupB { + t.Fatalf("MemberGroupIDs = %v, want [%s %s] (the envelope order)", got, testGroupA, testGroupB) + } +} + +// TestSharedNoGroupHeaderNoMembership: without the envelope the emitted event +// carries no membership (the drainer stores NULL). +func TestSharedNoGroupHeaderNoMembership(t *testing.T) { + up, _ := sharedGroupBackend(t) + s, em := proxyWithGroupScopes(t, nil) + + rr := httptest.NewRecorder() + s.Handler().ServeHTTP(rr, groupScopesRequest(t, up, "")) + if rr.Code != http.StatusOK { + t.Fatalf("status=%d, want 200", rr.Code) + } + evs := em.waitForEvents(1, time.Second) + if len(evs) != 1 { + t.Fatalf("emitted events=%d, want 1", len(evs)) + } + if len(evs[0].MemberGroupIDs) != 0 { + t.Fatalf("MemberGroupIDs = %v, want none", evs[0].MemberGroupIDs) + } +} + +// sharedGroupBackend is a shared-path upstream counting its hits. +func sharedGroupBackend(t *testing.T) (*url.URL, *int) { + t.Helper() + hits := 0 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits++ + _, _ = w.Write([]byte(`{"model":"model-a","usage":{"prompt_tokens":2,"completion_tokens":3}}`)) + })) + t.Cleanup(backend.Close) + up, err := url.Parse(backend.URL) + if err != nil { + t.Fatal(err) + } + return up, &hits +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 9dc1dbf..679bd52 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -603,6 +603,7 @@ func (s *Server) handleProxy(w http.ResponseWriter, r *http.Request) { tenantIdentity = "resource:" + id.ResourceID } var organizationLimits, ownerLimits admission.RateLimits + var groupScopes []admission.GroupScope var graph string if enforceAdmission { // A structurally broken policy — limit headers without their @@ -622,6 +623,17 @@ func (s *Server) handleProxy(w http.ResponseWriter, r *http.Request) { http.Error(w, "shared inference policy unavailable", http.StatusServiceUnavailable) return } + // The group quota envelope is a separate fail-closed parse: + // malformed, oversized, or over-count is a 503 with a loud log — + // the fail-visible default (the system never silently drops one + // group of a multi-group caller). Absent header → nil scopes. + var groupErr error + groupScopes, groupErr = parseTrustedGroupScopes(id) + if groupErr != nil { + s.log.Error.Printf("admission: invalid trusted group scope policy: %v request_id=%s", groupErr, requestID) + http.Error(w, "shared inference policy unavailable", http.StatusServiceUnavailable) + return + } graph = id.GraphK8sName if graph == "" { graph = graphFromUpstreamHost(upstream.Host) @@ -661,7 +673,11 @@ func (s *Server) handleProxy(w http.ResponseWriter, r *http.Request) { // default the client did not choose. The reservation is clamped // before prepareSharedDynamoRequest so the Dynamo output hint // still equals the reservation. - defaultOutput = undeclaredOutputReservation(defaultOutput, organizationLimits, ownerLimits) + groupLimits := make([]admission.RateLimits, len(groupScopes)) + for i, gs := range groupScopes { + groupLimits[i] = gs.Limits + } + defaultOutput = undeclaredOutputReservation(defaultOutput, append([]admission.RateLimits{organizationLimits, ownerLimits}, groupLimits...)...) } estimate, ok := admissionWork(body, defaultOutput) if !ok { @@ -692,12 +708,13 @@ func (s *Server) handleProxy(w http.ResponseWriter, r *http.Request) { // limit is unlimited in every scope, so it does not touch the store: no // Admit, no lease renewal, no completion. Every fail-closed check above // has already run. - if enforceAdmission && (s.settings.Admission.Enabled || organizationLimits.Any() || ownerLimits.Any()) { + if enforceAdmission && (s.settings.Admission.Enabled || organizationLimits.Any() || ownerLimits.Any() || len(groupScopes) > 0) { admitted, err = s.admitter.Admit(r.Context(), admission.Request{ Graph: graph, Organization: id.OrgID, Owner: id.OwnerID, Model: estimate.Model, PromptBytes: originalPromptBytes, EstimatedInputTokens: estimate.InputTokens, ReservedOutputTokens: estimate.OutputTokens, Adapter: id.Adapter != "", OrganizationLimits: organizationLimits, OwnerLimits: ownerLimits, + GroupScopes: groupScopes, }) if err != nil { if errors.Is(err, admission.ErrInvalidIdentity) { @@ -1088,6 +1105,13 @@ func (s *Server) emit(ctx context.Context, id identity.Identity, requestID, clie GroupID: id.GroupID, ResourceID: id.ResourceID, ResourceType: id.ResourceType, + // MemberGroupIDs is the caller's group list from the trusted group + // quota envelope (membership-aware group quotas, ruled 2026-10-07) — + // evidence the rater attributes against each group (plus GroupID, + // the group-token case) in the group_usage rollup. Empty when the + // caller belongs to no group; the drainer stores NULL then, like + // every other absent identity field. + MemberGroupIDs: id.MemberGroupIDs, // OrgID is the deployment-owning org (E2), stamped from the trusted // X-Saturn-Org-Id header. Captured here so the org rides the event to push; // empty is tolerated (a missing org is held + screamed at push, never gated From 5452f4c31f40c20789ba143cb71ec2d819d59edd Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 07:26:11 +0000 Subject: [PATCH 07/17] interceptor: wire the group spend store from DATABASE_URL The monthly spend cap reads group_usage in Postgres (the rater's attribution rollup), not the admission Valkey. With a DSN the admitter gets the spend store; without one (serving-only spokes) group rate limits still enforce and the missing spend check is a loud startup log, never a silent assumption. --- cmd/interceptor/main.go | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/cmd/interceptor/main.go b/cmd/interceptor/main.go index abcb816..5c07533 100644 --- a/cmd/interceptor/main.go +++ b/cmd/interceptor/main.go @@ -40,7 +40,7 @@ func main() { // R3 trusted-header registry: parse PHOEBE_TRUSTED_HEADERS (rendered // from the phoebe chart's ConfigMap) into the active set the parser's - // envelope reads resolve through. The pinned 13 already govern from + // envelope reads resolve through. The pinned 14 already govern from // package init; this engages the runtime config once, at startup, and // warns loudly if the chart render was empty/malformed. identity.LoadTrustedHeaders(log) @@ -100,6 +100,26 @@ func buildAdmission(s *config.Settings, log *logging.Logger) (admission.Admitter // docs/shared-tier-admission.md). log.Info.Printf("admission: admission.valkeyAddr is empty; using the metering Valkey from emit.valkeyAddr (%s) as the admission store", cfg.ValkeyAddr) } + // The monthly group spend cap reads group_usage in Postgres (the rater's + // attribution rollup), NOT the admission Valkey. Without a DATABASE_URL + // (a serving-only spoke install runs no Postgres at all) the cap cannot + // be checked: group rate limits still enforce, the spend cap is bypassed + // fail-open — said loudly here, never assumed silently. + closeSpend := func() {} + if dsn := os.Getenv("DATABASE_URL"); dsn != "" { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + store, err := admission.OpenPostgresSpendStore(ctx, dsn) + cancel() + if err != nil { + log.Error.Printf("admission: group spend caps DISABLED (spend-check store could not be built: %v); group rate limits still enforce", err) + } else { + admitter.WithGroupSpend(store, log) + closeSpend = func() { _ = store.Close() } + log.Info.Printf("admission: group spend caps enabled (spend check reads group_usage in Postgres)") + } + } else { + log.Error.Printf("admission: DATABASE_URL is unset; group spend caps are NOT enforced (group rate limits still enforce)") + } mode := "contract limits only (admission.enabled=false: operator capacity tiers off, envelope-less shared routes allowed)" if cfg.Enabled { mode = "contract limits + operator capacity tiers (admission.enabled=true: every shared request must carry the trusted envelope)" @@ -114,7 +134,10 @@ func buildAdmission(s *config.Settings, log *logging.Logger) (admission.Admitter } else { log.Info.Printf("admission: enforcing %s (valkey %s, lease ttl %s)", mode, cfg.ValkeyAddr, cfg.LeaseTTL) } - return admitter, func() { _ = client.Close() } + return admitter, func() { + closeSpend() + _ = client.Close() + } } // buildGateway constructs the TF gateway (org, model) resolver. DEFAULT: the From 6889d1c042414b27734f13d9290e421edca0d1d4 Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 08:12:05 +0000 Subject: [PATCH 08/17] admission: bound the group spend check; proxy: enforce the frozen envelope grammar Spend check: bound each Postgres read on the request path (250ms) and pool dials (2s) so a blackholed store fails open instead of hanging requests; compute the month boundary in UTC to match the rater's UTC hour buckets; expire cached verdicts at the UTC month boundary. Envelope parser: rate fields must match the frozen grammar (reject signed and zero-padded values) and spend caps are bounded to the NUMERIC(20,9) integer width. --- internal/admission/admission.go | 6 +- internal/admission/group_test.go | 102 ++++++++++++++++++++++++++++ internal/admission/spend.go | 56 +++++++++++---- internal/proxy/admission.go | 22 ++++-- internal/proxy/group_scopes_test.go | 17 +++++ 5 files changed, 185 insertions(+), 18 deletions(-) diff --git a/internal/admission/admission.go b/internal/admission/admission.go index 33d4093..607ae50 100644 --- a/internal/admission/admission.go +++ b/internal/admission/admission.go @@ -489,7 +489,11 @@ func (a *RedisAdmitter) checkGroupSpend(ctx context.Context, scopes []GroupScope } continue } - exhausted, err := a.spendStore.GroupSpendExhausted(ctx, gs.GroupID, gs.SpendCap) + // Bounded: a hung store read fails open (deadline exceeded is just + // another store error) instead of holding the request. + queryCtx, cancel := context.WithTimeout(ctx, groupSpendQueryBudget) + exhausted, err := a.spendStore.GroupSpendExhausted(queryCtx, gs.GroupID, gs.SpendCap) + cancel() if err != nil { a.spendBypass(gs.GroupID, err) continue diff --git a/internal/admission/group_test.go b/internal/admission/group_test.go index 7b897f8..1bafe14 100644 --- a/internal/admission/group_test.go +++ b/internal/admission/group_test.go @@ -1,11 +1,16 @@ package admission import ( + "bytes" "context" "errors" + "regexp" + "strings" "testing" "time" + sqlmock "github.com/DATA-DOG/go-sqlmock" + "github.com/saturncloud/phoebe/internal/config" "github.com/saturncloud/phoebe/internal/logging" ) @@ -255,6 +260,103 @@ func TestGroupSpendStoreErrorFailsOpen(t *testing.T) { } } +// blockingGroupSpendStore models a blackholed Postgres: the read never +// answers and returns only when its context is done. +type blockingGroupSpendStore struct{} + +func (blockingGroupSpendStore) GroupSpendExhausted(ctx context.Context, _, _ string) (bool, error) { + <-ctx.Done() + return false, ctx.Err() +} + +// TestGroupSpendHungStoreFailsOpenWithinBudget: a store read that never +// returns must not hold the request. The check abandons it after +// groupSpendQueryBudget, admits (fail open), and logs the bypass loudly. +func TestGroupSpendHungStoreFailsOpenWithinBudget(t *testing.T) { + a, _ := testAdmitter(t, config.AdmissionSettings{}) + log := logging.New(logging.ERROR) + var buf bytes.Buffer + log.Error.SetOutput(&buf) + a.WithGroupSpend(blockingGroupSpendStore{}, log) + + start := time.Now() + lease, err := a.Admit(context.Background(), groupRequest("org-h", "m", + groupScope(testGID, nil, "100"))) + elapsed := time.Since(start) + if err != nil { + t.Fatalf("hung-store admit: %v, want fail-open admission", err) + } + _ = lease.Complete(context.Background(), 0) + if elapsed > admitOperationBudget { + t.Fatalf("hung-store admit took %v, want it bounded by the spend query budget (%v)", elapsed, groupSpendQueryBudget) + } + out := buf.String() + if !strings.Contains(out, "monthly spend cap NOT enforced (fail open)") || !strings.Contains(out, context.DeadlineExceeded.Error()) { + t.Fatalf("bypass log = %q, want a fail-open ERROR naming the deadline", out) + } +} + +// TestGroupSpendMonthBoundaryIsUTC: the month predicate is computed in UTC, +// matching the rater's UTC hour buckets, not in the session TimeZone. +func TestGroupSpendMonthBoundaryIsUTC(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock: %v", err) + } + defer db.Close() + mock.ExpectQuery(regexp.QuoteMeta( + "window_start >= date_trunc('month', now() AT TIME ZONE 'UTC') AT TIME ZONE 'UTC'")). + WithArgs(testGID, "100"). + WillReturnRows(sqlmock.NewRows([]string{"exhausted"}).AddRow(true)) + + exhausted, err := NewPostgresSpendStore(db).GroupSpendExhausted(context.Background(), testGID, "100") + if err != nil { + t.Fatalf("GroupSpendExhausted: %v", err) + } + if !exhausted { + t.Fatalf("exhausted = false, want the row's verdict") + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("expectations: %v", err) + } +} + +// TestGroupSpendCacheExpiresAtUTCMonthBoundary: an exhausted verdict cached +// just before the UTC month turns must not deny requests in the new month — +// the rollover invalidates it even though the TTL has not elapsed. +func TestGroupSpendCacheExpiresAtUTCMonthBoundary(t *testing.T) { + a, store := spendTestAdmitter(t) + store.exhausted = true + now := time.Date(2026, time.October, 31, 23, 59, 50, 0, time.UTC) + a.spendCache.now = func() time.Time { return now } + scope := groupScope(testGID, nil, "100") + ctx := context.Background() + + var rejected *Rejected + if _, err := a.Admit(ctx, groupRequest("org-m", "m", scope)); !errors.As(err, &rejected) { + t.Fatalf("pre-boundary admit: %v, want a monthly_spend rejection", err) + } + if _, err := a.Admit(ctx, groupRequest("org-m", "m", scope)); !errors.As(err, &rejected) { + t.Fatalf("cached pre-boundary admit: %v, want a monthly_spend rejection", err) + } + if store.calls != 1 { + t.Fatalf("store calls = %d, want 1 (verdict cached inside the old month)", store.calls) + } + + // Twenty seconds later, inside the TTL but in November: the new month's + // spend is zero, so the cached verdict must be dropped and re-read. + now = now.Add(20 * time.Second) + store.exhausted = false + lease, err := a.Admit(ctx, groupRequest("org-m", "m", scope)) + if err != nil { + t.Fatalf("post-boundary admit: %v, want admission (new month)", err) + } + _ = lease.Complete(ctx, 0) + if store.calls != 2 { + t.Fatalf("store calls = %d, want 2 (month rollover re-reads the store)", store.calls) + } +} + // TestGroupSpendNoStoreFailsOpen: an admitter without a spend store (a // serving-only install) admits requests carrying a spend cap — the bypass is // the loud startup log's contract, not a per-request error. diff --git a/internal/admission/spend.go b/internal/admission/spend.go index 2d1c031..270abf3 100644 --- a/internal/admission/spend.go +++ b/internal/admission/spend.go @@ -7,11 +7,8 @@ import ( "sync" "time" - // pgx stdlib driver registers itself as "pgx" with database/sql — the same - // driver/DSN convention as the drainer and the rater. The spend check reads - // Postgres (group_usage, written by the rater), NOT the Valkey admission - // store, so the Valkey-outage ruling's store is not in this path. - _ "github.com/jackc/pgx/v5/stdlib" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/stdlib" "github.com/saturncloud/phoebe/internal/logging" ) @@ -24,6 +21,18 @@ import ( // rating cadence doesn't already bound. const groupSpendCacheTTL = time.Minute +// groupSpendQueryBudget bounds one spend-check read on the request path. The +// check runs synchronously before the Valkey reservation, so a slow or +// blackholed Postgres must not hold the request: past this budget the read is +// abandoned and the check fails open like any other store error. It sits well +// inside admitOperationBudget so the reservation keeps its own time. +const groupSpendQueryBudget = 250 * time.Millisecond + +// groupSpendConnectTimeout bounds establishing one pool connection, so a +// Postgres that drops packets cannot pin a dial (and with it one of the four +// pool slots) forever. +const groupSpendConnectTimeout = 2 * time.Second + // GroupSpendStore is the admission package's seam onto the group spend the // monthly cap compares against. It is an interface so admission can be tested // against a fake and the SQL tested in isolation via sqlmock, mirroring the @@ -31,7 +40,7 @@ const groupSpendCacheTTL = time.Minute type GroupSpendStore interface { // GroupSpendExhausted reports whether the group's spend in the CURRENT // CALENDAR MONTH has reached the cap: SUM(cost) over group_usage for - // window_start >= date_trunc('month', now()), compared in Postgres + // window_start >= the start of the current UTC month, compared in Postgres // (NUMERIC, never a Go number). cap is a plain decimal NUMERIC(20,9) // string straight from the trusted envelope. The verdict — not the spend // sum — crosses this seam, so no money value becomes a Go number. @@ -54,10 +63,21 @@ func OpenPostgresSpendStore(ctx context.Context, databaseURL string) (*PostgresS if databaseURL == "" { return nil, fmt.Errorf("admission: spend check: DATABASE_URL is empty (group_usage lives in phoebe's Postgres; the spend cap cannot be checked without it)") } - db, err := sql.Open("pgx", databaseURL) + // The pgx stdlib driver, opened from a parsed config so the connect + // timeout can be set — the same driver/DSN convention as the drainer and + // the rater. The spend check reads Postgres (group_usage, written by the + // rater), NOT the Valkey admission store, so the Valkey-outage ruling's + // store is not in this path. + cfg, err := pgx.ParseConfig(databaseURL) if err != nil { - return nil, fmt.Errorf("admission: spend check: open postgres: %w", err) + return nil, fmt.Errorf("admission: spend check: parse DATABASE_URL: %w", err) + } + // A connect_timeout given in the DSN wins; otherwise dials are bounded + // here rather than left to the kernel's TCP timeout. + if cfg.ConnectTimeout == 0 { + cfg.ConnectTimeout = groupSpendConnectTimeout } + db := stdlib.OpenDB(*cfg) // The hot path is served from the in-process verdict cache: the pool sees // at most one query per (group, cap) per TTL, so a handful of connections // is generous. @@ -79,14 +99,17 @@ func (s *PostgresSpendStore) Close() error { return s.db.Close() } // groupSpendExhaustedSQL compares the group's month-to-date attribution spend // against the cap in ONE round-trip, in Postgres, so money never becomes a Go -// number. An empty group_usage month (no rating yet, or a group with no +// number. The month boundary is computed in UTC explicitly: the rater fills +// window_start from UTC hour buckets, so a boundary taken in the session's +// TimeZone would shift the month by the session's UTC offset. An empty +// group_usage month (no rating yet, or a group with no // attributed usage) sums to NULL → COALESCE to 0 → the cap is unreached until // rated spend says otherwise. const groupSpendExhaustedSQL = ` SELECT COALESCE(SUM(cost), 0) >= $2::numeric FROM group_usage WHERE group_id = $1 - AND window_start >= date_trunc('month', now())` + AND window_start >= date_trunc('month', now() AT TIME ZONE 'UTC') AT TIME ZONE 'UTC'` func (s *PostgresSpendStore) GroupSpendExhausted(ctx context.Context, groupID, cap string) (bool, error) { var exhausted bool @@ -121,7 +144,9 @@ func newSpendVerdictCache() *spendVerdictCache { } // lookup returns the cached verdict for (groupID, cap) and whether it is -// fresh enough to trust. +// fresh enough to trust. An entry expires at the earlier of its TTL and the +// next UTC month boundary after its capture: the cap counts calendar-month +// spend, so a verdict taken in the old month says nothing about the new one. func (c *spendVerdictCache) lookup(groupID, cap string) (bool, bool) { c.mu.Lock() defer c.mu.Unlock() @@ -133,7 +158,8 @@ func (c *spendVerdictCache) lookup(groupID, cap string) (bool, bool) { if c.now != nil { now = c.now } - if now().Sub(e.at) >= groupSpendCacheTTL { + t := now() + if t.Sub(e.at) >= groupSpendCacheTTL || !t.Before(nextUTCMonth(e.at)) { return false, false } return e.exhausted, true @@ -150,6 +176,12 @@ func (c *spendVerdictCache) store(groupID, cap string, exhausted bool) { c.entries[groupID+"\x00"+cap] = spendVerdictEntry{exhausted: exhausted, at: now()} } +// nextUTCMonth returns the first instant of the UTC calendar month after t's. +func nextUTCMonth(t time.Time) time.Time { + u := t.UTC() + return time.Date(u.Year(), u.Month()+1, 1, 0, 0, 0, 0, time.UTC) +} + // spendFailureLog throttles the fail-open spend-check ERROR lines: the first // bypass of an incident logs immediately, then at most one line per // spendFailureLogInterval while the incident persists, carrying the count of diff --git a/internal/proxy/admission.go b/internal/proxy/admission.go index 196c206..1ed52bb 100644 --- a/internal/proxy/admission.go +++ b/internal/proxy/admission.go @@ -487,10 +487,18 @@ const ( var groupIDRe = regexp.MustCompile(`^[0-9a-f]{32}$`) // spendCapRe pins a monthly spend cap to a plain fixed-point decimal in the -// NUMERIC(20,9) unit: digits, an optional fraction of at most 9 digits, no -// sign, no exponent (mirrors rating.decimalRe's no-float discipline — the cap -// is money). "" parses to no cap; "0" is the explicit zero cap. -var spendCapRe = regexp.MustCompile(`^[0-9]+(\.[0-9]{1,9})?$`) +// NUMERIC(20,9) unit: at most 11 integer digits, an optional fraction of at +// most 9 digits, no sign, no exponent (mirrors rating.decimalRe's no-float +// discipline — the cap is money). "" parses to no cap; "0" is the explicit +// zero cap. The integer part is bounded because NUMERIC(20,9) holds at most +// 11 integer digits; a longer cap is out of contract and fails closed. +var spendCapRe = regexp.MustCompile(`^[0-9]{1,11}(\.[0-9]{1,9})?$`) + +// groupRateRe pins a rate field to the frozen grammar: "0" (zero cap) or a +// positive unpadded decimal. strconv.ParseInt alone also accepts "+30", +// "-0" and "0030", which the grammar forbids — an encoder bug there must +// fail closed (503), not be silently enforced. +var groupRateRe = regexp.MustCompile(`^(0|[1-9][0-9]*)$`) // parseTrustedGroupScopes parses the trusted group quota envelope // (X-Saturn-Scopes grammar: v1;:,,,,;...). @@ -573,11 +581,15 @@ func parseTrustedGroupScopes(id identity.Identity) ([]admission.GroupScope, erro // parseGroupRate parses one group envelope rate field: empty is unlimited // (nil), "0" is a zero cap, a positive integer is the per-minute cap. A -// negative or non-integer value is malformed and fails closed. +// signed, zero-padded, negative, or non-integer value is malformed and fails +// closed. func parseGroupRate(gid, dimension, value string) (*int64, error) { if value == "" { return nil, nil } + if !groupRateRe.MatchString(value) { + return nil, fmt.Errorf("trusted %s group %s has a malformed %s limit %q", identity.HeaderGroupScopes, gid, dimension, value) + } limit, err := strconv.ParseInt(value, 10, 64) if err != nil || limit < 0 { return nil, fmt.Errorf("trusted %s group %s has a malformed %s limit %q", identity.HeaderGroupScopes, gid, dimension, value) diff --git a/internal/proxy/group_scopes_test.go b/internal/proxy/group_scopes_test.go index b372d25..796238f 100644 --- a/internal/proxy/group_scopes_test.go +++ b/internal/proxy/group_scopes_test.go @@ -79,6 +79,16 @@ func TestParseTrustedGroupScopes(t *testing.T) { groupScopesValue(testGroupA, ",,,", "100.000000001"), []admission.GroupScope{{GroupID: testGroupA, SpendCap: "100.000000001"}}, }, + { + "spend cap 11-digit integer", + groupScopesValue(testGroupA, ",,,", "99999999999"), + []admission.GroupScope{{GroupID: testGroupA, SpendCap: "99999999999"}}, + }, + { + "spend cap at the numeric 20 9 bound", + groupScopesValue(testGroupA, ",,,", "99999999999.999999999"), + []admission.GroupScope{{GroupID: testGroupA, SpendCap: "99999999999.999999999"}}, + }, { "spend cap zero is a zero cap not absent", groupScopesValue(testGroupA, ",,,", "0"), @@ -151,6 +161,13 @@ func TestParseTrustedGroupScopes(t *testing.T) { {"too many rate fields", "v1;" + testGroupA + ":1,2,3,4,5,6"}, {"negative rate", groupScopesValue(testGroupA, "-1,,,", "")}, {"non-numeric rate", groupScopesValue(testGroupA, "x,,,", "")}, + {"signed requests rate", groupScopesValue(testGroupA, "+5,,,", "")}, + {"negative zero total rate", groupScopesValue(testGroupA, ",-0,,", "")}, + {"zero-padded uncached rate", groupScopesValue(testGroupA, ",,05,", "")}, + {"signed generated rate", groupScopesValue(testGroupA, ",,,+30", "")}, + {"zero-padded requests rate", groupScopesValue(testGroupA, "0030,,,", "")}, + {"spend cap 12-digit integer", groupScopesValue(testGroupA, ",,,", "100000000000")}, + {"spend cap 12-digit integer with full fraction", groupScopesValue(testGroupA, ",,,", "100000000000.000000000")}, {"uncached above total", groupScopesValue(testGroupA, ",100,101,", "")}, {"spend cap exponent", groupScopesValue(testGroupA, ",,,", "1e3")}, {"spend cap sign", groupScopesValue(testGroupA, ",,,", "+100")}, From 197668155cf6bf5de114e3d3ceaffac3f87bb70f Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 08:45:05 +0000 Subject: [PATCH 09/17] admission: collapse the spend-check stampede after a store error Cache a short-lived fail-open verdict after a spend-store error so a slow or blackholed Postgres is not re-queried once per request, collapse concurrent checks of one (group, cap) into a single in-flight read, and stop logging a client disconnect as a quota bypass, which consumed the throttled error slot a real outage needs. Add an integration-tagged test that runs the spend comparison against live Postgres loaded with the real migrations, so the zero-cap and spend-equals-cap rules cannot regress unseen. --- internal/admission/admission.go | 54 ++++-- internal/admission/group_test.go | 138 +++++++++++++++- internal/admission/spend.go | 96 ++++++++++- internal/admission/spend_integration_test.go | 163 +++++++++++++++++++ 4 files changed, 425 insertions(+), 26 deletions(-) create mode 100644 internal/admission/spend_integration_test.go diff --git a/internal/admission/admission.go b/internal/admission/admission.go index 607ae50..dbdf6f8 100644 --- a/internal/admission/admission.go +++ b/internal/admission/admission.go @@ -161,10 +161,11 @@ type RedisAdmitter struct { // (serving-only spokes): group RATE limits still enforce from the Valkey // store, spend caps are then unchecked — logged loudly at startup by // cmd/interceptor, never silently assumed. - spendStore GroupSpendStore - spendCache *spendVerdictCache - spendLog *logging.Logger - spendFails *spendFailureLog + spendStore GroupSpendStore + spendCache *spendVerdictCache + spendFlight *spendFlight + spendLog *logging.Logger + spendFails *spendFailureLog } // WithGroupSpend enables the monthly spend cap check against a GroupSpendStore @@ -176,6 +177,7 @@ func (a *RedisAdmitter) WithGroupSpend(store GroupSpendStore, log *logging.Logge a.spendStore = store a.spendLog = log a.spendCache = newSpendVerdictCache() + a.spendFlight = newSpendFlight() a.spendFails = &spendFailureLog{} return a } @@ -483,23 +485,45 @@ func (a *RedisAdmitter) checkGroupSpend(ctx context.Context, scopes []GroupScope if gs.SpendCap == "" { continue } - if exhausted, ok := a.spendCache.lookup(gs.GroupID, gs.SpendCap); ok { - if exhausted { + if verdict, ok := a.spendCache.lookup(gs.GroupID, gs.SpendCap); ok { + if verdict.failed { + // A recent store error already logged its bypass once + // (throttled): admit without re-querying and without waiting + // on the store. + continue + } + if verdict.exhausted { return &Rejected{Scope: "group:" + gs.GroupID, Dimension: "monthly_spend", RetryAfter: groupSpendCacheTTL, Contractual: true} } continue } - // Bounded: a hung store read fails open (deadline exceeded is just - // another store error) instead of holding the request. - queryCtx, cancel := context.WithTimeout(ctx, groupSpendQueryBudget) - exhausted, err := a.spendStore.GroupSpendExhausted(queryCtx, gs.GroupID, gs.SpendCap) - cancel() - if err != nil { - a.spendBypass(gs.GroupID, err) + // Collapse concurrent misses on one (group, cap) into a single store + // read: waiters take the leader's result (verdict, or fail-open on + // error) instead of each running their own query against a slow store. + result, leader := a.spendFlight.do(gs.GroupID+"\x00"+gs.SpendCap, func() (bool, error) { + // Bounded: a hung store read fails open (deadline exceeded is just + // another store error) instead of holding the request. + queryCtx, cancel := context.WithTimeout(ctx, groupSpendQueryBudget) + defer cancel() + return a.spendStore.GroupSpendExhausted(queryCtx, gs.GroupID, gs.SpendCap) + }) + if result.err != nil { + // A canceled request context (client gone) makes the check moot: + // it is not evidence the store is down and must not consume the + // throttled bypass slot that a genuine outage needs (KeepAlive + // guards the same race at renew time). The 250ms deadline + // expiring with a live client IS logged: the store is too slow. + // Only the leader records and logs — waiters share its outcome. + if leader && ctx.Err() == nil { + a.spendCache.storeFailure(gs.GroupID, gs.SpendCap) + a.spendBypass(gs.GroupID, result.err) + } continue } - a.spendCache.store(gs.GroupID, gs.SpendCap, exhausted) - if exhausted { + if leader { + a.spendCache.store(gs.GroupID, gs.SpendCap, result.exhausted) + } + if result.exhausted { return &Rejected{Scope: "group:" + gs.GroupID, Dimension: "monthly_spend", RetryAfter: groupSpendCacheTTL, Contractual: true} } } diff --git a/internal/admission/group_test.go b/internal/admission/group_test.go index 1bafe14..4c73564 100644 --- a/internal/admission/group_test.go +++ b/internal/admission/group_test.go @@ -6,6 +6,7 @@ import ( "errors" "regexp" "strings" + "sync" "testing" "time" @@ -19,6 +20,7 @@ import ( // are scripted per call so a test can observe exactly how often Postgres would // be read (the cache-behavior tests pin the once-per-TTL contract). type fakeGroupSpendStore struct { + mu sync.Mutex exhausted bool err error calls int @@ -27,6 +29,8 @@ type fakeGroupSpendStore struct { } func (f *fakeGroupSpendStore) GroupSpendExhausted(_ context.Context, groupID, cap string) (bool, error) { + f.mu.Lock() + defer f.mu.Unlock() f.calls++ f.gotGID = append(f.gotGID, groupID) f.gotCap = append(f.gotCap, cap) @@ -234,9 +238,13 @@ func TestGroupSpendCacheExpiry(t *testing.T) { // TestGroupSpendStoreErrorFailsOpen: a Postgres/store error must never deny — // the check fails open with a loud log (the 2026-09-24 posture: quotas are -// permissible; they never take inference down). +// permissible; they never take inference down). The error is negative-cached +// briefly so an outage is not re-queried per request, and the check re-probes +// once the short TTL lapses, so recovery is seen quickly. func TestGroupSpendStoreErrorFailsOpen(t *testing.T) { a, store := spendTestAdmitter(t) + now := time.Now() + a.spendCache.now = func() time.Time { return now } store.err = errors.New("postgres unreachable") lease, err := a.Admit(context.Background(), groupRequest("org-f", "m", @@ -246,17 +254,61 @@ func TestGroupSpendStoreErrorFailsOpen(t *testing.T) { } _ = lease.Complete(context.Background(), 0) - // A second request hits the failing store again: errors are NOT cached — - // caching a fail-open verdict would silently extend the bypass past - // recovery, and caching nothing keeps the check self-healing. + // Inside the negative TTL the second request admits without re-querying. lease2, err := a.Admit(context.Background(), groupRequest("org-f", "m", groupScope(testGID, nil, "100"))) if err != nil { t.Fatalf("second store-error admit: %v, want fail-open admission", err) } _ = lease2.Complete(context.Background(), 0) + if store.calls != 1 { + t.Fatalf("store calls = %d, want 1 (the error is negative-cached)", store.calls) + } + + // Past the negative TTL the check re-probes the store — and still admits + // while the error persists. + now = now.Add(groupSpendNegativeCacheTTL + time.Second) + lease3, err := a.Admit(context.Background(), groupRequest("org-f", "m", + groupScope(testGID, nil, "100"))) + if err != nil { + t.Fatalf("post-negative-TTL admit: %v, want fail-open admission", err) + } + _ = lease3.Complete(context.Background(), 0) if store.calls != 2 { - t.Fatalf("store calls = %d, want 2 (errors never cache)", store.calls) + t.Fatalf("store calls = %d, want 2 (recovery re-probes after the negative TTL)", store.calls) + } +} + +// TestGroupSpendStoreErrorNegativeCache: after one failing Admit, a second +// Admit within the negative-cache TTL makes no store call, does not wait on +// the store, and admits. +func TestGroupSpendStoreErrorNegativeCache(t *testing.T) { + a, store := spendTestAdmitter(t) + store.err = errors.New("postgres unreachable") + scope := groupScope(testGID, nil, "100") + ctx := context.Background() + + lease, err := a.Admit(ctx, groupRequest("org-nc", "m", scope)) + if err != nil { + t.Fatalf("first admit: %v", err) + } + _ = lease.Complete(ctx, 0) + if store.calls != 1 { + t.Fatalf("store calls = %d, want 1", store.calls) + } + + start := time.Now() + lease2, err := a.Admit(ctx, groupRequest("org-nc", "m", scope)) + elapsed := time.Since(start) + if err != nil { + t.Fatalf("negative-cached admit: %v, want fail-open admission", err) + } + _ = lease2.Complete(ctx, 0) + if store.calls != 1 { + t.Fatalf("store calls = %d, want 1 (no re-query inside the negative TTL)", store.calls) + } + if elapsed >= groupSpendQueryBudget { + t.Fatalf("negative-cached admit took %v, want no wait on the store (budget %v)", elapsed, groupSpendQueryBudget) } } @@ -296,6 +348,82 @@ func TestGroupSpendHungStoreFailsOpenWithinBudget(t *testing.T) { } } +// countedBlockingSpendStore models a blackholed Postgres that concurrent +// requests pile onto: the read never answers until its context is done, and +// every read is counted. +type countedBlockingSpendStore struct { + mu sync.Mutex + calls int +} + +func (s *countedBlockingSpendStore) GroupSpendExhausted(ctx context.Context, _, _ string) (bool, error) { + s.mu.Lock() + s.calls++ + s.mu.Unlock() + <-ctx.Done() + return false, ctx.Err() +} + +func (s *countedBlockingSpendStore) Calls() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.calls +} + +// TestGroupSpendConcurrentMissSingleflight: N concurrent Admits on one capped +// group with a blocking store collapse to a single store call — concurrent +// cache misses share one in-flight read — and every request still admits. +func TestGroupSpendConcurrentMissSingleflight(t *testing.T) { + a, _ := testAdmitter(t, config.AdmissionSettings{}) + store := &countedBlockingSpendStore{} + a.WithGroupSpend(store, logging.New(logging.ERROR)) + + const n = 20 + start := make(chan struct{}) + errs := make(chan error, n) + for i := 0; i < n; i++ { + go func() { + <-start + lease, err := a.Admit(context.Background(), groupRequest("org-sf", "m", + groupScope(testGID, nil, "100"))) + if err == nil { + err = lease.Complete(context.Background(), 0) + } + errs <- err + }() + } + close(start) + for i := 0; i < n; i++ { + if err := <-errs; err != nil { + t.Fatalf("concurrent admit %d: %v, want fail-open admission", i, err) + } + } + if calls := store.Calls(); calls != 1 { + t.Fatalf("store calls = %d, want 1 (concurrent misses share one in-flight read)", calls) + } +} + +// TestGroupSpendParentCancelNotLoggedAsBypass: a canceled request context +// (client gone) makes the spend check moot — it must NOT log the throttled +// 'monthly spend cap NOT enforced' ERROR, which a genuine outage needs. +func TestGroupSpendParentCancelNotLoggedAsBypass(t *testing.T) { + a, _ := testAdmitter(t, config.AdmissionSettings{}) + log := logging.New(logging.ERROR) + var buf bytes.Buffer + log.Error.SetOutput(&buf) + a.WithGroupSpend(blockingGroupSpendStore{}, log) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + // The Admit itself fails on the canceled context (the Valkey reservation + // cannot run); only the absence of the bypass line is asserted. + _, _ = a.Admit(ctx, groupRequest("org-pc", "m", + groupScope(testGID, nil, "100"))) + if out := buf.String(); strings.Contains(out, "monthly spend cap NOT enforced") { + t.Fatalf("canceled-request bypass logged: %q, want no bypass ERROR (the check is moot)", out) + } +} + // TestGroupSpendMonthBoundaryIsUTC: the month predicate is computed in UTC, // matching the rater's UTC hour buckets, not in the session TimeZone. func TestGroupSpendMonthBoundaryIsUTC(t *testing.T) { diff --git a/internal/admission/spend.go b/internal/admission/spend.go index 270abf3..bfa4057 100644 --- a/internal/admission/spend.go +++ b/internal/admission/spend.go @@ -21,6 +21,13 @@ import ( // rating cadence doesn't already bound. const groupSpendCacheTTL = time.Minute +// groupSpendNegativeCacheTTL is how long a fail-open verdict recorded after a +// store error is reused before re-reading Postgres. Short on purpose: the +// negative entry only absorbs the re-query storm while the store is down — +// past it the check re-probes, so recovery is seen within seconds rather than +// after a full verdict TTL. +const groupSpendNegativeCacheTTL = 10 * time.Second + // groupSpendQueryBudget bounds one spend-check read on the request path. The // check runs synchronously before the Valkey reservation, so a slow or // blackholed Postgres must not hold the request: past this budget the read is @@ -136,7 +143,14 @@ type spendVerdictCache struct { type spendVerdictEntry struct { exhausted bool - at time.Time + // failed marks a fail-open entry recorded after a store error: it admits + // without querying and without waiting, its bypass already logged once + // (throttled) by the request that took the error. It expires after the + // short groupSpendNegativeCacheTTL and is not subject to the UTC month + // rollover — a fail-open verdict never denies, so letting it lapse a few + // seconds into a new month is harmless. + failed bool + at time.Time } func newSpendVerdictCache() *spendVerdictCache { @@ -147,22 +161,28 @@ func newSpendVerdictCache() *spendVerdictCache { // fresh enough to trust. An entry expires at the earlier of its TTL and the // next UTC month boundary after its capture: the cap counts calendar-month // spend, so a verdict taken in the old month says nothing about the new one. -func (c *spendVerdictCache) lookup(groupID, cap string) (bool, bool) { +// The negative (store-error) entries use the shorter groupSpendNegativeCacheTTL +// and skip the month check — see spendVerdictEntry.failed. +func (c *spendVerdictCache) lookup(groupID, cap string) (spendVerdictEntry, bool) { c.mu.Lock() defer c.mu.Unlock() e, ok := c.entries[groupID+"\x00"+cap] if !ok { - return false, false + return spendVerdictEntry{}, false } now := time.Now if c.now != nil { now = c.now } t := now() - if t.Sub(e.at) >= groupSpendCacheTTL || !t.Before(nextUTCMonth(e.at)) { - return false, false + ttl := groupSpendCacheTTL + if e.failed { + ttl = groupSpendNegativeCacheTTL + } + if t.Sub(e.at) >= ttl || (!e.failed && !t.Before(nextUTCMonth(e.at))) { + return spendVerdictEntry{}, false } - return e.exhausted, true + return e, true } // store records a fresh verdict for (groupID, cap). @@ -176,6 +196,19 @@ func (c *spendVerdictCache) store(groupID, cap string, exhausted bool) { c.entries[groupID+"\x00"+cap] = spendVerdictEntry{exhausted: exhausted, at: now()} } +// storeFailure records a fail-open entry for (groupID, cap) after a store +// error: for groupSpendNegativeCacheTTL the check admits without querying and +// without waiting, then re-probes the store. +func (c *spendVerdictCache) storeFailure(groupID, cap string) { + c.mu.Lock() + defer c.mu.Unlock() + now := time.Now + if c.now != nil { + now = c.now + } + c.entries[groupID+"\x00"+cap] = spendVerdictEntry{failed: true, at: now()} +} + // nextUTCMonth returns the first instant of the UTC calendar month after t's. func nextUTCMonth(t time.Time) time.Time { u := t.UTC() @@ -213,3 +246,54 @@ func (l *spendFailureLog) logf(log *logging.Logger, format string, args ...inter } log.Error.Printf("%s", msg) } + +// spendFlight collapses concurrent in-flight spend-check reads per (group, +// cap): while one request reads Postgres for a key, every other request on +// the same key waits for that single read instead of issuing its own, so a +// blackholed store with N concurrent capped requests makes one query, not N. +// golang.org/x/sync is only an indirect dependency here, so this is a small +// hand-rolled in-flight map (a mutex and one call struct per key). +type spendFlight struct { + mu sync.Mutex + in map[string]*spendFlightCall +} + +type spendFlightCall struct { + done chan struct{} + result spendFlightResult +} + +type spendFlightResult struct { + exhausted bool + err error +} + +func newSpendFlight() *spendFlight { + return &spendFlight{in: map[string]*spendFlightCall{}} +} + +// do runs fn once per key while any caller is in flight and returns its +// result to every caller: concurrent waiters take the leader's verdict, or +// its error (the fail-open path), without re-running the query. The leader +// reports leader=true so exactly one caller records and logs the outcome. +func (f *spendFlight) do(key string, fn func() (bool, error)) (spendFlightResult, bool) { + f.mu.Lock() + if call, ok := f.in[key]; ok { + f.mu.Unlock() + // The leader's query is bounded by groupSpendQueryBudget, so this + // wait cannot hold the request past one query budget either. + <-call.done + return call.result, false + } + call := &spendFlightCall{done: make(chan struct{})} + f.in[key] = call + f.mu.Unlock() + + call.result.exhausted, call.result.err = fn() + close(call.done) + + f.mu.Lock() + delete(f.in, key) + f.mu.Unlock() + return call.result, true +} diff --git a/internal/admission/spend_integration_test.go b/internal/admission/spend_integration_test.go new file mode 100644 index 0000000..7f30b07 --- /dev/null +++ b/internal/admission/spend_integration_test.go @@ -0,0 +1,163 @@ +//go:build integration + +// Package admission integration test: runs the REAL group spend check +// (groupSpendExhaustedSQL, through PostgresSpendStore) against a LIVE Postgres +// loaded with the production migrations, and pins the verdicts the unit tests +// can only assume: the sqlmock test pins the query text and the fakes hard-code +// the boundary verdicts, so neither would notice the comparison drifting from +// >= to > or the COALESCE being dropped. +// +// Gated behind the `integration` build tag AND a non-empty PHOEBE_TEST_DATABASE_URL. +// Run with: +// +// PHOEBE_TEST_DATABASE_URL=postgres://... go test -tags=integration ./internal/admission/... +package admission + +import ( + "context" + "database/sql" + "os" + "strings" + "testing" + "time" +) + +// spendSchemaDDL returns the production schema applied before the spend test. +// It is loaded from the REAL migration .sql files (0001 through 0008, in apply +// order) rather than a hand-copied group_usage table, so the test runs against +// exactly the column types production has (group_usage.cost is NUMERIC(20,9)) +// and cannot silently drift from them. The DDL runs inside a per-test isolated +// schema (search_path is set by the caller), leaving no residue. +func spendSchemaDDL(t *testing.T) string { + t.Helper() + var b strings.Builder + for _, f := range []string{ + "../../migrations/0001_billing_event.up.sql", + "../../migrations/0002_rating.up.sql", + "../../migrations/0003_io_log.up.sql", + "../../migrations/0004_billing_event_serving_mode.up.sql", + "../../migrations/0005_invoice_grade_attempts.up.sql", + "../../migrations/0006_rollup_grain.up.sql", + "../../migrations/0007_serving_mode_explicit.up.sql", + // 0008 adds group_usage, the rollup the spend check sums. + "../../migrations/0008_group_scopes.up.sql", + } { + ddl, err := os.ReadFile(f) + if err != nil { + t.Fatalf("read migration %s: %v (the integration test applies the REAL "+ + "migration DDL so it can't drift from production)", f, err) + } + b.Write(ddl) + b.WriteString("\n") + } + return b.String() +} + +// TestIntegration_GroupSpendExhausted pins the monthly spend verdict as +// Postgres computes it. Every verdict is read through the store only — the +// test never sums spend in Go — so it proves the comparison itself runs in +// Postgres: +// +// - a zero cap with no group_usage rows is exhausted (the empty month sums +// to NULL, COALESCE makes it 0, and 0 >= 0 denies all paid work); +// - month spend exactly equal to the cap is exhausted (the >= boundary); +// - month spend below the cap is not exhausted; +// - a row in the PRIOR UTC month is excluded from the sum, even when it is +// the last hour before the boundary. +func TestIntegration_GroupSpendExhausted(t *testing.T) { + dsn := os.Getenv("PHOEBE_TEST_DATABASE_URL") + if dsn == "" { + t.Skip("PHOEBE_TEST_DATABASE_URL not set; skipping live-Postgres spend check") + } + ctx := context.Background() + db, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatalf("open: %v", err) + } + defer db.Close() + // search_path is a per-connection setting, so the pool is pinned to one + // connection: the store's query must see the same isolated schema the + // fixture was written into. + db.SetMaxOpenConns(1) + + const sch = "phoebe_admission_spend_it" + exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") + exec(t, db, "CREATE SCHEMA "+sch) + exec(t, db, "SET search_path TO "+sch) + defer func() { exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") }() + exec(t, db, spendSchemaDDL(t)) + + // UTC hour buckets, the same bucketing the rater writes window_start with. + // monthStart is the first hour of the current UTC month; the hour before it + // is the last hour of the prior month. + now := time.Now().UTC() + monthStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) + thisHour := now.Truncate(time.Hour) + priorMonthLastHour := monthStart.Add(-time.Hour) + + const ( + gEmpty = "00000000000000000000000000000000" + gSpend = "a1b2c3d4e5f60718293a4b5c6d7e8f90" + gOther = "00112233445566778899aabbccddeeff" + ) + seed := func(group string, windowStart time.Time, cost string) { + t.Helper() + if _, err := db.ExecContext(ctx, + `INSERT INTO group_usage (group_id, window_start, cost, event_count) + VALUES ($1, $2, $3::numeric, 1)`, + group, windowStart, cost); err != nil { + t.Fatalf("seed group_usage (%s, %s): %v", group, windowStart, err) + } + } + // gSpend's current-month spend is 1.500000001, split across two hours so + // the check must sum rows (when the test runs in the month's first hour + // both land in one bucket, so the second is folded into the first). The + // prior-month row is large enough that counting it would flip every verdict + // below. gOther's spend must not leak into gSpend's sum. + if thisHour.Equal(monthStart) { + seed(gSpend, monthStart, "1.500000001") + } else { + seed(gSpend, monthStart, "1.000000000") + seed(gSpend, thisHour, "0.500000001") + } + seed(gSpend, priorMonthLastHour, "1000") + seed(gOther, monthStart, "1000") + + store := NewPostgresSpendStore(db) + for _, tc := range []struct { + name string + group string + cap string + want bool + }{ + // 1. Zero cap, no rows at all for the group: exhausted. + {"zero cap, empty month denies", gEmpty, "0", true}, + // A positive cap over an empty month is unreached (COALESCE to 0). + {"positive cap, empty month admits", gEmpty, "0.000000001", false}, + // 2. Spend equal to the cap, to the ninth decimal: exhausted. Equality + // at this scale only holds if the comparison is NUMERIC in Postgres. + {"spend == cap denies", gSpend, "1.500000001", true}, + // 3. Spend below the cap by the smallest NUMERIC(20,9) step: admitted. + {"spend < cap admits", gSpend, "1.500000002", false}, + {"spend > cap denies", gSpend, "1.5", true}, + // 4. The prior-month row (1000) is excluded: were it counted, spend + // would be 1001.500000001 and this cap would be exhausted. + {"prior month excluded", gSpend, "1001", false}, + } { + got, err := store.GroupSpendExhausted(ctx, tc.group, tc.cap) + if err != nil { + t.Fatalf("%s: GroupSpendExhausted(%s, %s): %v", tc.name, tc.group, tc.cap, err) + } + if got != tc.want { + t.Errorf("%s: GroupSpendExhausted(%s, %s) = %v, want %v", + tc.name, tc.group, tc.cap, got, tc.want) + } + } +} + +func exec(t *testing.T, db *sql.DB, q string) { + t.Helper() + if _, err := db.Exec(q); err != nil { + t.Fatalf("exec %q: %v", q, err) + } +} From 2864330990b3882472113fc96b509a51f450906e Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 09:07:26 +0000 Subject: [PATCH 10/17] admission: one spend-query budget per request; proxy: skip admission for all-unlimited group envelopes; recovery: validate member_group_ids Bound the whole group spend check by a single 250ms budget instead of one per group, so many uncached capped groups against a slow store wait at most one budget. Detach the coalesced spend read from the leader's cancellation so one client's disconnect cannot end the read every waiter shares. Under admission.enabled=false, engage the admission store only when a group scope carries a rate limit or spend cap. At the recovery boundary, reject member_group_ids lists that violate the frozen envelope contract, since recovered evidence bypasses the proxy's strict parse and the stored list is what rating attributes into group_usage. --- internal/admission/admission.go | 22 ++++- internal/admission/group_test.go | 105 ++++++++++++++++++++++ internal/admission/spend.go | 15 ++-- internal/proxy/admission.go | 17 +++- internal/proxy/group_scopes_test.go | 110 ++++++++++++++++++++++- internal/proxy/proxy.go | 12 ++- internal/recovery/recovery.go | 52 +++++++++++ internal/recovery/recovery_test.go | 131 +++++++++++++++++++++++++--- 8 files changed, 439 insertions(+), 25 deletions(-) diff --git a/internal/admission/admission.go b/internal/admission/admission.go index dbdf6f8..8cdf10a 100644 --- a/internal/admission/admission.go +++ b/internal/admission/admission.go @@ -481,6 +481,12 @@ func (a *RedisAdmitter) checkGroupSpend(ctx context.Context, scopes []GroupScope } return nil } + // One query budget covers the whole check, not one per group: a request + // carrying many uncached capped groups against a slow store waits at most + // one groupSpendQueryBudget in total. A group whose read cannot finish + // before the shared deadline fails open like any other store error. The + // budget starts at the first cache miss, so all-cached checks pay nothing. + var budgetCtx context.Context for _, gs := range scopes { if gs.SpendCap == "" { continue @@ -497,13 +503,23 @@ func (a *RedisAdmitter) checkGroupSpend(ctx context.Context, scopes []GroupScope } continue } + if budgetCtx == nil { + var cancelBudget context.CancelFunc + budgetCtx, cancelBudget = context.WithTimeout(ctx, groupSpendQueryBudget) + defer cancelBudget() + } + deadline, _ := budgetCtx.Deadline() // Collapse concurrent misses on one (group, cap) into a single store // read: waiters take the leader's result (verdict, or fail-open on // error) instead of each running their own query against a slow store. - result, leader := a.spendFlight.do(gs.GroupID+"\x00"+gs.SpendCap, func() (bool, error) { + result, leader := a.spendFlight.do(budgetCtx, gs.GroupID+"\x00"+gs.SpendCap, func() (bool, error) { // Bounded: a hung store read fails open (deadline exceeded is just - // another store error) instead of holding the request. - queryCtx, cancel := context.WithTimeout(ctx, groupSpendQueryBudget) + // another store error) instead of holding the request. Detached + // from the leader's cancellation: other requests are waiting on + // this read, and one client disconnecting must not end it and + // hand every waiter a free admit. The shared deadline still + // bounds it. + queryCtx, cancel := context.WithDeadline(context.WithoutCancel(ctx), deadline) defer cancel() return a.spendStore.GroupSpendExhausted(queryCtx, gs.GroupID, gs.SpendCap) }) diff --git a/internal/admission/group_test.go b/internal/admission/group_test.go index 4c73564..af95488 100644 --- a/internal/admission/group_test.go +++ b/internal/admission/group_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "errors" + "fmt" "regexp" "strings" "sync" @@ -173,6 +174,17 @@ func TestGroupSpendCapAdmitsBelowCap(t *testing.T) { if store.calls != 1 { t.Fatalf("store calls = %d, want 1", store.calls) } + + // For contrast: a scope without a spend cap never reads the store. + lease, err = a.Admit(context.Background(), groupRequest("org-b", "m", + groupScope("ffffffffffffffffffffffffffffffff", nil, ""))) + if err != nil { + t.Fatalf("uncapped admit: %v, want admitted", err) + } + _ = lease.Complete(context.Background(), 0) + if store.calls != 1 { + t.Fatalf("store calls = %d, want still 1 (an empty SpendCap skips the store)", store.calls) + } } // TestGroupSpendVerdictCachedPerGroupAndCap: the verdict is cached ~60s per @@ -403,6 +415,99 @@ func TestGroupSpendConcurrentMissSingleflight(t *testing.T) { } } +// TestGroupSpendSharedBudgetAcrossGroups: many uncached capped groups against +// a store that never answers share ONE query budget — the request is held +// about one groupSpendQueryBudget in total, not one per group — and admits. +func TestGroupSpendSharedBudgetAcrossGroups(t *testing.T) { + a, _ := testAdmitter(t, config.AdmissionSettings{}) + a.WithGroupSpend(blockingGroupSpendStore{}, logging.New(logging.ERROR)) + + scopes := make([]GroupScope, 16) + for i := range scopes { + scopes[i] = groupScope(fmt.Sprintf("%032x", i+1), nil, "100") + } + start := time.Now() + lease, err := a.Admit(context.Background(), groupRequest("org-sb", "m", scopes...)) + elapsed := time.Since(start) + if err != nil { + t.Fatalf("16-group hung-store admit: %v, want fail-open admission", err) + } + _ = lease.Complete(context.Background(), 0) + if elapsed >= 2*groupSpendQueryBudget { + t.Fatalf("16-group admit took %v, want under %v (one shared query budget)", elapsed, 2*groupSpendQueryBudget) + } +} + +// releasedSpendStore answers exhausted=true once released, or the context's +// error if the context ends first; entered signals the first read began. +type releasedSpendStore struct { + entered chan struct{} + release chan struct{} + once sync.Once + mu sync.Mutex + calls int +} + +func (s *releasedSpendStore) GroupSpendExhausted(ctx context.Context, _, _ string) (bool, error) { + s.mu.Lock() + s.calls++ + s.mu.Unlock() + s.once.Do(func() { close(s.entered) }) + select { + case <-s.release: + return true, nil + case <-ctx.Done(): + return false, ctx.Err() + } +} + +// TestGroupSpendLeaderCancelDoesNotCancelCoalescedQuery: the leader's client +// disconnecting mid-read must not end the coalesced read. A waiter with a live +// context still gets the real verdict (the over-cap 429), from that one read. +func TestGroupSpendLeaderCancelDoesNotCancelCoalescedQuery(t *testing.T) { + a, _ := testAdmitter(t, config.AdmissionSettings{}) + store := &releasedSpendStore{entered: make(chan struct{}), release: make(chan struct{})} + a.WithGroupSpend(store, logging.New(logging.ERROR)) + + leaderCtx, cancelLeader := context.WithCancel(context.Background()) + leaderDone := make(chan struct{}) + go func() { + defer close(leaderDone) + if lease, err := a.Admit(leaderCtx, groupRequest("org-lc", "m", + groupScope(testGID, nil, "100"))); err == nil { + _ = lease.Complete(context.Background(), 0) + } + }() + <-store.entered + cancelLeader() + + waiterErr := make(chan error, 1) + go func() { + lease, err := a.Admit(context.Background(), groupRequest("org-lc", "m", + groupScope(testGID, nil, "100"))) + if err == nil { + _ = lease.Complete(context.Background(), 0) + } + waiterErr <- err + }() + // Let the waiter join the in-flight read, then let the read answer. + time.Sleep(20 * time.Millisecond) + close(store.release) + + err := <-waiterErr + <-leaderDone + var rejected *Rejected + if !errors.As(err, &rejected) || rejected.Dimension != "monthly_spend" || !rejected.Contractual { + t.Fatalf("waiter admit: %v, want the 429 monthly_spend rejection from the coalesced read", err) + } + store.mu.Lock() + calls := store.calls + store.mu.Unlock() + if calls != 1 { + t.Fatalf("store calls = %d, want 1 (the waiter shares the leader's read)", calls) + } +} + // TestGroupSpendParentCancelNotLoggedAsBypass: a canceled request context // (client gone) makes the spend check moot — it must NOT log the throttled // 'monthly spend cap NOT enforced' ERROR, which a genuine outage needs. diff --git a/internal/admission/spend.go b/internal/admission/spend.go index bfa4057..5e09033 100644 --- a/internal/admission/spend.go +++ b/internal/admission/spend.go @@ -276,14 +276,19 @@ func newSpendFlight() *spendFlight { // result to every caller: concurrent waiters take the leader's verdict, or // its error (the fail-open path), without re-running the query. The leader // reports leader=true so exactly one caller records and logs the outcome. -func (f *spendFlight) do(key string, fn func() (bool, error)) (spendFlightResult, bool) { +// A waiter stops waiting when its own ctx is done (its share of the spend +// query budget ran out, or its client left) and gets ctx's error, which the +// caller treats as fail-open; the leader's read keeps running for the others. +func (f *spendFlight) do(ctx context.Context, key string, fn func() (bool, error)) (spendFlightResult, bool) { f.mu.Lock() if call, ok := f.in[key]; ok { f.mu.Unlock() - // The leader's query is bounded by groupSpendQueryBudget, so this - // wait cannot hold the request past one query budget either. - <-call.done - return call.result, false + select { + case <-call.done: + return call.result, false + case <-ctx.Done(): + return spendFlightResult{err: ctx.Err()}, false + } } call := &spendFlightCall{done: make(chan struct{})} f.in[key] = call diff --git a/internal/proxy/admission.go b/internal/proxy/admission.go index 1ed52bb..24deae7 100644 --- a/internal/proxy/admission.go +++ b/internal/proxy/admission.go @@ -446,7 +446,8 @@ func anyScopedLimitPresent(id identity.Identity) bool { // malformed group scope envelope fails closed in parseTrustedGroupScopes), // and an underivable graph scope all answer 503. // Under admission.enabled=false the admission store itself is only consulted -// when a contract limit header is present (see the Admit call in proxy.go). +// when a limit is present to enforce — a contract limit header or a group +// scope carrying a rate limit or spend cap (see the Admit call in proxy.go). func trustedPolicyEnvelopePresent(id identity.Identity) bool { return id.OwnerID != "" || anyScopedLimitPresent(id) || id.GroupScopes != "" } @@ -597,6 +598,20 @@ func parseGroupRate(gid, dimension, value string) (*int64, error) { return &limit, nil } +// anyGroupScopeEnforced reports whether at least one group scope carries +// something to enforce: a non-empty spend cap or any non-nil per-minute rate +// limit. An all-unlimited group envelope (every rate empty, no spend cap) +// enforces nothing, so under admission.enabled=false it must not touch the +// admission store — exactly like a request with no envelope. +func anyGroupScopeEnforced(scopes []admission.GroupScope) bool { + for _, scope := range scopes { + if scope.SpendCap != "" || scope.Limits.Any() { + return true + } + } + return false +} + // legacyQuotaHeaderNames are the five single-scope quota headers that R8 // removed from the trusted envelope. Phoebe never reads them for a trust or // limit decision; they are listed here ONLY so the proxy can log that a diff --git a/internal/proxy/group_scopes_test.go b/internal/proxy/group_scopes_test.go index 796238f..bd67290 100644 --- a/internal/proxy/group_scopes_test.go +++ b/internal/proxy/group_scopes_test.go @@ -231,12 +231,28 @@ func groupScopesRequest(t *testing.T, upstream *url.URL, envelope string) *http. // proxyWithGroupScopes builds the admission-enabled proxy over a real // (frozen-clock) miniredis, with an optional spend store. func proxyWithGroupScopes(t *testing.T, spend admission.GroupSpendStore) (*Server, *recordingEmitter) { + t.Helper() + return proxyWithGroupScopesEnabled(t, true, spend) +} + +// proxyWithGroupScopesEnabled builds the proxy over a real (frozen-clock) +// miniredis with the given admission.enabled flag. The admitter runs with the +// effective settings — operator tiers cleared under admission.enabled=false, +// exactly as a default chart install renders them — and an optional spend +// store. +func proxyWithGroupScopesEnabled(t *testing.T, enabled bool, spend admission.GroupSpendStore) (*Server, *recordingEmitter) { t.Helper() mr := frozenMiniredis(t) cfg := proxyAdmissionConfig(10) + cfg.Enabled = enabled + cfg.ValkeyAddr = mr.Addr() + eff, ok := (&config.Settings{Admission: cfg}).EffectiveAdmission() + if !ok { + t.Fatal("effective admission settings: no store configured") + } client := redis.NewClient(&redis.Options{Addr: mr.Addr()}) t.Cleanup(func() { _ = client.Close() }) - admitter := admission.New(client, cfg) + admitter := admission.New(client, eff) if spend != nil { admitter.WithGroupSpend(spend, logging.New(logging.ERROR)) } @@ -422,3 +438,95 @@ func sharedGroupBackend(t *testing.T) (*url.URL, *int) { } return up, &hits } + +// TestSharedGroupScopesAdmissionTriggerWithAdmissionFlagOff pins the trigger: +// under admission.enabled=false the Admit round-trip is engaged by a scope +// that carries enforcement, not by the mere presence of the group envelope. +// Every org/owner scoped limit header is absent, so only the group scopes can +// engage the trigger. An all-unlimited group envelope performs no admission +// call; an envelope with a rate limit or a spend cap still performs exactly +// one, unchanged from the enabled=true path. +func TestSharedGroupScopesAdmissionTriggerWithAdmissionFlagOff(t *testing.T) { + up, _ := sharedGroupBackend(t) + cfg := proxyAdmissionConfig(10) + cfg.Enabled = false + send := func(a *countingAdmitter, envelope string) *httptest.ResponseRecorder { + s := New(&config.Settings{Admission: cfg}, logging.New(logging.ERROR), &recordingEmitter{}).WithAdmitter(a) + req := groupScopesRequest(t, up, envelope) + for _, header := range identity.ScopedRateLimitHeaders { + req.Header.Del(header) + } + rr := httptest.NewRecorder() + s.Handler().ServeHTTP(rr, req) + return rr + } + + // The all-unlimited envelope (every rate empty, no spend cap) enforces + // nothing: no Admit, no Valkey reservation, no completion round-trip. + a := &countingAdmitter{} + if rr := send(a, groupScopesHeader); rr.Code != http.StatusOK || a.calls != 0 { + t.Fatalf("all-unlimited envelope: status=%d Admit calls=%d, want 200 and no admission round-trip", rr.Code, a.calls) + } + + // A rate-limited group scope engages admission exactly as a contract + // limit header does. + a = &countingAdmitter{} + if rr := send(a, groupScopesValue(testGroupA, "30,,,", "")); rr.Code != http.StatusOK || a.calls != 1 { + t.Fatalf("rate-limited group scope: status=%d Admit calls=%d, want 200 and exactly one Admit", rr.Code, a.calls) + } + + // A group carrying only a spend cap engages admission too. + a = &countingAdmitter{} + if rr := send(a, groupScopesValue(testGroupA, ",,,", "100")); rr.Code != http.StatusOK || a.calls != 1 { + t.Fatalf("spend-capped group scope: status=%d Admit calls=%d, want 200 and exactly one Admit", rr.Code, a.calls) + } +} + +// TestSharedGroupScopesEnforcedWithAdmissionFlagOff is the envelope-driven +// enforcement test for the item-13 posture: admission.enabled=false with ONLY +// X-Saturn-Group-Scopes stamped (no org/owner limit headers). A group zero +// cap answers 429 + Retry-After before upstream, and a request within the +// group limits is admitted — the group-scope mirror of the contract-limits +// enabled=false cases in contract_limits_test.go. +func TestSharedGroupScopesEnforcedWithAdmissionFlagOff(t *testing.T) { + up, hits := sharedGroupBackend(t) + s, em := proxyWithGroupScopesEnabled(t, false, nil) + + // Only the group envelope: every org/owner scoped limit header absent. + envelopeOnly := func(envelope string) *http.Request { + req := groupScopesRequest(t, up, envelope) + for _, header := range identity.ScopedRateLimitHeaders { + req.Header.Del(header) + } + return req + } + + // A group zero cap blocks every request with the contractual 429. + rr := httptest.NewRecorder() + s.Handler().ServeHTTP(rr, envelopeOnly(groupScopesValue(testGroupA, "0,,,", ""))) + if rr.Code != http.StatusTooManyRequests { + t.Fatalf("zero-cap status=%d, want 429", rr.Code) + } + if rr.Header().Get("Retry-After") == "" { + t.Fatal("zero-cap response missing Retry-After") + } + + // A request within the group limits is admitted and forwarded. + rr = httptest.NewRecorder() + s.Handler().ServeHTTP(rr, envelopeOnly(groupScopesValue(testGroupA, "30,,,", ""))) + if rr.Code != http.StatusOK { + t.Fatalf("within-limits status=%d, want 200", rr.Code) + } + if *hits != 1 { + t.Fatalf("upstream hits=%d, want 1 (the zero cap never forwards)", *hits) + } + evs := em.waitForEvents(1, time.Second) + if len(evs) != 1 { + t.Fatalf("emitted events=%d, want 1", len(evs)) + } + // Membership attribution rides the emitted event whether or not admission + // ran: skipping the store for an unlimited envelope never drops it. + if got := evs[0].MemberGroupIDs; len(got) != 1 || got[0] != testGroupA { + t.Fatalf("MemberGroupIDs = %v, want [%s]", got, testGroupA) + } +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 679bd52..10c6864 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -576,7 +576,9 @@ func (s *Server) handleProxy(w http.ResponseWriter, r *http.Request) { // Contract limits therefore engage whenever Atlas stamped them, with or // without operator tiers (ruling R12 clarification). Under // admission.enabled=false the store itself is consulted only when a - // contract limit header is present (the Admit call below); a limit + // contract limit header is present (the Admit call below; a group + // envelope engages it only when a group scope carries a rate limit or + // spend cap); a limit // header without the owner-id anchor still fails closed in // parseTrustedRateLimits. Only a shared request with NO envelope at all, // under admission.enabled=false, skips admission entirely: the @@ -706,9 +708,11 @@ func (s *Server) handleProxy(w http.ResponseWriter, r *http.Request) { } // Under admission.enabled=false a request whose envelope carries no // limit is unlimited in every scope, so it does not touch the store: no - // Admit, no lease renewal, no completion. Every fail-closed check above - // has already run. - if enforceAdmission && (s.settings.Admission.Enabled || organizationLimits.Any() || ownerLimits.Any() || len(groupScopes) > 0) { + // Admit, no lease renewal, no completion. A group envelope engages the + // store only when at least one group scope carries enforcement — a + // spend cap or a rate limit; an all-unlimited envelope enforces + // nothing. Every fail-closed check above has already run. + if enforceAdmission && (s.settings.Admission.Enabled || organizationLimits.Any() || ownerLimits.Any() || anyGroupScopeEnforced(groupScopes)) { admitted, err = s.admitter.Admit(r.Context(), admission.Request{ Graph: graph, Organization: id.OrgID, Owner: id.OwnerID, Model: estimate.Model, PromptBytes: originalPromptBytes, EstimatedInputTokens: estimate.InputTokens, diff --git a/internal/recovery/recovery.go b/internal/recovery/recovery.go index dee58e0..90a2abc 100644 --- a/internal/recovery/recovery.go +++ b/internal/recovery/recovery.go @@ -33,6 +33,13 @@ type Evidence struct { Duplicates int } +const ( + // maxMemberGroupIDs and groupIDHexLength are the frozen envelope contract + // (ruling 3) applied to recovered evidence. + maxMemberGroupIDs = 16 + groupIDHexLength = 32 +) + // Digest returns a stable SHA-256 digest binding the COMPLETE validated, // de-duplicated event set — not merely its request-id set. -apply compares this // value so that the artifact an operator reviewed during dry-run is provably @@ -344,9 +351,54 @@ func validate(ev metering.Event) error { if ev.StatusCode != 0 && (ev.StatusCode < 100 || ev.StatusCode > 599) { return fmt.Errorf("status_code %d is outside the database range 100..599", ev.StatusCode) } + if err := validateMemberGroupIDs(ev.MemberGroupIDs); err != nil { + return err + } + return nil +} + +// validateMemberGroupIDs checks a recovered event's membership list against +// the frozen envelope contract (ruling 3): at most 16 entries, each unique, +// each exactly 32 lowercase hexadecimal characters. The normal request path +// validates the envelope strictly at the proxy, but recovered evidence +// bypasses that parse, and the stored list is what the rater attributes into +// group_usage — an attacker-chosen or malformed list would land in the +// rollup and could deny real spend caps. Reject it here, on the same error +// path as every other malformed recovered event. +func validateMemberGroupIDs(ids []string) error { + if len(ids) > maxMemberGroupIDs { + return fmt.Errorf("member_group_ids has %d entries; the envelope contract allows at most %d", + len(ids), maxMemberGroupIDs) + } + seen := make(map[string]struct{}, len(ids)) + for _, id := range ids { + if !isGroupID(id) { + return fmt.Errorf("member_group_ids entry %q is not exactly 32 lowercase hexadecimal characters", id) + } + if _, dup := seen[id]; dup { + return fmt.Errorf("member_group_ids contains duplicate entry %q", id) + } + seen[id] = struct{}{} + } return nil } +// isGroupID reports whether s is a 32-character lowercase hex group id, the +// frozen envelope grammar (ruling 3). Uppercase or any other spelling is +// not a group id. +func isGroupID(s string) bool { + if len(s) != groupIDHexLength { + return false + } + for i := 0; i < len(s); i++ { + c := s[i] + if !('0' <= c && c <= '9' || 'a' <= c && c <= 'f') { + return false + } + } + return true +} + func printableID(name, value string) error { if value == "" { return fmt.Errorf("%s is required", name) diff --git a/internal/recovery/recovery_test.go b/internal/recovery/recovery_test.go index db72330..474224b 100644 --- a/internal/recovery/recovery_test.go +++ b/internal/recovery/recovery_test.go @@ -57,18 +57,127 @@ func TestLoadJSONLAndFloorLogsDedupeIdenticalEvents(t *testing.T) { } func TestLoadRejectsConflictingDuplicate(t *testing.T) { - ev1 := testEvent("phoebe-same") - ev2 := ev1 - ev2.CompletionTokens = 1 - one, _ := json.Marshal(ev1) - two, _ := json.Marshal(ev2) - path := filepath.Join(t.TempDir(), "evidence.jsonl") - if err := os.WriteFile(path, append(append(one, '\n'), two...), 0o600); err != nil { - t.Fatal(err) + gidA := strings.Repeat("a", 32) + gidB := strings.Repeat("b", 32) + + t.Run("differing token counts", func(t *testing.T) { + ev1 := testEvent("phoebe-same") + ev2 := ev1 + ev2.CompletionTokens = 1 + one, _ := json.Marshal(ev1) + two, _ := json.Marshal(ev2) + path := filepath.Join(t.TempDir(), "evidence.jsonl") + if err := os.WriteFile(path, append(append(one, '\n'), two...), 0o600); err != nil { + t.Fatal(err) + } + _, err := Load(path) + if err == nil || !strings.Contains(err.Error(), "conflicting duplicate") { + t.Fatalf("expected conflicting duplicate error, got %v", err) + } + }) + + // A regression that dropped MemberGroupIDs from the canonical comparison + // would still pass the token-count case above, so the membership list + // needs its own conflicting-duplicate case. + t.Run("differing member group ids", func(t *testing.T) { + ev1 := testEvent("phoebe-same") + ev1.MemberGroupIDs = []string{gidA} + ev2 := ev1 + ev2.MemberGroupIDs = []string{gidB} + one, _ := json.Marshal(ev1) + two, _ := json.Marshal(ev2) + path := filepath.Join(t.TempDir(), "evidence.jsonl") + if err := os.WriteFile(path, append(append(one, '\n'), two...), 0o600); err != nil { + t.Fatal(err) + } + _, err := Load(path) + if err == nil || !strings.Contains(err.Error(), "conflicting duplicate") { + t.Fatalf("expected conflicting duplicate error, got %v", err) + } + }) + + // Identical records, membership list included, must still dedupe as + // before rather than be reported as a conflict. + t.Run("identical events with identical memberships dedupe", func(t *testing.T) { + ev1 := testEvent("phoebe-same") + ev1.MemberGroupIDs = []string{gidA} + ev2 := ev1 + one, _ := json.Marshal(ev1) + two, _ := json.Marshal(ev2) + path := filepath.Join(t.TempDir(), "evidence.jsonl") + if err := os.WriteFile(path, append(append(one, '\n'), two...), 0o600); err != nil { + t.Fatal(err) + } + evidence, err := Load(path) + if err != nil { + t.Fatalf("identical records must dedupe, got %v", err) + } + if evidence.Duplicates != 1 || len(evidence.Events) != 1 { + t.Fatalf("evidence = %+v, want one event with one duplicate", evidence) + } + }) +} + +// TestLoadRejectsMalformedMemberGroupIDs pins the envelope contract (ruling 3) +// at the recovery boundary: recovered evidence bypasses the proxy's strict +// envelope parse, and the stored list is what the rater attributes into +// group_usage, so a malformed, duplicate, or oversize list must be refused +// here rather than replayed. +func TestLoadRejectsMalformedMemberGroupIDs(t *testing.T) { + validGID := func(n int) string { + return fmt.Sprintf("%032x", n) + } + sixteenValid := func() []string { + ids := make([]string, 16) + for i := range ids { + ids[i] = validGID(i) + } + return ids } - _, err := Load(path) - if err == nil || !strings.Contains(err.Error(), "conflicting duplicate") { - t.Fatalf("expected conflicting duplicate error, got %v", err) + + cases := []struct { + name string + ids []string + wantErr bool + }{ + {"no memberships is valid", nil, false}, + {"wrong length id", []string{strings.Repeat("a", 31)}, true}, + {"long id", []string{strings.Repeat("a", 33)}, true}, + {"non-hex id", []string{strings.Repeat("g", 32)}, true}, + {"uppercase id", []string{strings.ToUpper(strings.Repeat("a", 32))}, true}, + {"duplicate ids within one event", []string{validGID(1), validGID(1)}, true}, + {"more than 16 ids", append(sixteenValid(), validGID(16)), true}, + {"16 valid ids", sixteenValid(), false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + ev := testEvent("req-groups") + ev.MemberGroupIDs = tc.ids + data, err := json.Marshal(ev) + if err != nil { + t.Fatal(err) + } + path := filepath.Join(t.TempDir(), "evidence.jsonl") + if err := os.WriteFile(path, data, 0o600); err != nil { + t.Fatal(err) + } + evidence, err := Load(path) + if tc.wantErr { + if err == nil { + t.Fatalf("member_group_ids %v was accepted; the envelope contract rejects it", tc.ids) + } + if !strings.Contains(err.Error(), "member_group_ids") { + t.Fatalf("error %q does not name member_group_ids", err) + } + return + } + if err != nil { + t.Fatalf("member_group_ids %v must be accepted: %v", tc.ids, err) + } + if len(evidence.Events) != 1 { + t.Fatalf("events = %+v, want one event", evidence.Events) + } + }) } } From f24050ef23ad7c995a13fcf7663df019d03f3f64 Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 09:56:18 +0000 Subject: [PATCH 11/17] tests: prove the group-attribution gates and spend-check boundaries; docs: correct the 0008 rollout and rollback claims Extend the group_usage attribution integration test so every per-event withheld gate is exercised against a shared money grain instead of vacuously, run the spend integration test under a non-UTC session timezone, pin the throttled fail-open spend log, disambiguate the group settlement tests from a kept reservation, and add a real-Valkey group scope check. Correct the 0008 rollout text (an old rater runs on the new schema and writes no group_usage) and mark the 0008 rollback lossy: dropping member_group_ids destroys the membership evidence and month-to-date group spend restarts lower. --- internal/admission/group_test.go | 113 +++++++++ .../admission/real_valkey_integration_test.go | 75 ++++++ internal/admission/spend.go | 9 +- internal/admission/spend_integration_test.go | 20 +- .../rating/group_rollup_integration_test.go | 231 ++++++++++++++++++ internal/rating/rollback_integration_test.go | 136 +++++++++++ migrations/0008_group_scopes.down.sql | 10 +- migrations/0008_group_scopes.up.sql | 13 +- migrations/README.md | 24 +- 9 files changed, 616 insertions(+), 15 deletions(-) create mode 100644 internal/rating/rollback_integration_test.go diff --git a/internal/admission/group_test.go b/internal/admission/group_test.go index af95488..d040da7 100644 --- a/internal/admission/group_test.go +++ b/internal/admission/group_test.go @@ -656,3 +656,116 @@ func TestGroupScopesSettleThroughCompleteUsage(t *testing.T) { t.Fatalf("third admit (window at 50/50): %v, want a contractual rejection", err) } } + +// TestGroupScopesSettleZeroUsageRefundsReservation: the settlement test above +// cannot tell settlement from a reservation left charged (each request +// reserves 20 output tokens, so an unsettled window would also read 20, 40, +// 60). Here every request completes with zero generated tokens: against a cap +// of 50, five requests admit only if each settlement refunds its 20-token +// reservation — without settlement the third would be rejected at 60/50. +func TestGroupScopesSettleZeroUsageRefundsReservation(t *testing.T) { + a, _ := spendTestAdmitter(t) + cap := int64(50) + r := groupRequest("org-z0", "m", GroupScope{GroupID: testGID, + Limits: RateLimits{GeneratedTokens: &cap}}) + + for i := 0; i < 5; i++ { + lease, err := a.Admit(context.Background(), r) + if err != nil { + t.Fatalf("admit %d: %v, want admitted (each zero-usage Complete refunds the reservation)", i+1, err) + } + if err := lease.Complete(context.Background(), 0); err != nil { + t.Fatalf("Complete %d: %v", i+1, err) + } + } +} + +// TestGroupScopesSettleChargesActualUsage: settling 45 generated tokens leaves +// the window at 45/50, so the next 20-token reservation is rejected — even +// though an unsettled window (holding only the first 20-token reservation) +// would have fit it at 40/50. +func TestGroupScopesSettleChargesActualUsage(t *testing.T) { + a, _ := spendTestAdmitter(t) + cap := int64(50) + r := groupRequest("org-l45", "m", GroupScope{GroupID: testGID, + Limits: RateLimits{GeneratedTokens: &cap}}) + + lease, err := a.Admit(context.Background(), r) + if err != nil { + t.Fatalf("first admit: %v", err) + } + if err := lease.Complete(context.Background(), 45); err != nil { + t.Fatalf("Complete: %v", err) + } + + _, err = a.Admit(context.Background(), r) + var rejected *Rejected + if !errors.As(err, &rejected) || !rejected.Contractual || rejected.Scope != "group:"+testGID { + t.Fatalf("admit after settling 45/50: %v, want a contractual group rejection", err) + } +} + +// TestSpendFailureLogThrottle: the first fail-open line of an incident writes +// immediately; repeats inside spendFailureLogInterval are counted and silent; +// the first call after the interval writes and carries the suppressed count. +func TestSpendFailureLogThrottle(t *testing.T) { + log := logging.New(logging.ERROR) + var buf bytes.Buffer + log.Error.SetOutput(&buf) + now := time.Now() + l := &spendFailureLog{now: func() time.Time { return now }} + + l.logf(log, "bypass %d", 1) + if got := strings.Count(buf.String(), "bypass"); got != 1 { + t.Fatalf("lines after first logf = %d (%q), want 1", got, buf.String()) + } + if strings.Contains(buf.String(), "suppressed") { + t.Fatalf("first line = %q, want no suppressed count", buf.String()) + } + + for i := 0; i < 3; i++ { + now = now.Add(spendFailureLogInterval / 4) + l.logf(log, "bypass %d", i+2) + } + if got := strings.Count(buf.String(), "bypass"); got != 1 { + t.Fatalf("lines inside the interval = %d (%q), want still 1 (repeats are silent)", got, buf.String()) + } + + now = now.Add(spendFailureLogInterval) + l.logf(log, "bypass %d", 5) + out := buf.String() + if got := strings.Count(out, "bypass"); got != 2 { + t.Fatalf("lines after the interval = %d (%q), want 2", got, out) + } + if !strings.Contains(out, "bypass 5 (+3 similar suppressed)") { + t.Fatalf("post-interval line = %q, want it to carry the 3 suppressed repeats", out) + } +} + +// TestGroupSpendStoreErrorLogsOnce: a plain (non-timeout) store error writes +// exactly one fail-open ERROR line for the incident; the negative-cached +// repeats do not add more. +func TestGroupSpendStoreErrorLogsOnce(t *testing.T) { + a, store := spendTestAdmitter(t) + log := logging.New(logging.ERROR) + var buf bytes.Buffer + log.Error.SetOutput(&buf) + a.WithGroupSpend(store, log) + store.err = errors.New("postgres unreachable") + + for i := 0; i < 3; i++ { + lease, err := a.Admit(context.Background(), groupRequest("org-lo", "m", + groupScope(testGID, nil, "100"))) + if err != nil { + t.Fatalf("store-error admit %d: %v, want fail-open admission", i+1, err) + } + _ = lease.Complete(context.Background(), 0) + } + out := buf.String() + if got := strings.Count(out, "monthly spend cap NOT enforced (fail open)"); got != 1 { + t.Fatalf("fail-open ERROR lines = %d (%q), want exactly 1", got, out) + } + if !strings.Contains(out, "postgres unreachable") { + t.Fatalf("fail-open line = %q, want it to name the store error", out) + } +} diff --git a/internal/admission/real_valkey_integration_test.go b/internal/admission/real_valkey_integration_test.go index 856af8a..745c2d7 100644 --- a/internal/admission/real_valkey_integration_test.go +++ b/internal/admission/real_valkey_integration_test.go @@ -414,3 +414,78 @@ func TestRealValkeyContractOnlyBurst(t *testing.T) { }) } } + +// TestRealValkeyGroupScopes runs the group fan-out through real Valkey's Lua: +// one request carrying 16 rate-limited group scopes (the envelope maximum) +// admits and settles, each group's window is charged the settled usage rather +// than the reservation, and a zero-capped group is a contractual rejection. +func TestRealValkeyGroupScopes(t *testing.T) { + addr := os.Getenv("PHOEBE_TEST_ADMISSION_VALKEY_ADDR") + if addr == "" { + t.Skip("PHOEBE_TEST_ADMISSION_VALKEY_ADDR not set; skipping real-Valkey group scopes") + } + ctx, cancel := context.WithTimeout(context.Background(), 75*time.Second) + defer cancel() + client := redis.NewClient(&redis.Options{Addr: addr}) + if err := client.Ping(ctx).Err(); err != nil { + t.Fatalf("ping real Valkey: %v", err) + } + cfg := config.AdmissionSettings{ + KeyPrefix: fmt.Sprintf("phoebe-admission-groups-%d", time.Now().UnixNano()), + LeaseTTL: time.Second, + Platform: limits(64), + } + a := New(client, cfg) + t.Cleanup(func() { + _ = client.Del(context.Background(), a.counters, a.leases, a.expiries, a.windowExpiries).Err() + _ = client.Close() + }) + + // The windows are fixed per minute; start clear of a boundary so the + // settle and the per-group checks below land in the same window. + if s := time.Now().Second(); s >= 50 { + time.Sleep(time.Duration(61-s) * time.Second) + } + + cap := int64(50) + scopes := make([]GroupScope, 16) + for i := range scopes { + scopes[i] = GroupScope{GroupID: fmt.Sprintf("%032x", i+1), + Limits: RateLimits{GeneratedTokens: &cap}} + } + r := request("org-g", "model") + r.GroupScopes = scopes + + lease, err := a.Admit(ctx, r) + if err != nil { + t.Fatalf("admit 16-group request: %v", err) + } + // Settle 45 generated tokens: every group window now reads 45/50. A window + // left holding only the 20-token reservation would read 20/50 instead. + if err := lease.Complete(ctx, 45); err != nil { + t.Fatalf("settle 16-group request: %v", err) + } + + // Each group alone: a 6-token reservation fits an unsettled 20/50 window + // but not a settled 45/50 one, so a rejection naming that group proves its + // own window settled. + for _, gs := range scopes { + probe := request("org-g", "model") + probe.ReservedOutputTokens = 6 + probe.GroupScopes = []GroupScope{gs} + _, err := a.Admit(ctx, probe) + rejected, ok := err.(*Rejected) + if !ok || !rejected.Contractual || rejected.Scope != "group:"+gs.GroupID { + t.Fatalf("probe for group %s after settling 45/50 = %T %v, want a contractual rejection naming the group", gs.GroupID, err, err) + } + } + + zero := int64(0) + zr := request("org-g", "model") + zr.GroupScopes = []GroupScope{{GroupID: fmt.Sprintf("%032x", 99), + Limits: RateLimits{Requests: &zero}}} + _, err = a.Admit(ctx, zr) + if rejected, ok := err.(*Rejected); !ok || !rejected.Contractual { + t.Fatalf("zero-capped group admit = %T %v, want a contractual rejection", err, err) + } +} diff --git a/internal/admission/spend.go b/internal/admission/spend.go index 5e09033..ab9abc6 100644 --- a/internal/admission/spend.go +++ b/internal/admission/spend.go @@ -227,12 +227,19 @@ type spendFailureLog struct { mu sync.Mutex last time.Time suppressed int + // now is the clock; nil means time.Now. Tests inject a fixed clock to + // step across spendFailureLogInterval without sleeping. + now func() time.Time } func (l *spendFailureLog) logf(log *logging.Logger, format string, args ...interface{}) { l.mu.Lock() defer l.mu.Unlock() - now := time.Now() + clock := time.Now + if l.now != nil { + clock = l.now + } + now := clock() if now.Sub(l.last) < spendFailureLogInterval { l.suppressed++ return diff --git a/internal/admission/spend_integration_test.go b/internal/admission/spend_integration_test.go index 7f30b07..2ef9308 100644 --- a/internal/admission/spend_integration_test.go +++ b/internal/admission/spend_integration_test.go @@ -63,7 +63,15 @@ func spendSchemaDDL(t *testing.T) string { // - month spend exactly equal to the cap is exhausted (the >= boundary); // - month spend below the cap is not exhausted; // - a row in the PRIOR UTC month is excluded from the sum, even when it is -// the last hour before the boundary. +// the last hour before the boundary; +// - a row in the FIRST hour of the current UTC month is counted. +// +// The session runs with TimeZone set to UTC+14 (Pacific/Kiritimati), so a +// month boundary taken in the session's time zone instead of UTC lands 14 +// hours away from the UTC boundary: it either pulls the prior month's last +// hour into the sum or, during the last 14 hours of a UTC month, pushes the +// first hour of the current month out of it. Postgres in CI runs with +// TimeZone=UTC, where the two boundaries coincide and neither mistake shows. func TestIntegration_GroupSpendExhausted(t *testing.T) { dsn := os.Getenv("PHOEBE_TEST_DATABASE_URL") if dsn == "" { @@ -84,6 +92,9 @@ func TestIntegration_GroupSpendExhausted(t *testing.T) { exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") exec(t, db, "CREATE SCHEMA "+sch) exec(t, db, "SET search_path TO "+sch) + // Like search_path, TimeZone is per connection, so it holds for every + // query below on the single pinned connection. + exec(t, db, "SET TIME ZONE 'Pacific/Kiritimati'") defer func() { exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") }() exec(t, db, spendSchemaDDL(t)) @@ -99,6 +110,7 @@ func TestIntegration_GroupSpendExhausted(t *testing.T) { gEmpty = "00000000000000000000000000000000" gSpend = "a1b2c3d4e5f60718293a4b5c6d7e8f90" gOther = "00112233445566778899aabbccddeeff" + gStart = "ffeeddccbbaa99887766554433221100" ) seed := func(group string, windowStart time.Time, cost string) { t.Helper() @@ -122,6 +134,8 @@ func TestIntegration_GroupSpendExhausted(t *testing.T) { } seed(gSpend, priorMonthLastHour, "1000") seed(gOther, monthStart, "1000") + // gStart has a single row, in the first hour of the current UTC month. + seed(gStart, monthStart, "2.5") store := NewPostgresSpendStore(db) for _, tc := range []struct { @@ -143,6 +157,10 @@ func TestIntegration_GroupSpendExhausted(t *testing.T) { // 4. The prior-month row (1000) is excluded: were it counted, spend // would be 1001.500000001 and this cap would be exhausted. {"prior month excluded", gSpend, "1001", false}, + // 5. The first hour of the current UTC month is counted: spend equal + // to the cap is exhausted. A boundary taken in the session's UTC+14 + // time zone would drop this row late in the month and admit. + {"current month first hour counted", gStart, "2.5", true}, } { got, err := store.GroupSpendExhausted(ctx, tc.group, tc.cap) if err != nil { diff --git a/internal/rating/group_rollup_integration_test.go b/internal/rating/group_rollup_integration_test.go index 1450499..1d4d3c6 100644 --- a/internal/rating/group_rollup_integration_test.go +++ b/internal/rating/group_rollup_integration_test.go @@ -207,6 +207,237 @@ func TestIntegration_GroupUsageAttribution(t *testing.T) { } } +// TestIntegration_GroupUsageWithheldGatesSharedGrain pins the per-event gates +// the attributed CTE re-applies on the resolved side (store.go:819-826) against +// real Postgres. The sibling test above withholds only a usage_found event, and +// every seeded event has its own auth_id, so every event sits in its own money +// grain. That leaves the other gates unexercised: the attributed CTE joins +// priced rollups back onto resolved events on the full grain, and with no +// withheld event sharing a priced rollup's grain, dropping any one of those +// gates would change nothing — the test stays green while a withheld event +// rides its priced twin's rollup into group_usage and inflates the group spend +// the admission check reads. +// +// So every withheld event below SHARES its (auth, resource, model, +// serving_mode, hour) money grain with a priced event: +// +// - unpriced twin: model 'm-unp' prices through base_model 'b' for one event +// (C4 rung (c)); its twin names an unknown base, so its prompt_price is +// NULL and only the prompt_price gate keeps it out of the gA row; +// - owner-conflict twin: one event carries user_id AND group_id, collapses +// to the same user/group owner bucket as its clean twin, and is dropped per +// event — only the owner_conflict gate keeps its memberships (and its own +// group_id) out of group_usage; +// - ambiguous_org: one resource carrying two distinct orgs is withheld at +// the rollup level, so neither clean event may attribute to gC at all; +// - identical pair: two genuinely different but identical-looking billable +// events in one rollup must BOTH attribute to gD — the DISTINCT ON +// (request_id, gid) dedupe must not collapse them into one. +// +// For each withheld case the group's group_usage row must carry exactly its +// priced event(s), and the money rollup must be what a baseline WITHOUT the +// withheld events would write: same rollups, same events rated, same total. +func TestIntegration_GroupUsageWithheldGatesSharedGrain(t *testing.T) { + dsn := os.Getenv("PHOEBE_TEST_DATABASE_URL") + if dsn == "" { + t.Skip("PHOEBE_TEST_DATABASE_URL not set; skipping live-Postgres conformance") + } + ctx := context.Background() + db, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatalf("open: %v", err) + } + defer db.Close() + + const sch = "phoebe_rating_group_gates_it" + exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") + exec(t, db, "CREATE SCHEMA "+sch) + exec(t, db, "SET search_path TO "+sch) + defer func() { exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") }() + exec(t, db, ratingSchemaDDL(t)) + + hour := mustTime("2026-06-08T10:00:00Z") + book := newTestBook( + map[string]Rate3{"b": rate3("0.000005", "0.0000005", "0.00002")}, + nil, PolicyIdentity, Dec{}, Dec{}, + ) + rate, err := book.Resolve("b") + if err != nil { + t.Fatalf("resolve: %v", err) + } + billed := rate.Quantized() + + const ( + gA = "a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a" + gB = "b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b" + gC = "c2c2c2c2c2c2c2c2c2c2c2c2c2c2c2c2c" + gD = "d3d3d3d3d3d3d3d3d3d3d3d3d3d3d3d3d" + ) + + // Every event is dedicated with authoritative usage; each withheld event is + // kept out of money by exactly ONE classification (its price, its owner + // conflict, or its rollup's two orgs) and shares its money grain with the + // priced event that must still attribute. + type gateEvent struct { + req string + auth string + user string // billing_event.user_id ('' writes NULL) + group string // billing_event.group_id (the token's own group) + members []string // billing_event.member_group_ids + resource string + org string // billing_event.org_id ('' writes NULL) + model string + base string // billing_event.base_model ('' writes NULL) + prompt int + } + events := []gateEvent{ + // (a) the unpriced twin: 'ga-ok' prices model 'm-unp' through base 'b' + // (C4 rung (c)); 'ga-bad' names base 'zz-unpriced-base', which has no + // price row, so its prompt_price is NULL and it is withheld as unpriced. + {req: "ga-ok", auth: "a-unp", members: []string{gA}, resource: "r-unp", model: "m-unp", base: "b", prompt: 40}, + {req: "ga-bad", auth: "a-unp", members: []string{gA}, resource: "r-unp", model: "m-unp", base: "zz-unpriced-base", prompt: 999}, + // (b) the owner-conflict twin: 'gb-bad' carries user_id AND group_id, + // which collapses its owner to ''/'' — the same bucket as 'gb-ok' — and + // drops it from money per event. Its memberships and its own group must + // not reach group_usage. + {req: "gb-ok", auth: "a-conf", members: []string{gB}, resource: "r-conf", model: "b", prompt: 25}, + {req: "gb-bad", auth: "a-conf", user: "u-x", group: "g-conflicted", members: []string{gB}, resource: "r-conf", model: "b", prompt: 777}, + // (c) ambiguous_org: one resource, two distinct orgs — the rollup is + // withheld as a whole, so neither clean event may attribute to gC. + {req: "gc-1", auth: "a-org", members: []string{gC}, resource: "r-org", org: "org-x", model: "b", prompt: 50}, + {req: "gc-2", auth: "a-org", members: []string{gC}, resource: "r-org", org: "org-y", model: "b", prompt: 50}, + // (d) the identical pair: same grain, same counts, same group — both + // events must attribute to gD; dropping request_id from the DISTINCT ON + // key would collapse them into one. + {req: "gd-1", auth: "a-twin", members: []string{gD}, resource: "r-twin", model: "b", prompt: 60}, + {req: "gd-2", auth: "a-twin", members: []string{gD}, resource: "r-twin", model: "b", prompt: 60}, + } + for i, e := range events { + _, err := db.ExecContext(ctx, + `INSERT INTO billing_event (request_id, auth_id, user_id, group_id, member_group_ids, resource_id, org_id, model, base_model, serving_mode, usage_found, prompt_tokens, cached_tokens, completion_tokens, event_ts) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,'dedicated',true,$10,0,0,$11)`, + e.req, nullableStr(e.auth), nullableStr(e.user), nullableStr(e.group), groupIDsOrNil(e.members), + nullableStr(e.resource), nullableStr(e.org), nullableStr(e.model), nullableStr(e.base), + e.prompt, hour.Add(time.Duration(i)*time.Minute)) + if err != nil { + t.Fatalf("seed event %d (%s): %v", i, e.req, err) + } + } + + // Oracle for the money baseline WITHOUT the withheld events: 'ga-ok' prices + // at the base 'b' rate through C4 rung (c), the others price model 'b' + // directly — all four resolve to the same quantized rate. + moneyEvents := []RatedEvent{ + {AuthID: "a-unp", ResourceID: "r-unp", ModelID: "m-unp", BaseModel: "b", ServingMode: "dedicated", PromptTokens: 40}, + {AuthID: "a-conf", ResourceID: "r-conf", ModelID: "b", ServingMode: "dedicated", PromptTokens: 25}, + {AuthID: "a-twin", ResourceID: "r-twin", ModelID: "b", ServingMode: "dedicated", PromptTokens: 60}, + {AuthID: "a-twin", ResourceID: "r-twin", ModelID: "b", ServingMode: "dedicated", PromptTokens: 60}, + } + costs := make([]Dec, len(moneyEvents)) + for i, e := range moneyEvents { + costs[i] = Rate(e, billed) + } + + store := NewPostgresStore(db) + res, err := store.RateWindow(ctx, book, hour, hour.Add(time.Hour)) + if err != nil { + t.Fatalf("RateWindow: %v", err) + } + + // Money is exactly the priced baseline: the three priced rollups (one per + // grain), four rated events, and a total equal to the sum of the four + // oracle costs. The withheld events changed nothing. + if res.RollupsWritten != 3 || res.EventsRated != 4 { + t.Fatalf("rollups/events = %d/%d, want 3/4 (the priced baseline only)", res.RollupsWritten, res.EventsRated) + } + wantTotal := costs[0].Add(costs[1]).Add(costs[2]).Add(costs[3]).String() + if MustDec(res.TotalCost).String() != wantTotal { + t.Fatalf("total cost = %s, want %s (the priced baseline only)", res.TotalCost, wantTotal) + } + // The withheld events are classified, not silently dropped: the anomaly + // partition accounts for all 8 seeded events exactly once. + if res.UnpricedEvents != 1 || res.OwnerConflictEvents != 1 || res.AmbiguousOrgEvents != 2 { + t.Fatalf("unpriced/conflict/ambig-org = %d/%d/%d, want 1/1/2", + res.UnpricedEvents, res.OwnerConflictEvents, res.AmbiguousOrgEvents) + } + if got := res.EventsRated + res.MissingUsageEvents + res.InvalidUsageEvents + + res.UnpricedEvents + res.UnattributableEvents + res.InvalidServingModeEvents + + res.AmbiguousBaseEvents + res.AmbiguousOrgEvents + res.OwnerConflictEvents; got != 8 { + t.Fatalf("anomaly partition = %d, want 8 (all seeded events accounted exactly once)", got) + } + if res.GroupRollupsWritten != 3 { + t.Fatalf("group rollups = %d, want 3 (groups %s, %s, %s)", + res.GroupRollupsWritten, gA[:8], gB[:8], gD[:8]) + } + + type groupRow struct { + prompt, billable, count int64 + cost string + } + readRow := func(gid string) (groupRow, bool) { + var r groupRow + err := db.QueryRowContext(ctx, + `SELECT prompt_tokens, billable_prompt_tokens, event_count, cost::text + FROM group_usage WHERE group_id=$1 AND window_start=$2`, gid, hour). + Scan(&r.prompt, &r.billable, &r.count, &r.cost) + if errors.Is(err, sql.ErrNoRows) { + return groupRow{}, false + } + if err != nil { + t.Fatalf("read group_usage %s: %v", gid, err) + } + return r, true + } + + // (a) gA carries 'ga-ok' alone: dropping the prompt_price gate would let + // 'ga-bad' ride its twin's rollup in (count 2, prompt 1039). + rA, ok := readRow(gA) + if !ok { + t.Fatalf("no group_usage row for %s", gA) + } + if rA.count != 1 || rA.prompt != 40 || rA.billable != 40 || MustDec(rA.cost).String() != costs[0].String() { + t.Fatalf("group %s row = %+v, want count 1 / prompt 40 / billable 40 / cost %s (the priced twin only)", + gA, rA, costs[0]) + } + + // (b) gB carries 'gb-ok' alone: dropping the owner_conflict gate would add + // 'gb-bad' (count 2, prompt 802) and create a row for its own group. + rB, ok := readRow(gB) + if !ok { + t.Fatalf("no group_usage row for %s", gB) + } + if rB.count != 1 || rB.prompt != 25 || MustDec(rB.cost).String() != costs[1].String() { + t.Fatalf("group %s row = %+v, want count 1 / prompt 25 / cost %s (the clean twin only)", gB, rB, costs[1]) + } + if _, ok := readRow("g-conflicted"); ok { + t.Fatalf("the conflicted event's own group reached group_usage; the owner_conflict twin must attribute nowhere") + } + + // (c) the ambiguous-org rollup has no money row and no attribution row. + if _, ok := readRow(gC); ok { + t.Fatalf("group %s row exists, want none (its rollup is ambiguous-org withheld)", gC) + } + var nOrgRollups int + if err := db.QueryRowContext(ctx, + `SELECT COUNT(*) FROM rated_usage WHERE resource_id='r-org'`).Scan(&nOrgRollups); err != nil { + t.Fatalf("count r-org rollups: %v", err) + } + if nOrgRollups != 0 { + t.Fatalf("rated_usage has %d rows for r-org, want 0 (ambiguous-org withheld)", nOrgRollups) + } + + // (d) gD carries BOTH identical events: per-event dedupe would show count 1. + rD, ok := readRow(gD) + if !ok { + t.Fatalf("no group_usage row for %s", gD) + } + twinCost := costs[2].Add(costs[3]).String() + if rD.count != 2 || rD.prompt != 120 || MustDec(rD.cost).String() != twinCost { + t.Fatalf("group %s row = %+v, want count 2 / prompt 120 / cost %s (both events attribute)", + gD, rD, twinCost) + } +} + // groupIDsOrNil binds a membership list as a driver value (nil for none — // mirrors the drainer's nullGroupIDs so the fixture writes what production // writes). diff --git a/internal/rating/rollback_integration_test.go b/internal/rating/rollback_integration_test.go new file mode 100644 index 0000000..3fb41a2 --- /dev/null +++ b/internal/rating/rollback_integration_test.go @@ -0,0 +1,136 @@ +//go:build integration + +package rating + +import ( + "context" + "database/sql" + "os" + "testing" +) + +// TestIntegration_GroupScopesRollbackIsLossy pins the LOSSY claim in the +// header of migrations/0008_group_scopes.down.sql against real Postgres, so +// the comment can never be "fixed" back to the false claim that re-rating +// after a down/up cycle rebuilds the group rows: +// +// - the schema comes back: after 0008 down then up again, +// billing_event.member_group_ids exists again and group_usage exists +// again; +// - the evidence does NOT come back: a row seeded with membership evidence +// reads member_group_ids = NULL after the cycle -- the down migration +// dropped the only copy of the membership evidence, and the up migration +// only re-adds an empty column. Re-rating afterwards attributes group +// usage from group-token usage (billing_event.group_id) only, so +// month-to-date group spend restarts lower than it really is. +// +// It loads the REAL migration files 0001-0008 from disk (like +// TestIntegration_InvalidUsageEvidenceNeverEntersMoney's apply loop) rather +// than a hand-copied DDL constant, so it tests exactly what cmd/migrate ships. +func TestIntegration_GroupScopesRollbackIsLossy(t *testing.T) { + dsn := os.Getenv("PHOEBE_TEST_DATABASE_URL") + if dsn == "" { + t.Skip("PHOEBE_TEST_DATABASE_URL not set; skipping live-Postgres conformance") + } + ctx := context.Background() + db, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatalf("open: %v", err) + } + defer db.Close() + + const sch = "phoebe_rating_rollback_it" + exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") + exec(t, db, "CREATE SCHEMA "+sch) + exec(t, db, "SET search_path TO "+sch) + defer func() { exec(t, db, "DROP SCHEMA IF EXISTS "+sch+" CASCADE") }() + + apply := func(name string) { + t.Helper() + ddl, readErr := os.ReadFile("../../migrations/" + name) + if readErr != nil { + t.Fatalf("read migration %s: %v (the integration test applies the REAL "+ + "migration DDL so it can't drift from production)", name, readErr) + } + exec(t, db, string(ddl)) + } + // The full production chain 0001-0008, in cmd/migrate apply order. + for _, name := range []string{ + "0001_billing_event.up.sql", + "0002_rating.up.sql", + "0003_io_log.up.sql", + "0004_billing_event_serving_mode.up.sql", + "0005_invoice_grade_attempts.up.sql", + "0006_rollup_grain.up.sql", + "0007_serving_mode_explicit.up.sql", + "0008_group_scopes.up.sql", + } { + apply(name) + } + + // Seed one row carrying membership evidence -- the column 0008 added. + if _, err := db.ExecContext(ctx, + `INSERT INTO billing_event (request_id, auth_id, member_group_ids) + VALUES ('member-evidence','a', ARRAY['g1','g2']::text[])`); err != nil { + t.Fatalf("seed membership evidence: %v", err) + } + var seeded int + if err := db.QueryRowContext(ctx, + `SELECT COUNT(*) FROM billing_event + WHERE request_id = 'member-evidence' AND member_group_ids = ARRAY['g1','g2']::text[]`). + Scan(&seeded); err != nil { + t.Fatalf("read seeded evidence: %v", err) + } + if seeded != 1 { + t.Fatalf("seeded membership evidence rows = %d, want 1 (the fixture must carry evidence before the rollback)", seeded) + } + + // The rollback cycle: 0008 down, then 0008 up again. + apply("0008_group_scopes.down.sql") + apply("0008_group_scopes.up.sql") + + // The schema is restored: the column and the table both exist again. + var colExists bool + if err := db.QueryRowContext(ctx, + `SELECT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = current_schema() + AND table_name = 'billing_event' + AND column_name = 'member_group_ids')`).Scan(&colExists); err != nil { + t.Fatalf("check member_group_ids column: %v", err) + } + if !colExists { + t.Fatal("billing_event.member_group_ids is missing after the down/up cycle -- the schema must be restored") + } + var tblExists bool + if err := db.QueryRowContext(ctx, + `SELECT EXISTS ( + SELECT 1 FROM information_schema.tables + WHERE table_schema = current_schema() + AND table_name = 'group_usage')`).Scan(&tblExists); err != nil { + t.Fatalf("check group_usage table: %v", err) + } + if !tblExists { + t.Fatal("group_usage is missing after the down/up cycle -- the schema must be restored") + } + + // The evidence is NOT restored: the seeded row survives (the raw ledger is + // never touched by 0008), but its membership evidence reads back NULL. + // This is the lossy part: if the down migration ever "preserved" the + // column (or the header reverted to claiming it did), this assertion goes + // red. + var rowCount, evidenceRows int + if err := db.QueryRowContext(ctx, + `SELECT COUNT(*), COUNT(*) FILTER (WHERE member_group_ids IS NOT NULL) + FROM billing_event WHERE request_id = 'member-evidence'`). + Scan(&rowCount, &evidenceRows); err != nil { + t.Fatalf("read evidence after rollback: %v", err) + } + if rowCount != 1 { + t.Fatalf("billing_event rows for the seeded event = %d, want 1 (the raw row must survive the cycle)", rowCount) + } + if evidenceRows != 0 { + t.Fatalf("member_group_ids still holds evidence on %d row(s) after the down/up cycle -- "+ + "the rollback is LOSSY and must read NULL (see 0008_group_scopes.down.sql's header)", evidenceRows) + } +} diff --git a/migrations/0008_group_scopes.down.sql b/migrations/0008_group_scopes.down.sql index a034c65..06118b1 100644 --- a/migrations/0008_group_scopes.down.sql +++ b/migrations/0008_group_scopes.down.sql @@ -1,7 +1,15 @@ -- Roll back 0008. group_usage is fully owned by this migration (drop the -- table); member_group_ids is evidence only (drop the column). Rated windows -- keep their rated_usage money rows either way -- this rollback never touches --- money. Re-rating a window after a down/up cycle rebuilds the group rows. +-- money. +-- +-- THIS ROLLBACK IS LOSSY. member_group_ids is the only copy of the membership +-- evidence: dropping the column destroys it, and re-applying the up migration +-- only re-adds an empty (NULL) column. After a down/up cycle, re-rating a +-- window rebuilds group_usage from group-token usage (billing_event.group_id) +-- only -- membership attribution is permanently lost, so month-to-date group +-- spend restarts lower than it really is and spend caps can over-admit until +-- real usage re-accumulates. Operators should avoid a mid-month rollback. DROP TABLE IF EXISTS group_usage; diff --git a/migrations/0008_group_scopes.up.sql b/migrations/0008_group_scopes.up.sql index 5184e1b..561fc26 100644 --- a/migrations/0008_group_scopes.up.sql +++ b/migrations/0008_group_scopes.up.sql @@ -17,10 +17,15 @@ -- under every group it belongs to, and withheld (non-money) events -- contribute nothing. -- --- ROLLOUT: the rater's statement references member_group_ids and group_usage, --- so a pre-0008 rater cannot run against this schema and a post-0008 rater --- cannot run against the old one (SQLSTATE 42703) — roll code and schema --- together, in the order used for migration 0007 (migrate up, then deploy). +-- ROLLOUT: the failure is asymmetric, so the order is migration first, then +-- code. A post-0008 rater cannot run against the pre-0008 schema (SQLSTATE +-- 42703: its statements reference member_group_ids and group_usage, which do +-- not exist there yet). A pre-0008 rater, by contrast, runs cleanly against +-- the 0008 schema — 0008 only adds a nullable column and a new table, so the +-- old statements still execute. The old rater writes no group_usage rows, so +-- group spend caps read an empty rollup (zero spend) until the new rater +-- deploys. Safe order: apply migration 0008 first, then deploy the new rater, +-- which then starts writing group_usage. -- --------------------------------------------------------------------------- -- billing_event: membership evidence diff --git a/migrations/README.md b/migrations/README.md index 8eb5e93..2b3ba84 100644 --- a/migrations/README.md +++ b/migrations/README.md @@ -68,14 +68,22 @@ migrate version # print the current applied version ### Rollout order for migration 0008 Migration 0008 adds `billing_event.member_group_ids` and the `group_usage` -table, both of which the rater's single rating statement reads/writes. A -pre-0008 rater fails against this schema and a post-0008 rater fails against -the old one (SQLSTATE 42703), so roll code and schema together, in the -migration-0007 order: run `cmd/migrate up`, then deploy the new drainer and -rater (the drainer's INSERT gains the column; an old drainer against schema -0008 simply writes NULL membership). The interceptor's group scope enforcement -is envelope-driven, so it is safe across the cutover: no stamped envelope, no -group check. +table, both of which the rater's single rating statement reads/writes. The +failure modes are asymmetric, so the order matters: + +- A pre-0008 rater still runs cleanly against the 0008 schema: 0008 only adds + a nullable column and a new table, neither of which the old rating + statement references. The old rater rates and writes money as before, but + writes no `group_usage`, so the admission group spend check reads an empty + rollup (zero group spend) until the new rater covers those hours. +- A post-0008 rater fails against the old schema (SQLSTATE 42703): its rating + statement references `member_group_ids` and `group_usage`. + +So roll code and schema together, in the migration-0007 order: run +`cmd/migrate up`, then deploy the new drainer and rater (the drainer's INSERT +gains the column; an old drainer against schema 0008 simply writes NULL +membership). The interceptor's group scope enforcement is envelope-driven, so +it is safe across the cutover: no stamped envelope, no group check. `group_usage` rows for already-rated hours exist only after a re-rate (the rollup is written by the rater, not backfilled). Until a window is re-rated, From b3a7924580a3bd217a58a92e35c2d99b7b8d0b1f Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 10:42:55 +0000 Subject: [PATCH 12/17] admission: a starved spend read is not a store outage; tests: e2e group spend path, usage-gate twins, deterministic coalescing A spend read that begins with almost none of the shared 250ms budget left and then times out was starved by earlier groups in the same request, not evidence of an outage: fail it open without the ERROR line and without a negative-cache entry so the group's next request re-queries the store. Add an end-to-end test that carries a group envelope from the proxy through drain and rating to the spend verdict, cover the usage_found and valid_usage attribution gates with unbillable twins, use grammar-valid group ids in the withheld-gates fixture, make the hung-store assertion bound against the spend budget, replace the sleep-based coalescing rendezvous with a deterministic signal, document the post-0008 drainer poison-drop against the pre-0008 schema, and correct the window_start column comment to the UTC boundary the code uses. --- internal/admission/admission.go | 14 +- internal/admission/group_test.go | 93 ++++++++++- internal/admission/spend.go | 15 ++ internal/e2e/e2e_test.go | 154 ++++++++++++++++++ internal/proxy/group_scopes_test.go | 90 +++++++++- .../rating/group_rollup_integration_test.go | 104 ++++++++---- migrations/0008_group_scopes.up.sql | 3 +- migrations/README.md | 6 + 8 files changed, 436 insertions(+), 43 deletions(-) diff --git a/internal/admission/admission.go b/internal/admission/admission.go index 8cdf10a..18670b6 100644 --- a/internal/admission/admission.go +++ b/internal/admission/admission.go @@ -509,6 +509,9 @@ func (a *RedisAdmitter) checkGroupSpend(ctx context.Context, scopes []GroupScope defer cancelBudget() } deadline, _ := budgetCtx.Deadline() + // startBudget is what was left of the shared budget when this read + // began. Only the leader runs the read, so only the leader sets it. + var startBudget time.Duration // Collapse concurrent misses on one (group, cap) into a single store // read: waiters take the leader's result (verdict, or fail-open on // error) instead of each running their own query against a slow store. @@ -519,6 +522,7 @@ func (a *RedisAdmitter) checkGroupSpend(ctx context.Context, scopes []GroupScope // this read, and one client disconnecting must not end it and // hand every waiter a free admit. The shared deadline still // bounds it. + startBudget = time.Until(deadline) queryCtx, cancel := context.WithDeadline(context.WithoutCancel(ctx), deadline) defer cancel() return a.spendStore.GroupSpendExhausted(queryCtx, gs.GroupID, gs.SpendCap) @@ -530,7 +534,15 @@ func (a *RedisAdmitter) checkGroupSpend(ctx context.Context, scopes []GroupScope // guards the same race at renew time). The 250ms deadline // expiring with a live client IS logged: the store is too slow. // Only the leader records and logs — waiters share its outcome. - if leader && ctx.Err() == nil { + // + // A read that began with almost none of the shared budget left + // and then ran out of time was starved by the reads of earlier + // groups in this request; that is not evidence of an outage. It + // fails open without the ERROR line and without a negative cache + // entry, so the group's next request queries the store again + // instead of skipping its cap for groupSpendNegativeCacheTTL. + starved := errors.Is(result.err, context.DeadlineExceeded) && startBudget < groupSpendStarvedFloor + if leader && ctx.Err() == nil && !starved { a.spendCache.storeFailure(gs.GroupID, gs.SpendCap) a.spendBypass(gs.GroupID, result.err) } diff --git a/internal/admission/group_test.go b/internal/admission/group_test.go index d040da7..68d4274 100644 --- a/internal/admission/group_test.go +++ b/internal/admission/group_test.go @@ -351,8 +351,10 @@ func TestGroupSpendHungStoreFailsOpenWithinBudget(t *testing.T) { t.Fatalf("hung-store admit: %v, want fail-open admission", err) } _ = lease.Complete(context.Background(), 0) - if elapsed > admitOperationBudget { - t.Fatalf("hung-store admit took %v, want it bounded by the spend query budget (%v)", elapsed, groupSpendQueryBudget) + // Bounded by the spend budget, not the wider admit budget — and it did + // wait for that budget, so an instant answer cannot pass this test. + if elapsed >= 2*groupSpendQueryBudget || elapsed < groupSpendQueryBudget/2 { + t.Fatalf("hung-store admit took %v, want about one spend query budget (%v)", elapsed, groupSpendQueryBudget) } out := buf.String() if !strings.Contains(out, "monthly spend cap NOT enforced (fail open)") || !strings.Contains(out, context.DeadlineExceeded.Error()) { @@ -469,6 +471,9 @@ func TestGroupSpendLeaderCancelDoesNotCancelCoalescedQuery(t *testing.T) { store := &releasedSpendStore{entered: make(chan struct{}), release: make(chan struct{})} a.WithGroupSpend(store, logging.New(logging.ERROR)) + joined := make(chan struct{}, 1) + a.spendFlight.joined = func(string) { joined <- struct{}{} } + leaderCtx, cancelLeader := context.WithCancel(context.Background()) leaderDone := make(chan struct{}) go func() { @@ -490,8 +495,8 @@ func TestGroupSpendLeaderCancelDoesNotCancelCoalescedQuery(t *testing.T) { } waiterErr <- err }() - // Let the waiter join the in-flight read, then let the read answer. - time.Sleep(20 * time.Millisecond) + // Let the read answer only once the waiter has joined it. + <-joined close(store.release) err := <-waiterErr @@ -508,6 +513,86 @@ func TestGroupSpendLeaderCancelDoesNotCancelCoalescedQuery(t *testing.T) { } } +// slowSpendStore models a healthy but slow Postgres: it answers below-cap +// after its latency, and counts reads per group. So that the test does not +// depend on scheduler timing, a read that has enough time left answers with +// room to spare (it never takes more than half the time remaining), and a +// read that starts under groupSpendStarvedFloor runs out its deadline — the +// starved case the check must not report as an outage. +type slowSpendStore struct { + latency time.Duration + mu sync.Mutex + calls map[string]int +} + +func (s *slowSpendStore) GroupSpendExhausted(ctx context.Context, groupID, _ string) (bool, error) { + s.mu.Lock() + s.calls[groupID]++ + s.mu.Unlock() + wait := s.latency + if deadline, ok := ctx.Deadline(); ok { + left := time.Until(deadline) + if left < groupSpendStarvedFloor { + <-ctx.Done() + return false, ctx.Err() + } + wait = min(wait, left/2) + } + select { + case <-time.After(wait): + return false, nil + case <-ctx.Done(): + return false, ctx.Err() + } +} + +func (s *slowSpendStore) Calls(groupID string) int { + s.mu.Lock() + defer s.mu.Unlock() + return s.calls[groupID] +} + +// TestGroupSpendBudgetExhaustionNotLoggedAsOutage: six uncached capped groups +// against a store that answers in ~60ms cannot all fit in one shared 250ms +// budget. The later groups run out of budget because of their siblings, not +// because the store is down: the request admits (fail open) without the +// outage ERROR line, and no negative cache entry is recorded, so the next +// request for a starved group queries the store again. +func TestGroupSpendBudgetExhaustionNotLoggedAsOutage(t *testing.T) { + a, _ := testAdmitter(t, config.AdmissionSettings{}) + log := logging.New(logging.ERROR) + var buf bytes.Buffer + log.Error.SetOutput(&buf) + store := &slowSpendStore{latency: 60 * time.Millisecond, calls: map[string]int{}} + a.WithGroupSpend(store, log) + + scopes := make([]GroupScope, 6) + for i := range scopes { + scopes[i] = groupScope(fmt.Sprintf("%032x", i+1), nil, "100") + } + lease, err := a.Admit(context.Background(), groupRequest("org-be", "m", scopes...)) + if err != nil { + t.Fatalf("6-group slow-store admit: %v, want fail-open admission", err) + } + _ = lease.Complete(context.Background(), 0) + if out := buf.String(); strings.Contains(out, "monthly spend cap NOT enforced") { + t.Fatalf("budget-starved groups logged as an outage: %q, want no bypass ERROR", out) + } + + // The last group can never fit (6 x 60ms > 250ms). With a fresh budget + // its own request reads the store instead of hitting a negative entry. + last := scopes[len(scopes)-1] + before := store.Calls(last.GroupID) + lease, err = a.Admit(context.Background(), groupRequest("org-be", "m", last)) + if err != nil { + t.Fatalf("starved-group admit: %v, want admission (below cap)", err) + } + _ = lease.Complete(context.Background(), 0) + if got := store.Calls(last.GroupID); got != before+1 { + t.Fatalf("starved group store calls = %d, want %d (no negative cache entry after starvation)", got, before+1) + } +} + // TestGroupSpendParentCancelNotLoggedAsBypass: a canceled request context // (client gone) makes the spend check moot — it must NOT log the throttled // 'monthly spend cap NOT enforced' ERROR, which a genuine outage needs. diff --git a/internal/admission/spend.go b/internal/admission/spend.go index ab9abc6..7ff239f 100644 --- a/internal/admission/spend.go +++ b/internal/admission/spend.go @@ -35,6 +35,13 @@ const groupSpendNegativeCacheTTL = 10 * time.Second // inside admitOperationBudget so the reservation keeps its own time. const groupSpendQueryBudget = 250 * time.Millisecond +// groupSpendStarvedFloor is the remaining shared budget below which a read +// that times out is treated as starved by the earlier reads in the same +// request rather than as a slow store. A read that starts with this little +// time cannot finish against a healthy Postgres either, so its timeout says +// nothing about the store and is neither logged nor negative-cached. +const groupSpendStarvedFloor = 25 * time.Millisecond + // groupSpendConnectTimeout bounds establishing one pool connection, so a // Postgres that drops packets cannot pin a dial (and with it one of the four // pool slots) forever. @@ -263,6 +270,10 @@ func (l *spendFailureLog) logf(log *logging.Logger, format string, args ...inter type spendFlight struct { mu sync.Mutex in map[string]*spendFlightCall + // joined, when set, is called each time a caller joins an in-flight read + // as a waiter. It exists so a test can release the read only after a + // waiter has joined it; nil in production. + joined func(key string) } type spendFlightCall struct { @@ -289,7 +300,11 @@ func newSpendFlight() *spendFlight { func (f *spendFlight) do(ctx context.Context, key string, fn func() (bool, error)) (spendFlightResult, bool) { f.mu.Lock() if call, ok := f.in[key]; ok { + joined := f.joined f.mu.Unlock() + if joined != nil { + joined(key) + } select { case <-call.done: return call.result, false diff --git a/internal/e2e/e2e_test.go b/internal/e2e/e2e_test.go index 5ae0c5a..37c6ca7 100644 --- a/internal/e2e/e2e_test.go +++ b/internal/e2e/e2e_test.go @@ -41,6 +41,7 @@ import ( _ "github.com/jackc/pgx/v5/stdlib" "github.com/redis/go-redis/v9" + "github.com/saturncloud/phoebe/internal/admission" "github.com/saturncloud/phoebe/internal/config" "github.com/saturncloud/phoebe/internal/drain" "github.com/saturncloud/phoebe/internal/emit" @@ -1065,3 +1066,156 @@ func TestE2E_AdapterHeaderLandsInBillingEventAndTriggersPremium(t *testing.T) { h.assertNumericEqual(t, appliedPrompt, "0.000006", "rated_usage.applied_prompt_rate (base x premium frozen on row)") h.assertNumericEqual(t, cost, wantCost, "rated_usage.cost") } + +// testGroupID is the 32-hex group id the group-scopes pipeline test stamps on +// its X-Saturn-Group-Scopes envelope (uuid4 hex, as Atlas mints). +const testGroupID = "a1b2c3d4e5f60718293a4b5c6d7e8f90" + +// TestE2E_GroupScopedRequestBecomesGroupSpend is the group attribution +// pipeline test. Every hop of the membership-aware group quota path (proxy +// envelope parse and emit, drainer binding of member_group_ids, rater +// attribution into group_usage, admission spend verdict) has unit tests over +// hand-seeded rows, so a shape mismatch BETWEEN hops — for example the drainer +// binding the group list in a form the rater's unnest does not read as text[] — +// would leave group spend silently at zero with every unit test green. This +// test runs one shared request carrying a real v1 envelope through the real +// components and asserts the far end: +// +// - membership contract: billing_event.member_group_ids holds the envelope's +// group id, captured at meter time; +// - attribution contract: group_usage has a row for that group whose cost +// equals the event's money cost (the same oracle as the money test); +// - spend verdict: admission.PostgresSpendStore reads that row back, so a +// cap equal to the cost is exhausted and a higher cap is not. +// +// The request passes through a real admitter (admission.enabled=false, as a +// default chart install renders it) over the harness's miniredis, with the +// real Postgres spend store wired in: the envelope's requests-per-minute rate +// and its spend cap engage Admit, and the spend check reads the still-empty +// group_usage on the request path before rating fills it. +func TestE2E_GroupScopedRequestBecomesGroupSpend(t *testing.T) { + h := newHarness(t, "phoebe_e2e_group_scopes") + + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fl, _ := w.(http.Flusher) + for _, chunk := range strings.SplitAfter(vllmStream, "\n\n") { + if chunk == "" { + continue + } + _, _ = io.WriteString(w, chunk) + if fl != nil { + fl.Flush() + } + } + })) + defer backend.Close() + + // A real admitter over the harness's miniredis (its own key prefix, so it + // never touches the metering stream), with the real Postgres spend store. + admissionCfg := config.AdmissionSettings{ + Enabled: false, + KeyPrefix: "phoebe-admission-e2e", + LeaseTTL: time.Minute, + DefaultMaxOutputTokens: config.DefaultMaxOutputTokens, + } + spendStore := admission.NewPostgresSpendStore(h.db) + admitter := admission.New(h.rdb, admissionCfg).WithGroupSpend(spendStore, h.log) + srv := proxy.New(&config.Settings{ListenAddr: ":0", Admission: admissionCfg}, h.log, h.emitter). + WithAdmitter(admitter) + + // A shared route as Atlas stamps it, plus a one-group envelope: 60 + // requests per minute, the token rates unlimited, and a spend cap well + // above this request's cost so the request is admitted. + rr := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", + strings.NewReader(`{"model":"`+testModelName+`","stream":true,"max_tokens":300,"messages":[]}`)) + req.Header.Set(identity.HeaderUpstream, backend.URL) + req.Header.Set(identity.HeaderServingMode, identity.ServingModeShared) + req.Header.Set(identity.HeaderServedModel, testModelName) + req.Header.Set(identity.HeaderAuthID, testAuthID) + req.Header.Set(identity.HeaderResourceID, testResourceID) + req.Header.Set(identity.HeaderResourceType, "deployment") + req.Header.Set(identity.HeaderUserID, "user-e2e") + req.Header.Set(identity.HeaderOwnerID, "user-e2e") + req.Header.Set(identity.HeaderOrgID, testOrgID) + req.Header.Set(identity.HeaderGroupScopes, "v1;"+testGroupID+":60,,,,100") + srv.Handler().ServeHTTP(rr, req) + + if rr.Code != http.StatusOK { + t.Fatalf("proxy status = %d, want 200 (body: %s)", rr.Code, rr.Body.String()) + } + + h.waitForStreamLen(t, 1, 5*time.Second) + h.drainUntilRows(t, 1, 10*time.Second) + + // MEMBERSHIP CONTRACT at the drain layer: the group list is stored as a + // Postgres text[] holding exactly the envelope's group id. + var members sql.NullString + if err := h.db.QueryRow( + `SELECT array_to_string(member_group_ids, ',') FROM billing_event`).Scan(&members); err != nil { + t.Fatalf("read billing_event.member_group_ids: %v", err) + } + if !members.Valid || members.String != testGroupID { + t.Fatalf("billing_event.member_group_ids = %v, want {%s} (from X-Saturn-Group-Scopes)", members, testGroupID) + } + + res := h.rateEventHour(t, h.priceBook(t)) + + // Same fixture and price file as the money test, so the same hand-derived + // cost: 86*0.000005 + 1920*0.0000005 + 300*0.00002 = 0.00739. + const wantCost = "0.00739" + if res.EventsRated != 1 || res.UnpricedEvents != 0 || res.UnattributableEvents != 0 { + t.Fatalf("rater Result = %+v, want 1 event rated with no anomalies", res) + } + h.assertNumericEqual(t, res.TotalCost, wantCost, "Result.TotalCost") + + // ATTRIBUTION CONTRACT: exactly one group_usage row, for the envelope's + // group, carrying the event's cost and tokens. + var nGroupRows int + if err := h.db.QueryRow("SELECT COUNT(*) FROM group_usage").Scan(&nGroupRows); err != nil { + t.Fatalf("count group_usage: %v", err) + } + if nGroupRows != 1 { + t.Fatalf("group_usage rows = %d, want exactly 1 (the envelope's group); zero means group spend is silently unattributed", nGroupRows) + } + var ( + guGroupID, guCost string + guPrompt, guCached, guCompletion int64 + guBillable, guEventCount int64 + ) + if err := h.db.QueryRow( + `SELECT group_id, cost::text, prompt_tokens, cached_tokens, completion_tokens, billable_prompt_tokens, event_count + FROM group_usage`). + Scan(&guGroupID, &guCost, &guPrompt, &guCached, &guCompletion, &guBillable, &guEventCount); err != nil { + t.Fatalf("read group_usage: %v", err) + } + if guGroupID != testGroupID { + t.Errorf("group_usage.group_id = %q, want %q", guGroupID, testGroupID) + } + if guPrompt != 2006 || guCached != 1920 || guCompletion != 300 || guBillable != 86 || guEventCount != 1 { + t.Errorf("group_usage tokens = %d/%d/%d billable=%d events=%d, want 2006/1920/300 billable=86 events=1", + guPrompt, guCached, guCompletion, guBillable, guEventCount) + } + h.assertNumericEqual(t, guCost, wantCost, "group_usage.cost") + + // SPEND VERDICT: the admission spend store reads the rated row back. The + // cap is inclusive (spend >= cap is exhausted), so a cap equal to the cost + // is exhausted and any higher cap is not. + ctx := context.Background() + exhausted, err := spendStore.GroupSpendExhausted(ctx, testGroupID, wantCost) + if err != nil { + t.Fatalf("GroupSpendExhausted(cap=%s): %v", wantCost, err) + } + if !exhausted { + t.Errorf("GroupSpendExhausted(cap=%s) = false, want true — the spend check does not see the rated group spend", wantCost) + } + exhausted, err = spendStore.GroupSpendExhausted(ctx, testGroupID, "0.007390001") + if err != nil { + t.Fatalf("GroupSpendExhausted(cap=0.007390001): %v", err) + } + if exhausted { + t.Error("GroupSpendExhausted(cap=0.007390001) = true, want false — a cap above the spend must not be exhausted") + } +} diff --git a/internal/proxy/group_scopes_test.go b/internal/proxy/group_scopes_test.go index bd67290..54e63bd 100644 --- a/internal/proxy/group_scopes_test.go +++ b/internal/proxy/group_scopes_test.go @@ -194,20 +194,37 @@ func TestParseTrustedGroupScopes(t *testing.T) { } }) - t.Run("over the size bound", func(t *testing.T) { - // 16 maximal entries (32-hex gid + ":,,,,") plus the v1; prefix stays - // under 4 KiB, so pad one entry's (unused) spend field to cross it. + t.Run("maximal valid envelope stays under the size bound", func(t *testing.T) { + // K4B: a >4 KiB fixture built only from grammatically valid entries + // does not exist. The 16-entry cap bounds the count and every field + // is length-bounded by its own grammar, so the largest envelope the + // grammar admits is the one below — 16 entries, every rate at the 18 + // nines that stay inside int64, every spend cap at spendCapRe's + // 11+9 digits, about 2 KiB, half the 4 KiB bound. Any header long + // enough to reach 4 KiB is malformed in at least one other way (an + // over-long spend field fails spendCapRe first), so the parser's + // size check is defense-in-depth ordering, not a bound a valid + // producer can ever hit — no valid-but-oversize fixture can pin it. + // The transport-level ceiling on any single header is the HTTP + // server's MaxHeaderBytes-style limit (net/http's default header + // block, which server.go does not override). The bounds the grammar + // actually enforces are pinned here and in the siblings above: this + // maximal envelope parses, and the shapes that could approach 4 KiB + // fail on the 16-entry cap (the "over the entry bound" case) or the + // per-entry grammar (the malformed table) — remove either check and + // this test's siblings fail. + maxRate := strings.Repeat("9", 18) // largest rate that still fits int64 + maxSpend := "99999999999.999999999" // largest spend spendCapRe admits entries := make([]string, 0, maxGroupScopeEntries) for i := 0; i < maxGroupScopeEntries; i++ { - entries = append(entries, groupIDFor(i)+":,,,,") + entries = append(entries, groupIDFor(i)+":"+maxRate+","+maxRate+","+maxRate+","+maxRate+","+maxSpend) } in := "v1;" + strings.Join(entries, ";") - in += strings.Repeat("0", maxGroupScopeBytes-len(in)) - if len(in) != maxGroupScopeBytes { - t.Fatalf("fixture size = %d, want exactly the %d-byte bound", len(in), maxGroupScopeBytes) + if len(in) >= maxGroupScopeBytes { + t.Fatalf("maximal valid envelope = %d bytes, want under the %d-byte bound — the fixture must stay the largest the grammar admits", len(in), maxGroupScopeBytes) } - if _, err := parseTrustedGroupScopes(identity.Identity{GroupScopes: in}); err == nil { - t.Fatalf("parse of a %d-byte header succeeded, want the size bound enforced", len(in)) + if _, err := parseTrustedGroupScopes(identity.Identity{GroupScopes: in}); err != nil { + t.Fatalf("parse of the maximal valid envelope (%d bytes) failed: %v", len(in), err) } }) } @@ -530,3 +547,58 @@ func TestSharedGroupScopesEnforcedWithAdmissionFlagOff(t *testing.T) { t.Fatalf("MemberGroupIDs = %v, want [%s]", got, testGroupA) } } + +// TestSharedGroupScopesOnlyEnvelopeFailsClosedWithAdmissionFlagOff closes the +// trigger hole in the enabled=false test above: that one builds on +// sharedRequest, which always stamps X-Saturn-Owner-Id, so its owner-id +// anchor alone makes trustedPolicyEnvelopePresent true — a regression +// dropping id.GroupScopes != "" from that function would leave every case +// above green. This test sends the group envelope with NO other trusted +// policy header: the owner-id anchor and every scoped rate-limit header are +// deleted, so only the group clause of trustedPolicyEnvelopePresent can +// engage admission. +// +// What the proxy actually answers for that shape, and why it is not 429 or +// 200: parseTrustedRateLimits runs first inside the admission block and +// fails closed on the missing contract anchor (errNoTrustedRateLimitPolicy), +// so the request is refused 503 before the group scopes are even parsed — a +// group zero cap can never be the rejecting reason here, and a within-limits +// envelope is never admitted either. Group enforcement under +// admission.enabled=false is only reachable WITH the owner-id anchor (the +// tests above); this shape pins the trigger, not the enforcement. +// +// The pin: with the group envelope as the only trusted header, both a zero +// cap and a within-limits envelope answer 503 and never reach the upstream. +// Drop id.GroupScopes != "" from trustedPolicyEnvelopePresent and both are +// forwarded (200), failing this test. +func TestSharedGroupScopesOnlyEnvelopeFailsClosedWithAdmissionFlagOff(t *testing.T) { + up, hits := sharedGroupBackend(t) + s, _ := proxyWithGroupScopesEnabled(t, false, nil) + + // Only the group envelope: the owner-id anchor and every scoped + // rate-limit header deleted. + groupOnly := func(envelope string) *http.Request { + req := groupScopesRequest(t, up, envelope) + req.Header.Del(identity.HeaderOwnerID) + for _, header := range identity.ScopedRateLimitHeaders { + req.Header.Del(header) + } + return req + } + + rr := httptest.NewRecorder() + s.Handler().ServeHTTP(rr, groupOnly(groupScopesValue(testGroupA, "0,,,", ""))) + if rr.Code != http.StatusServiceUnavailable { + t.Fatalf("group-only zero-cap status=%d, want 503 (the contract-anchor parse fails closed before group scopes)", rr.Code) + } + + rr = httptest.NewRecorder() + s.Handler().ServeHTTP(rr, groupOnly(groupScopesValue(testGroupA, "30,,,", ""))) + if rr.Code != http.StatusServiceUnavailable { + t.Fatalf("group-only within-limits status=%d, want 503", rr.Code) + } + + if *hits != 0 { + t.Fatalf("upstream hits=%d, want 0 — a group-only request never forwards; with the group clause dropped from trustedPolicyEnvelopePresent both requests are admitted and this fails", *hits) + } +} diff --git a/internal/rating/group_rollup_integration_test.go b/internal/rating/group_rollup_integration_test.go index 1d4d3c6..85fb850 100644 --- a/internal/rating/group_rollup_integration_test.go +++ b/internal/rating/group_rollup_integration_test.go @@ -232,7 +232,11 @@ func TestIntegration_GroupUsageAttribution(t *testing.T) { // the rollup level, so neither clean event may attribute to gC at all; // - identical pair: two genuinely different but identical-looking billable // events in one rollup must BOTH attribute to gD — the DISTINCT ON -// (request_id, gid) dedupe must not collapse them into one. +// (request_id, gid) dedupe must not collapse them into one; +// - missing-usage twin: one event has usage_found=false, so only the +// usage_found gate keeps it out of the gE row; +// - invalid-usage twin: one event reports more cached than prompt tokens, +// so only the valid_usage gate keeps it out of the gF row. // // For each withheld case the group's group_usage row must carry exactly its // priced event(s), and the money rollup must be what a baseline WITHOUT the @@ -268,16 +272,20 @@ func TestIntegration_GroupUsageWithheldGatesSharedGrain(t *testing.T) { billed := rate.Quantized() const ( - gA = "a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a" - gB = "b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b" - gC = "c2c2c2c2c2c2c2c2c2c2c2c2c2c2c2c2c" - gD = "d3d3d3d3d3d3d3d3d3d3d3d3d3d3d3d3d" + gA = "a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0" + gB = "b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1" + gC = "c2c2c2c2c2c2c2c2c2c2c2c2c2c2c2c2" + gD = "d3d3d3d3d3d3d3d3d3d3d3d3d3d3d3d3" + gE = "e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4e4" + gF = "f5f5f5f5f5f5f5f5f5f5f5f5f5f5f5f5" + // gConf is the owner-conflicted event's own group (its group_id). + gConf = "9a9a9a9a9a9a9a9a9a9a9a9a9a9a9a9a" ) - // Every event is dedicated with authoritative usage; each withheld event is - // kept out of money by exactly ONE classification (its price, its owner - // conflict, or its rollup's two orgs) and shares its money grain with the - // priced event that must still attribute. + // Every event is dedicated; each withheld event is kept out of money by + // exactly ONE classification (its price, its owner conflict, its rollup's + // two orgs, its missing usage, or its invalid usage) and shares its money + // grain with the priced event that must still attribute. type gateEvent struct { req string auth string @@ -289,6 +297,8 @@ func TestIntegration_GroupUsageWithheldGatesSharedGrain(t *testing.T) { model string base string // billing_event.base_model ('' writes NULL) prompt int + cached int + noUsage bool // writes usage_found=false (the zero value is authoritative usage) } events := []gateEvent{ // (a) the unpriced twin: 'ga-ok' prices model 'm-unp' through base 'b' @@ -301,7 +311,7 @@ func TestIntegration_GroupUsageWithheldGatesSharedGrain(t *testing.T) { // drops it from money per event. Its memberships and its own group must // not reach group_usage. {req: "gb-ok", auth: "a-conf", members: []string{gB}, resource: "r-conf", model: "b", prompt: 25}, - {req: "gb-bad", auth: "a-conf", user: "u-x", group: "g-conflicted", members: []string{gB}, resource: "r-conf", model: "b", prompt: 777}, + {req: "gb-bad", auth: "a-conf", user: "u-x", group: gConf, members: []string{gB}, resource: "r-conf", model: "b", prompt: 777}, // (c) ambiguous_org: one resource, two distinct orgs — the rollup is // withheld as a whole, so neither clean event may attribute to gC. {req: "gc-1", auth: "a-org", members: []string{gC}, resource: "r-org", org: "org-x", model: "b", prompt: 50}, @@ -311,14 +321,22 @@ func TestIntegration_GroupUsageWithheldGatesSharedGrain(t *testing.T) { // key would collapse them into one. {req: "gd-1", auth: "a-twin", members: []string{gD}, resource: "r-twin", model: "b", prompt: 60}, {req: "gd-2", auth: "a-twin", members: []string{gD}, resource: "r-twin", model: "b", prompt: 60}, + // (e) the missing-usage twin: 'ge-bad' has no authoritative usage, so + // it is withheld as missing usage and must not reach gE. + {req: "ge-ok", auth: "a-miss", members: []string{gE}, resource: "r-miss", model: "b", prompt: 30}, + {req: "ge-bad", auth: "a-miss", members: []string{gE}, resource: "r-miss", model: "b", prompt: 888, noUsage: true}, + // (f) the invalid-usage twin: 'gf-bad' reports more cached than prompt + // tokens, so it is withheld as invalid usage and must not reach gF. + {req: "gf-ok", auth: "a-inv", members: []string{gF}, resource: "r-inv", model: "b", prompt: 35}, + {req: "gf-bad", auth: "a-inv", members: []string{gF}, resource: "r-inv", model: "b", prompt: 10, cached: 20}, } for i, e := range events { _, err := db.ExecContext(ctx, `INSERT INTO billing_event (request_id, auth_id, user_id, group_id, member_group_ids, resource_id, org_id, model, base_model, serving_mode, usage_found, prompt_tokens, cached_tokens, completion_tokens, event_ts) - VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,'dedicated',true,$10,0,0,$11)`, + VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,'dedicated',$10,$11,$12,0,$13)`, e.req, nullableStr(e.auth), nullableStr(e.user), nullableStr(e.group), groupIDsOrNil(e.members), nullableStr(e.resource), nullableStr(e.org), nullableStr(e.model), nullableStr(e.base), - e.prompt, hour.Add(time.Duration(i)*time.Minute)) + !e.noUsage, e.prompt, e.cached, hour.Add(time.Duration(i)*time.Minute)) if err != nil { t.Fatalf("seed event %d (%s): %v", i, e.req, err) } @@ -326,12 +344,14 @@ func TestIntegration_GroupUsageWithheldGatesSharedGrain(t *testing.T) { // Oracle for the money baseline WITHOUT the withheld events: 'ga-ok' prices // at the base 'b' rate through C4 rung (c), the others price model 'b' - // directly — all four resolve to the same quantized rate. + // directly — all six resolve to the same quantized rate. moneyEvents := []RatedEvent{ {AuthID: "a-unp", ResourceID: "r-unp", ModelID: "m-unp", BaseModel: "b", ServingMode: "dedicated", PromptTokens: 40}, {AuthID: "a-conf", ResourceID: "r-conf", ModelID: "b", ServingMode: "dedicated", PromptTokens: 25}, {AuthID: "a-twin", ResourceID: "r-twin", ModelID: "b", ServingMode: "dedicated", PromptTokens: 60}, {AuthID: "a-twin", ResourceID: "r-twin", ModelID: "b", ServingMode: "dedicated", PromptTokens: 60}, + {AuthID: "a-miss", ResourceID: "r-miss", ModelID: "b", ServingMode: "dedicated", PromptTokens: 30}, + {AuthID: "a-inv", ResourceID: "r-inv", ModelID: "b", ServingMode: "dedicated", PromptTokens: 35}, } costs := make([]Dec, len(moneyEvents)) for i, e := range moneyEvents { @@ -344,30 +364,36 @@ func TestIntegration_GroupUsageWithheldGatesSharedGrain(t *testing.T) { t.Fatalf("RateWindow: %v", err) } - // Money is exactly the priced baseline: the three priced rollups (one per - // grain), four rated events, and a total equal to the sum of the four + // Money is exactly the priced baseline: the five priced rollups (one per + // grain), six rated events, and a total equal to the sum of the six // oracle costs. The withheld events changed nothing. - if res.RollupsWritten != 3 || res.EventsRated != 4 { - t.Fatalf("rollups/events = %d/%d, want 3/4 (the priced baseline only)", res.RollupsWritten, res.EventsRated) + if res.RollupsWritten != 5 || res.EventsRated != 6 { + t.Fatalf("rollups/events = %d/%d, want 5/6 (the priced baseline only)", res.RollupsWritten, res.EventsRated) } - wantTotal := costs[0].Add(costs[1]).Add(costs[2]).Add(costs[3]).String() + wantTotalDec := costs[0] + for _, c := range costs[1:] { + wantTotalDec = wantTotalDec.Add(c) + } + wantTotal := wantTotalDec.String() if MustDec(res.TotalCost).String() != wantTotal { t.Fatalf("total cost = %s, want %s (the priced baseline only)", res.TotalCost, wantTotal) } // The withheld events are classified, not silently dropped: the anomaly - // partition accounts for all 8 seeded events exactly once. - if res.UnpricedEvents != 1 || res.OwnerConflictEvents != 1 || res.AmbiguousOrgEvents != 2 { - t.Fatalf("unpriced/conflict/ambig-org = %d/%d/%d, want 1/1/2", - res.UnpricedEvents, res.OwnerConflictEvents, res.AmbiguousOrgEvents) + // partition accounts for all 12 seeded events exactly once. + if res.UnpricedEvents != 1 || res.OwnerConflictEvents != 1 || res.AmbiguousOrgEvents != 2 || + res.MissingUsageEvents != 1 || res.InvalidUsageEvents != 1 { + t.Fatalf("unpriced/conflict/ambig-org/missing/invalid = %d/%d/%d/%d/%d, want 1/1/2/1/1", + res.UnpricedEvents, res.OwnerConflictEvents, res.AmbiguousOrgEvents, + res.MissingUsageEvents, res.InvalidUsageEvents) } if got := res.EventsRated + res.MissingUsageEvents + res.InvalidUsageEvents + res.UnpricedEvents + res.UnattributableEvents + res.InvalidServingModeEvents + - res.AmbiguousBaseEvents + res.AmbiguousOrgEvents + res.OwnerConflictEvents; got != 8 { - t.Fatalf("anomaly partition = %d, want 8 (all seeded events accounted exactly once)", got) + res.AmbiguousBaseEvents + res.AmbiguousOrgEvents + res.OwnerConflictEvents; got != 12 { + t.Fatalf("anomaly partition = %d, want 12 (all seeded events accounted exactly once)", got) } - if res.GroupRollupsWritten != 3 { - t.Fatalf("group rollups = %d, want 3 (groups %s, %s, %s)", - res.GroupRollupsWritten, gA[:8], gB[:8], gD[:8]) + if res.GroupRollupsWritten != 5 { + t.Fatalf("group rollups = %d, want 5 (groups %s, %s, %s, %s, %s)", + res.GroupRollupsWritten, gA[:8], gB[:8], gD[:8], gE[:8], gF[:8]) } type groupRow struct { @@ -409,7 +435,7 @@ func TestIntegration_GroupUsageWithheldGatesSharedGrain(t *testing.T) { if rB.count != 1 || rB.prompt != 25 || MustDec(rB.cost).String() != costs[1].String() { t.Fatalf("group %s row = %+v, want count 1 / prompt 25 / cost %s (the clean twin only)", gB, rB, costs[1]) } - if _, ok := readRow("g-conflicted"); ok { + if _, ok := readRow(gConf); ok { t.Fatalf("the conflicted event's own group reached group_usage; the owner_conflict twin must attribute nowhere") } @@ -436,6 +462,28 @@ func TestIntegration_GroupUsageWithheldGatesSharedGrain(t *testing.T) { t.Fatalf("group %s row = %+v, want count 2 / prompt 120 / cost %s (both events attribute)", gD, rD, twinCost) } + + // (e) gE carries 'ge-ok' alone: dropping the usage_found gate would add + // 'ge-bad' (count 2, prompt 918). + rE, ok := readRow(gE) + if !ok { + t.Fatalf("no group_usage row for %s", gE) + } + if rE.count != 1 || rE.prompt != 30 || rE.billable != 30 || MustDec(rE.cost).String() != costs[4].String() { + t.Fatalf("group %s row = %+v, want count 1 / prompt 30 / billable 30 / cost %s (the priced twin only)", + gE, rE, costs[4]) + } + + // (f) gF carries 'gf-ok' alone: dropping the valid_usage gate would add + // 'gf-bad' (count 2, prompt 45). + rF, ok := readRow(gF) + if !ok { + t.Fatalf("no group_usage row for %s", gF) + } + if rF.count != 1 || rF.prompt != 35 || rF.billable != 35 || MustDec(rF.cost).String() != costs[5].String() { + t.Fatalf("group %s row = %+v, want count 1 / prompt 35 / billable 35 / cost %s (the priced twin only)", + gF, rF, costs[5]) + } } // groupIDsOrNil binds a membership list as a driver value (nil for none — diff --git a/migrations/0008_group_scopes.up.sql b/migrations/0008_group_scopes.up.sql index 561fc26..89e2cbe 100644 --- a/migrations/0008_group_scopes.up.sql +++ b/migrations/0008_group_scopes.up.sql @@ -49,7 +49,8 @@ CREATE TABLE group_usage ( group_id TEXT NOT NULL, -- The UTC hour bucket, same bucketing expression the rater uses for -- rated_usage.window_start, so a group's monthly spend sum reads - -- window_start >= date_trunc('month', now()) over aligned hours. + -- window_start >= date_trunc('month', now() AT TIME ZONE 'UTC') AT TIME ZONE 'UTC' + -- over aligned hours. window_start TIMESTAMPTZ NOT NULL, prompt_tokens BIGINT NOT NULL DEFAULT 0, cached_tokens BIGINT NOT NULL DEFAULT 0, diff --git a/migrations/README.md b/migrations/README.md index 2b3ba84..2bc04f9 100644 --- a/migrations/README.md +++ b/migrations/README.md @@ -78,6 +78,12 @@ failure modes are asymmetric, so the order matters: rollup (zero group spend) until the new rater covers those hours. - A post-0008 rater fails against the old schema (SQLSTATE 42703): its rating statement references `member_group_ids` and `group_usage`. +- A post-0008 drainer fails against the old schema the same way (SQLSTATE + 42703): its INSERT names `member_group_ids`, which does not exist there yet, + so EVERY insert fails and every event takes the poison path — logged, ACK'd, + and dropped, served but never billed. `cmd/migrate up` must come first, + before BOTH the drainer and the rater deploy; the drainer is not safe in + either order. So roll code and schema together, in the migration-0007 order: run `cmd/migrate up`, then deploy the new drainer and rater (the drainer's INSERT From f784bae9f2a974371a16164f88c08236d3f19abf Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 11:21:25 +0000 Subject: [PATCH 13/17] proxy: pin the group generated-token cap in the undeclared-max_tokens reservation clamp --- internal/proxy/group_scopes_test.go | 100 ++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/internal/proxy/group_scopes_test.go b/internal/proxy/group_scopes_test.go index 54e63bd..0830745 100644 --- a/internal/proxy/group_scopes_test.go +++ b/internal/proxy/group_scopes_test.go @@ -3,6 +3,7 @@ package proxy import ( "context" "errors" + "io" "net/http" "net/http/httptest" "net/url" @@ -602,3 +603,102 @@ func TestSharedGroupScopesOnlyEnvelopeFailsClosedWithAdmissionFlagOff(t *testing t.Fatalf("upstream hits=%d, want 0 — a group-only request never forwards; with the group clause dropped from trustedPolicyEnvelopePresent both requests are admitted and this fails", *hits) } } + +// TestSharedGroupGeneratedCapClampsUndeclaredMaxTokensReservation: the +// undeclared-max_tokens output reservation is clamped to the tightest +// positive generated-token limit across EVERY scope class on the request — +// organization, owner, and each stamped group scope. A group whose generated +// cap sits below the configured default reserves the cap, not the default, so +// its undeclared requests admit instead of reserving the default and hitting +// the unsatisfiable 400 of +// TestMaxTokensAboveGeneratedLimitIsUnsatisfiable400. Dropping the group +// limits from the undeclaredOutputReservation call at proxy.go leaves every +// case here green only if it also breaks the group cap below the default — +// the end-to-end half fails: the unclamped 20-token default reservation over +// a 10-token group cap is rejected before upstream. +func TestSharedGroupGeneratedCapClampsUndeclaredMaxTokensReservation(t *testing.T) { + up, _ := sharedGroupBackend(t) + cfg := proxyAdmissionConfig(10) // DefaultMaxOutputTokens: 20 + + // noMaxTokens swaps the shared body ({"model":"model-a","max_tokens":20}) + // for an undeclared-output one, the shape the reservation clamp exists + // for. + noMaxTokens := func(req *http.Request) { + req.Body = io.NopCloser(strings.NewReader(`{"model":"model-a"}`)) + req.ContentLength = -1 + } + noMaxTokensReq := func(req *http.Request) *http.Request { + noMaxTokens(req) + return req + } + send := func(recording *recordingAdmitter, envelope string, mutate func(*http.Request)) *httptest.ResponseRecorder { + s := New(&config.Settings{Admission: cfg}, logging.New(logging.ERROR), &recordingEmitter{}).WithAdmitter(recording) + req := groupScopesRequest(t, up, envelope) + if mutate != nil { + mutate(req) + } + rr := httptest.NewRecorder() + s.Handler().ServeHTTP(rr, req) + return rr + } + + // The group generated cap below the default drives the reservation: 10, + // not the 20-token default. + recording := &recordingAdmitter{} + rr := send(recording, groupScopesValue(testGroupA, ",,,10", ""), noMaxTokens) + if rr.Code != http.StatusOK || recording.last.ReservedOutputTokens != 10 { + t.Fatalf("group cap 10: status=%d reserved=%d, want 200 with a 10-token reservation (the group cap, not the 20-token default)", rr.Code, recording.last.ReservedOutputTokens) + } + + // A group cap above the default changes nothing: the default stands. + recording = &recordingAdmitter{} + rr = send(recording, groupScopesValue(testGroupA, ",,,30", ""), noMaxTokens) + if rr.Code != http.StatusOK || recording.last.ReservedOutputTokens != 20 { + t.Fatalf("group cap 30: status=%d reserved=%d, want 200 with the 20-token default (a cap above the default does not raise the reservation)", rr.Code, recording.last.ReservedOutputTokens) + } + + // Every scope class feeds the same clamp and the tightest wins: org 5, + // owner 8, group 10 → 5. + recording = &recordingAdmitter{} + rr = send(recording, groupScopesValue(testGroupA, ",,,10", ""), func(req *http.Request) { + req.Header.Set(identity.HeaderOrgRateLimitGeneratedTokens, "5") + req.Header.Set(identity.HeaderOwnerRateLimitGeneratedTokens, "8") + noMaxTokens(req) + }) + if rr.Code != http.StatusOK || recording.last.ReservedOutputTokens != 5 { + t.Fatalf("org 5 / owner 8 / group 10: status=%d reserved=%d, want 200 with the tightest (5)", rr.Code, recording.last.ReservedOutputTokens) + } + + // End to end against the real store: the undeclared request from a + // group capped below the default is ADMITTED. Without the group limits in + // the clamp the reservation stays 20 over the 10-token cap and this is + // the unsatisfiable 400 — this half fails on that mutation. + s, _ := proxyWithGroupScopes(t, nil) + rr = httptest.NewRecorder() + s.Handler().ServeHTTP(rr, noMaxTokensReq(groupScopesRequest(t, up, groupScopesValue(testGroupA, ",,,10", "")))) + if rr.Code != http.StatusOK { + t.Fatalf("undeclared request, group cap 10: status=%d body=%q, want 200 admitted (the reservation clamped to the cap)", rr.Code, strings.TrimSpace(rr.Body.String())) + } + + // The declared path is unchanged: max_tokens above the group cap can + // never be admitted — 400 without Retry-After, the group-scope twin of + // TestMaxTokensAboveGeneratedLimitIsUnsatisfiable400. + hits := 0 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits++ + _, _ = w.Write([]byte(`{"model":"model-a","usage":{"prompt_tokens":2,"completion_tokens":3}}`)) + })) + t.Cleanup(backend.Close) + declUp, err := url.Parse(backend.URL) + if err != nil { + t.Fatal(err) + } + rr = httptest.NewRecorder() + s.Handler().ServeHTTP(rr, groupScopesRequest(t, declUp, groupScopesValue(testGroupA, ",,,10", ""))) + if rr.Code != http.StatusBadRequest || rr.Header().Get("Retry-After") != "" || hits != 0 { + t.Fatalf("declared max_tokens 20 over group cap 10: status=%d Retry-After=%q hits=%d, want 400 without Retry-After before upstream", rr.Code, rr.Header().Get("Retry-After"), hits) + } + if got := strings.TrimSpace(rr.Body.String()); got != "max_tokens exceeds the per-window generated-token limit" { + t.Fatalf("body=%q", got) + } +} From accf4d600dce50542856e7d2e25289420d7d219e Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 11:21:25 +0000 Subject: [PATCH 14/17] interceptor: test the buildAdmission DATABASE_URL spend-store wiring --- cmd/interceptor/main_integration_test.go | 137 +++++++++++++++++++++++ cmd/interceptor/main_test.go | 79 +++++++++++++ 2 files changed, 216 insertions(+) create mode 100644 cmd/interceptor/main_integration_test.go diff --git a/cmd/interceptor/main_integration_test.go b/cmd/interceptor/main_integration_test.go new file mode 100644 index 0000000..111f260 --- /dev/null +++ b/cmd/interceptor/main_integration_test.go @@ -0,0 +1,137 @@ +//go:build integration + +// Package main integration test: runs buildAdmission's DATABASE_URL spend +// wiring against a LIVE Postgres loaded with the production migrations, so the +// "group spend caps enabled" startup line is proven to describe a working +// check — the unit test pins the two failure logs, but only this half +// exercises OpenPostgresSpendStore + WithGroupSpend together. +// +// Gated behind the `integration` build tag AND a non-empty +// PHOEBE_TEST_DATABASE_URL. Run with: +// +// PHOEBE_TEST_DATABASE_URL=postgres://... go test -tags=integration ./cmd/interceptor/... +package main + +import ( + "bytes" + "context" + "database/sql" + "errors" + "fmt" + "io" + "log" + "os" + "sort" + "strings" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + _ "github.com/jackc/pgx/v5/stdlib" + + "github.com/saturncloud/phoebe/internal/admission" + "github.com/saturncloud/phoebe/internal/logging" + "github.com/saturncloud/phoebe/migrations" +) + +// newSpendWiringHarness creates an isolated schema — named per test PROCESS so +// concurrent runs against a shared database stop stomping each other — with +// ALL production migrations applied (discovered from the embedded +// migrations.FS, ordered by filename; a hand-maintained list could go stale +// while green), and returns a pool pinned to it plus the schema-pinned DSN for +// handing to buildAdmission as DATABASE_URL. +func newSpendWiringHarness(t *testing.T) (*sql.DB, string) { + t.Helper() + dsn := os.Getenv("PHOEBE_TEST_DATABASE_URL") + if dsn == "" { + t.Skip("PHOEBE_TEST_DATABASE_URL not set; skipping interceptor spend-wiring integration test") + } + schema := fmt.Sprintf("phoebe_interceptor_spend_it_%d", os.Getpid()) + + admin, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatalf("open admin pool: %v", err) + } + t.Cleanup(func() { + _, _ = admin.Exec("DROP SCHEMA IF EXISTS " + schema + " CASCADE") + _ = admin.Close() + }) + execSpendWiring(t, admin, "DROP SCHEMA IF EXISTS "+schema+" CASCADE") + execSpendWiring(t, admin, "CREATE SCHEMA "+schema) + + sep := "?" + if strings.Contains(dsn, "?") { + sep = "&" + } + pinned := dsn + sep + "search_path=" + schema + db, err := sql.Open("pgx", pinned) + if err != nil { + t.Fatalf("open schema pool: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + entries, err := migrations.FS.ReadDir(".") + if err != nil { + t.Fatalf("read migrations FS: %v", err) + } + names := make([]string, 0, len(entries)) + for _, e := range entries { + if !e.IsDir() && strings.HasSuffix(e.Name(), ".up.sql") { + names = append(names, e.Name()) + } + } + sort.Strings(names) + for _, name := range names { + b, err := migrations.FS.ReadFile(name) + if err != nil { + t.Fatalf("read migration %s: %v", err, name) + } + execSpendWiring(t, db, string(b)) + } + return db, pinned +} + +func execSpendWiring(t *testing.T, db *sql.DB, stmt string) { + t.Helper() + if _, err := db.Exec(stmt); err != nil { + t.Fatalf("exec failed: %v\nstatement: %s", err, stmt) + } +} + +// TestIntegration_BuildAdmissionGroupSpendCapsEnabled is the positive half of +// the DATABASE_URL wiring: against live Postgres, buildAdmission logs the +// "enabled" line AND the admitter it returns enforces the monthly spend cap — +// a spend-capped Admit whose group_usage spend reached the cap is the +// contractual monthly_spend rejection. A regression deleting the +// WithGroupSpend call admits this request, so the wiring's production switch +// is now mutation-covered, not just logged. +func TestIntegration_BuildAdmissionGroupSpendCapsEnabled(t *testing.T) { + db, pinned := newSpendWiringHarness(t) + t.Setenv("DATABASE_URL", pinned) + mr := miniredis.RunT(t) + var buf bytes.Buffer + logger := &logging.Logger{Debug: log.New(io.Discard, "", 0), Info: log.New(&buf, "", 0), Warn: log.New(&buf, "", 0), Error: log.New(&buf, "", 0)} + admitter, closeAdmission := buildAdmission(loadTestSettings(t, "emit:\n valkeyAddr: "+mr.Addr()+"\n"), logger) + defer closeAdmission() + if admitter == nil { + t.Fatal("no admitter built") + } + if out := buf.String(); !strings.Contains(out, "group spend caps enabled") { + t.Fatalf("startup log missing the enabled line:\n%s", out) + } + + // The group's month-to-date attribution spend is 10 — over the cap of 1 — + // in the current UTC hour bucket, so the month predicate cannot skip it. + _, err := db.Exec(`INSERT INTO group_usage (group_id, window_start, cost, event_count) + VALUES ($1, $2, '10', 1)`, spendWiringGID, time.Now().UTC().Truncate(time.Hour)) + if err != nil { + t.Fatalf("seed group_usage: %v", err) + } + + req := spendWiringRequest(admission.GroupScope{GroupID: spendWiringGID, SpendCap: "1"}) + _, err = admitter.Admit(context.Background(), req) + var rejected *admission.Rejected + if !errors.As(err, &rejected) || !rejected.Contractual || rejected.Dimension != "monthly_spend" { + t.Fatalf("spend-capped admit = %v, want a contractual monthly_spend rejection (the wired spend check reads group_usage)", err) + } +} diff --git a/cmd/interceptor/main_test.go b/cmd/interceptor/main_test.go index 797875f..8f1fba0 100644 --- a/cmd/interceptor/main_test.go +++ b/cmd/interceptor/main_test.go @@ -3,6 +3,7 @@ package main import ( "bytes" "context" + "errors" "io" "log" "os" @@ -91,3 +92,81 @@ func TestBuildAdmissionLogsMeteringStoreFallback(t *testing.T) { t.Fatalf("fallback logged although admission.valkeyAddr is set: %q", explicit) } } + +const spendWiringGID = "a1b2c3d4e5f60718293a4b5c6d7e8f90" + +// spendWiringRequest is a shared-shaped admit request carrying one group +// scope; scope carries either the rate limit or the spend cap under test. +func spendWiringRequest(scope admission.GroupScope) admission.Request { + return admission.Request{ + Graph: "graph-a", Organization: "org-a", Owner: "owner-a", Model: "m", + PromptBytes: 10, EstimatedInputTokens: 10, ReservedOutputTokens: 20, + GroupScopes: []admission.GroupScope{scope}, + } +} + +// TestBuildAdmissionGroupSpendCapsDatabaseURL pins the production DATABASE_URL +// switch that turns the monthly group spend cap on. This wiring is the +// feature's only production call site, and every other spend test wires the +// store by hand — so a regression here (WithGroupSpend not called on a +// successful open, or the open error made fatal) disables group spend caps on +// every interceptor, or crash-loops it, with the rest of the suite green. +func TestBuildAdmissionGroupSpendCapsDatabaseURL(t *testing.T) { + build := func(t *testing.T) (admission.Admitter, *bytes.Buffer) { + t.Helper() + mr := miniredis.RunT(t) + var buf bytes.Buffer + logger := &logging.Logger{Debug: log.New(io.Discard, "", 0), Info: log.New(&buf, "", 0), Warn: log.New(&buf, "", 0), Error: log.New(&buf, "", 0)} + admitter, closeAdmission := buildAdmission(loadTestSettings(t, "emit:\n valkeyAddr: "+mr.Addr()+"\n"), logger) + t.Cleanup(closeAdmission) + return admitter, &buf + } + + t.Run("unset logs NOT enforced and group rate limits still enforce", func(t *testing.T) { + t.Setenv("DATABASE_URL", "") + admitter, buf := build(t) + if admitter == nil { + t.Fatal("no admitter with DATABASE_URL unset: contract and group rate limits would not be enforced") + } + if out := buf.String(); !strings.Contains(out, "group spend caps are NOT enforced") { + t.Fatalf("startup log missing the loud NOT-enforced line:\n%s", out) + } + + // The startup line's promise, pinned behaviorally: group rate limits + // still enforce — a capped window rejects the second request with the + // contractual 429 mapping. + one := int64(1) + rateScope := admission.GroupScope{GroupID: spendWiringGID, Limits: admission.RateLimits{Requests: &one}} + lease, err := admitter.Admit(context.Background(), spendWiringRequest(rateScope)) + if err != nil { + t.Fatalf("first group-rate admit: %v, want admitted", err) + } + _ = lease.Complete(context.Background(), 0) + _, err = admitter.Admit(context.Background(), spendWiringRequest(rateScope)) + var rejected *admission.Rejected + if !errors.As(err, &rejected) || !rejected.Contractual { + t.Fatalf("second group-rate admit: %v, want a contractual rejection (group rate limits still enforce without DATABASE_URL)", err) + } + + // ... and the spend cap does not: a spend-capped group admits + // fail-open with no store behind it. + cappedScope := admission.GroupScope{GroupID: spendWiringGID, SpendCap: "0"} + if _, err := admitter.Admit(context.Background(), spendWiringRequest(cappedScope)); err != nil { + t.Fatalf("spend-capped admit with no store: %v, want admitted (spend caps are NOT enforced without DATABASE_URL)", err) + } + }) + + t.Run("unroutable DSN logs DISABLED and serving continues", func(t *testing.T) { + // 127.0.0.1:1 answers with connection refused, so the 5s open budget + // is not what bounds this test. + t.Setenv("DATABASE_URL", "postgres://postgres:pw@127.0.0.1:1/postgres?sslmode=disable") + admitter, buf := build(t) + if admitter == nil { + t.Fatal("no admitter with an unroutable DATABASE_URL: serving must continue without spend caps") + } + out := buf.String() + if !strings.Contains(out, "group spend caps DISABLED") || !strings.Contains(out, "group rate limits still enforce") { + t.Fatalf("startup log missing the loud DISABLED line naming the surviving rate limits:\n%s", out) + } + }) +} From 3b0dbe847b5fa451ee0149309ced97ff90680b8a Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 11:43:54 +0000 Subject: [PATCH 15/17] admission: rename cap parameter to spendCap (revive redefines-builtin-id) --- internal/admission/group_test.go | 4 ++-- internal/admission/spend.go | 18 +++++++++--------- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/internal/admission/group_test.go b/internal/admission/group_test.go index 68d4274..5474fef 100644 --- a/internal/admission/group_test.go +++ b/internal/admission/group_test.go @@ -29,12 +29,12 @@ type fakeGroupSpendStore struct { gotCap []string } -func (f *fakeGroupSpendStore) GroupSpendExhausted(_ context.Context, groupID, cap string) (bool, error) { +func (f *fakeGroupSpendStore) GroupSpendExhausted(_ context.Context, groupID, spendCap string) (bool, error) { f.mu.Lock() defer f.mu.Unlock() f.calls++ f.gotGID = append(f.gotGID, groupID) - f.gotCap = append(f.gotCap, cap) + f.gotCap = append(f.gotCap, spendCap) return f.exhausted, f.err } diff --git a/internal/admission/spend.go b/internal/admission/spend.go index 7ff239f..6b3f614 100644 --- a/internal/admission/spend.go +++ b/internal/admission/spend.go @@ -58,7 +58,7 @@ type GroupSpendStore interface { // (NUMERIC, never a Go number). cap is a plain decimal NUMERIC(20,9) // string straight from the trusted envelope. The verdict — not the spend // sum — crosses this seam, so no money value becomes a Go number. - GroupSpendExhausted(ctx context.Context, groupID, cap string) (bool, error) + GroupSpendExhausted(ctx context.Context, groupID, spendCap string) (bool, error) } // PostgresSpendStore answers the monthly spend check against group_usage @@ -125,9 +125,9 @@ FROM group_usage WHERE group_id = $1 AND window_start >= date_trunc('month', now() AT TIME ZONE 'UTC') AT TIME ZONE 'UTC'` -func (s *PostgresSpendStore) GroupSpendExhausted(ctx context.Context, groupID, cap string) (bool, error) { +func (s *PostgresSpendStore) GroupSpendExhausted(ctx context.Context, groupID, spendCap string) (bool, error) { var exhausted bool - if err := s.db.QueryRowContext(ctx, groupSpendExhaustedSQL, groupID, cap).Scan(&exhausted); err != nil { + if err := s.db.QueryRowContext(ctx, groupSpendExhaustedSQL, groupID, spendCap).Scan(&exhausted); err != nil { return false, fmt.Errorf("admission: spend check: group %s: %w", groupID, err) } return exhausted, nil @@ -170,10 +170,10 @@ func newSpendVerdictCache() *spendVerdictCache { // spend, so a verdict taken in the old month says nothing about the new one. // The negative (store-error) entries use the shorter groupSpendNegativeCacheTTL // and skip the month check — see spendVerdictEntry.failed. -func (c *spendVerdictCache) lookup(groupID, cap string) (spendVerdictEntry, bool) { +func (c *spendVerdictCache) lookup(groupID, spendCap string) (spendVerdictEntry, bool) { c.mu.Lock() defer c.mu.Unlock() - e, ok := c.entries[groupID+"\x00"+cap] + e, ok := c.entries[groupID+"\x00"+spendCap] if !ok { return spendVerdictEntry{}, false } @@ -193,27 +193,27 @@ func (c *spendVerdictCache) lookup(groupID, cap string) (spendVerdictEntry, bool } // store records a fresh verdict for (groupID, cap). -func (c *spendVerdictCache) store(groupID, cap string, exhausted bool) { +func (c *spendVerdictCache) store(groupID, spendCap string, exhausted bool) { c.mu.Lock() defer c.mu.Unlock() now := time.Now if c.now != nil { now = c.now } - c.entries[groupID+"\x00"+cap] = spendVerdictEntry{exhausted: exhausted, at: now()} + c.entries[groupID+"\x00"+spendCap] = spendVerdictEntry{exhausted: exhausted, at: now()} } // storeFailure records a fail-open entry for (groupID, cap) after a store // error: for groupSpendNegativeCacheTTL the check admits without querying and // without waiting, then re-probes the store. -func (c *spendVerdictCache) storeFailure(groupID, cap string) { +func (c *spendVerdictCache) storeFailure(groupID, spendCap string) { c.mu.Lock() defer c.mu.Unlock() now := time.Now if c.now != nil { now = c.now } - c.entries[groupID+"\x00"+cap] = spendVerdictEntry{failed: true, at: now()} + c.entries[groupID+"\x00"+spendCap] = spendVerdictEntry{failed: true, at: now()} } // nextUTCMonth returns the first instant of the UTC calendar month after t's. From 354034c8d366fbeed62523f72dfe6ba4931fbdc1 Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 12:12:40 +0000 Subject: [PATCH 16/17] admission: rename test cap locals to spendCap (revive redefines-builtin-id) --- internal/admission/group_test.go | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/internal/admission/group_test.go b/internal/admission/group_test.go index 5474fef..9918f4d 100644 --- a/internal/admission/group_test.go +++ b/internal/admission/group_test.go @@ -76,8 +76,8 @@ func TestGroupRateLimitsEnforcedLikeContractScope(t *testing.T) { // Only the requests window is capped, so the request's own estimates // (the helper's 10-in/20-out shape) can never trip the unsatisfiable // check: the third request is rejected by the window counter itself. - cap := int64(2) - scope := GroupScope{GroupID: testGID, Limits: RateLimits{Requests: &cap}} + spendCap := int64(2) + scope := GroupScope{GroupID: testGID, Limits: RateLimits{Requests: &spendCap}} for i := 0; i < 2; i++ { lease, err := a.Admit(context.Background(), groupRequest("org-g", "m", scope)) @@ -693,9 +693,9 @@ func TestGroupSpendNoStoreFailsOpen(t *testing.T) { // distinct from the store-outage bypass. func TestGroupScopeMissingGroupIDFailsClosed(t *testing.T) { a, _ := spendTestAdmitter(t) - cap := int64(10) + spendCap := int64(10) _, err := a.Admit(context.Background(), groupRequest("org-i", "m", - GroupScope{Limits: RateLimits{Requests: &cap}})) + GroupScope{Limits: RateLimits{Requests: &spendCap}})) if !errors.Is(err, ErrInvalidIdentity) { t.Fatalf("anonymous group scope: %v, want ErrInvalidIdentity (fail closed)", err) } @@ -715,9 +715,9 @@ func TestGroupScopeMissingGroupIDFailsClosed(t *testing.T) { // reservation that would push the window past 50 is rejected. func TestGroupScopesSettleThroughCompleteUsage(t *testing.T) { a, _ := spendTestAdmitter(t) - cap := int64(50) + spendCap := int64(50) r := groupRequest("org-l", "m", GroupScope{GroupID: testGID, - Limits: RateLimits{GeneratedTokens: &cap}}) + Limits: RateLimits{GeneratedTokens: &spendCap}}) lease, err := a.Admit(context.Background(), r) if err != nil { @@ -750,9 +750,9 @@ func TestGroupScopesSettleThroughCompleteUsage(t *testing.T) { // reservation — without settlement the third would be rejected at 60/50. func TestGroupScopesSettleZeroUsageRefundsReservation(t *testing.T) { a, _ := spendTestAdmitter(t) - cap := int64(50) + spendCap := int64(50) r := groupRequest("org-z0", "m", GroupScope{GroupID: testGID, - Limits: RateLimits{GeneratedTokens: &cap}}) + Limits: RateLimits{GeneratedTokens: &spendCap}}) for i := 0; i < 5; i++ { lease, err := a.Admit(context.Background(), r) @@ -771,9 +771,9 @@ func TestGroupScopesSettleZeroUsageRefundsReservation(t *testing.T) { // would have fit it at 40/50. func TestGroupScopesSettleChargesActualUsage(t *testing.T) { a, _ := spendTestAdmitter(t) - cap := int64(50) + spendCap := int64(50) r := groupRequest("org-l45", "m", GroupScope{GroupID: testGID, - Limits: RateLimits{GeneratedTokens: &cap}}) + Limits: RateLimits{GeneratedTokens: &spendCap}}) lease, err := a.Admit(context.Background(), r) if err != nil { From edada951cd0ba5f30160392b610fb5ef1faf7981 Mon Sep 17 00:00:00 2001 From: hugo Date: Wed, 7 Oct 2026 12:13:52 +0000 Subject: [PATCH 17/17] admission: rename remaining cap identifiers in integration tests (revive) --- internal/admission/real_valkey_integration_test.go | 4 ++-- internal/admission/spend_integration_test.go | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/internal/admission/real_valkey_integration_test.go b/internal/admission/real_valkey_integration_test.go index 745c2d7..b4c858d 100644 --- a/internal/admission/real_valkey_integration_test.go +++ b/internal/admission/real_valkey_integration_test.go @@ -447,11 +447,11 @@ func TestRealValkeyGroupScopes(t *testing.T) { time.Sleep(time.Duration(61-s) * time.Second) } - cap := int64(50) + spendCap := int64(50) scopes := make([]GroupScope, 16) for i := range scopes { scopes[i] = GroupScope{GroupID: fmt.Sprintf("%032x", i+1), - Limits: RateLimits{GeneratedTokens: &cap}} + Limits: RateLimits{GeneratedTokens: &spendCap}} } r := request("org-g", "model") r.GroupScopes = scopes diff --git a/internal/admission/spend_integration_test.go b/internal/admission/spend_integration_test.go index 2ef9308..13a19d6 100644 --- a/internal/admission/spend_integration_test.go +++ b/internal/admission/spend_integration_test.go @@ -141,7 +141,7 @@ func TestIntegration_GroupSpendExhausted(t *testing.T) { for _, tc := range []struct { name string group string - cap string + spendCap string want bool }{ // 1. Zero cap, no rows at all for the group: exhausted. @@ -162,13 +162,13 @@ func TestIntegration_GroupSpendExhausted(t *testing.T) { // time zone would drop this row late in the month and admit. {"current month first hour counted", gStart, "2.5", true}, } { - got, err := store.GroupSpendExhausted(ctx, tc.group, tc.cap) + got, err := store.GroupSpendExhausted(ctx, tc.group, tc.spendCap) if err != nil { - t.Fatalf("%s: GroupSpendExhausted(%s, %s): %v", tc.name, tc.group, tc.cap, err) + t.Fatalf("%s: GroupSpendExhausted(%s, %s): %v", tc.name, tc.group, tc.spendCap, err) } if got != tc.want { t.Errorf("%s: GroupSpendExhausted(%s, %s) = %v, want %v", - tc.name, tc.group, tc.cap, got, tc.want) + tc.name, tc.group, tc.spendCap, got, tc.want) } } }