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 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) + } + }) +} diff --git a/internal/admission/admission.go b/internal/admission/admission.go index 97ad289..18670b6 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,30 @@ 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 + spendFlight *spendFlight + 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.spendFlight = newSpendFlight() + a.spendFails = &spendFailureLog{} + return a } func New(client redis.Cmdable, cfg config.AdmissionSettings) *RedisAdmitter { @@ -199,9 +244,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 +446,130 @@ 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 + } + // 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 + } + 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 + } + if budgetCtx == nil { + var cancelBudget context.CancelFunc + budgetCtx, cancelBudget = context.WithTimeout(ctx, groupSpendQueryBudget) + 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. + 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. 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. + startBudget = time.Until(deadline) + queryCtx, cancel := context.WithDeadline(context.WithoutCancel(ctx), deadline) + 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. + // + // 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) + } + continue + } + 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} + } + } + 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..9918f4d --- /dev/null +++ b/internal/admission/group_test.go @@ -0,0 +1,856 @@ +package admission + +import ( + "bytes" + "context" + "errors" + "fmt" + "regexp" + "strings" + "sync" + "testing" + "time" + + sqlmock "github.com/DATA-DOG/go-sqlmock" + + "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 { + mu sync.Mutex + exhausted bool + err error + calls int + gotGID []string + gotCap []string +} + +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, spendCap) + 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. + 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)) + 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) + } + + // 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 +// (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). 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", + groupScope(testGID, nil, "100"))) + if err != nil { + t.Fatalf("store-error admit: %v, want fail-open admission", err) + } + _ = lease.Complete(context.Background(), 0) + + // 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 (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) + } +} + +// 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) + // 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()) { + t.Fatalf("bypass log = %q, want a fail-open ERROR naming the deadline", out) + } +} + +// 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) + } +} + +// 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)) + + joined := make(chan struct{}, 1) + a.spendFlight.joined = func(string) { joined <- struct{}{} } + + 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 read answer only once the waiter has joined it. + <-joined + 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) + } +} + +// 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. +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) { + 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. +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) + spendCap := int64(10) + _, err := a.Admit(context.Background(), groupRequest("org-i", "m", + GroupScope{Limits: RateLimits{Requests: &spendCap}})) + 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) + spendCap := int64(50) + r := groupRequest("org-l", "m", GroupScope{GroupID: testGID, + Limits: RateLimits{GeneratedTokens: &spendCap}}) + + 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) + } +} + +// 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) + spendCap := int64(50) + r := groupRequest("org-z0", "m", GroupScope{GroupID: testGID, + Limits: RateLimits{GeneratedTokens: &spendCap}}) + + 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) + spendCap := int64(50) + r := groupRequest("org-l45", "m", GroupScope{GroupID: testGID, + Limits: RateLimits{GeneratedTokens: &spendCap}}) + + 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..b4c858d 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) + } + + spendCap := int64(50) + scopes := make([]GroupScope, 16) + for i := range scopes { + scopes[i] = GroupScope{GroupID: fmt.Sprintf("%032x", i+1), + Limits: RateLimits{GeneratedTokens: &spendCap}} + } + 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 new file mode 100644 index 0000000..6b3f614 --- /dev/null +++ b/internal/admission/spend.go @@ -0,0 +1,326 @@ +package admission + +import ( + "context" + "database/sql" + "fmt" + "sync" + "time" + + "github.com/jackc/pgx/v5" + "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 + +// 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 +// 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 + +// 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. +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 +// 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 >= 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. + GroupSpendExhausted(ctx context.Context, groupID, spendCap 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)") + } + // 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: 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. + 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. 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() AT TIME ZONE 'UTC') AT TIME ZONE 'UTC'` + +func (s *PostgresSpendStore) GroupSpendExhausted(ctx context.Context, groupID, spendCap string) (bool, error) { + var exhausted bool + 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 +} + +// 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 + // 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 { + return &spendVerdictCache{entries: map[string]spendVerdictEntry{}} +} + +// lookup returns the cached verdict for (groupID, cap) and whether it is +// 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. +// The negative (store-error) entries use the shorter groupSpendNegativeCacheTTL +// and skip the month check — see spendVerdictEntry.failed. +func (c *spendVerdictCache) lookup(groupID, spendCap string) (spendVerdictEntry, bool) { + c.mu.Lock() + defer c.mu.Unlock() + e, ok := c.entries[groupID+"\x00"+spendCap] + if !ok { + return spendVerdictEntry{}, false + } + now := time.Now + if c.now != nil { + now = c.now + } + t := now() + 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, true +} + +// store records a fresh verdict for (groupID, cap). +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"+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, spendCap string) { + c.mu.Lock() + defer c.mu.Unlock() + now := time.Now + if c.now != nil { + now = c.now + } + c.entries[groupID+"\x00"+spendCap] = 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() + 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 +// 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 + // 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() + clock := time.Now + if l.now != nil { + clock = l.now + } + now := clock() + 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) +} + +// 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 + // 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 { + 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. +// 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 { + joined := f.joined + f.mu.Unlock() + if joined != nil { + joined(key) + } + 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 + 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..13a19d6 --- /dev/null +++ b/internal/admission/spend_integration_test.go @@ -0,0 +1,181 @@ +//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; +// - 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 == "" { + 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) + // 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)) + + // 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" + gStart = "ffeeddccbbaa99887766554433221100" + ) + 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") + // 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 { + name string + group string + spendCap 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}, + // 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.spendCap) + if err != nil { + 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.spendCap, 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) + } +} 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/e2e/e2e_test.go b/internal/e2e/e2e_test.go index a4c0eab..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" @@ -155,6 +156,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 { @@ -1062,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/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 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/proxy/admission.go b/internal/proxy/admission.go index 81dc316..24deae7 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,18 @@ 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). +// 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) + return id.OwnerID != "" || anyScopedLimitPresent(id) || id.GroupScopes != "" } // undeclaredOutputReservation is the output-token reservation for a request @@ -467,6 +472,146 @@ 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: 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;:,,,,;...). +// 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 +// 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) + } + 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 new file mode 100644 index 0000000..0830745 --- /dev/null +++ b/internal/proxy/group_scopes_test.go @@ -0,0 +1,704 @@ +package proxy + +import ( + "context" + "errors" + "io" + "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 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"), + []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,,,", "")}, + {"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")}, + {"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("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)+":"+maxRate+","+maxRate+","+maxRate+","+maxRate+","+maxSpend) + } + in := "v1;" + strings.Join(entries, ";") + 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 the maximal valid envelope (%d bytes) failed: %v", len(in), err) + } + }) +} + +// 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() + 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, eff) + 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 +} + +// 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) + } +} + +// 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) + } +} + +// 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) + } +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 9dc1dbf..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 @@ -603,6 +605,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 +625,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 +675,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 { @@ -690,14 +708,17 @@ 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()) { + // 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, ReservedOutputTokens: estimate.OutputTokens, Adapter: id.Adapter != "", OrganizationLimits: organizationLimits, OwnerLimits: ownerLimits, + GroupScopes: groupScopes, }) if err != nil { if errors.Is(err, admission.ErrInvalidIdentity) { @@ -1088,6 +1109,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 diff --git a/internal/rating/group_rollup_integration_test.go b/internal/rating/group_rollup_integration_test.go new file mode 100644 index 0000000..85fb850 --- /dev/null +++ b/internal/rating/group_rollup_integration_test.go @@ -0,0 +1,497 @@ +//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) + } +} + +// 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; +// - 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 +// 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 = "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; 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 + 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 + 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' + // (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: 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}, + {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}, + // (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',$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.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) + } + } + + // 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 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 { + 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 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 != 5 || res.EventsRated != 6 { + t.Fatalf("rollups/events = %d/%d, want 5/6 (the priced baseline only)", res.RollupsWritten, res.EventsRated) + } + 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 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 != 12 { + t.Fatalf("anomaly partition = %d, want 12 (all seeded events accounted exactly once)", got) + } + 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 { + 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(gConf); 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) + } + + // (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 — +// 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/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/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. diff --git a/internal/recovery/recovery.go b/internal/recovery/recovery.go index 7fc5d7d..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 @@ -274,7 +281,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++ @@ -335,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 073bba1..474224b 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" @@ -56,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) + } + }) } } @@ -175,8 +285,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) } } } diff --git a/migrations/0008_group_scopes.down.sql b/migrations/0008_group_scopes.down.sql new file mode 100644 index 0000000..06118b1 --- /dev/null +++ b/migrations/0008_group_scopes.down.sql @@ -0,0 +1,16 @@ +-- 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. +-- +-- 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; + +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..89e2cbe --- /dev/null +++ b/migrations/0008_group_scopes.up.sql @@ -0,0 +1,69 @@ +-- 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 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 +-- --------------------------------------------------------------------------- + +-- 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() 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, + 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..2bc04f9 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,38 @@ 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. 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`. +- 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 +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