diff --git a/server/db/queries/billing_rows.sql b/server/db/queries/billing_rows.sql index 1e8e28a..f49fef4 100644 --- a/server/db/queries/billing_rows.sql +++ b/server/db/queries/billing_rows.sql @@ -45,6 +45,48 @@ WHERE attributed_user_id = ? AND pioneer_created_at < ? ORDER BY pioneer_created_at ASC; +-- name: ListAllUsersLiveSpendToday :many +-- Live spend for ALL users since dayStart, grouped by user. +-- Used by RunSmartAllocation so it doesn't depend on the stale +-- user_daily_spend cache. +SELECT + b.attributed_user_id AS user_id, + COALESCE(SUM(CASE + WHEN pk.user_id = b.attributed_user_id AND pk.max_micros > pk.shared_micros + THEN b.cost_micros ELSE 0 + END), 0) AS private_spent_micros, + COALESCE(SUM(CASE + WHEN NOT (pk.user_id = b.attributed_user_id AND pk.max_micros > pk.shared_micros) + THEN b.cost_micros ELSE 0 + END), 0) AS shared_spent_micros +FROM pool_key_billing_rows b +JOIN pool_keys pk ON pk.id = b.pool_key_id +WHERE b.is_duplicate = 0 + AND b.pioneer_created_at >= ? + AND b.attributed_user_id IS NOT NULL +GROUP BY b.attributed_user_id; + +-- name: GetUserLiveSpendToday :one +-- Live spend for a user since the start of the current UTC day, read +-- directly from billing rows. Used by PoolGate for accurate pre-flight +-- checks without waiting for the reconciler to flush user_daily_spend. +-- Splits into shared vs private: private = rows where the user owns the +-- key and the key has a non-zero private reservation (max > shared). +SELECT + COALESCE(SUM(CASE + WHEN pk.user_id = b.attributed_user_id AND pk.max_micros > pk.shared_micros + THEN b.cost_micros ELSE 0 + END), 0) AS private_spent_micros, + COALESCE(SUM(CASE + WHEN NOT (pk.user_id = b.attributed_user_id AND pk.max_micros > pk.shared_micros) + THEN b.cost_micros ELSE 0 + END), 0) AS shared_spent_micros +FROM pool_key_billing_rows b +JOIN pool_keys pk ON pk.id = b.pool_key_id +WHERE b.attributed_user_id = ? + AND b.is_duplicate = 0 + AND b.pioneer_created_at >= ?; + -- name: UserHistoryProfile :many -- Per-user historical spend profile over a time window. -- Returns one row per user, with total non-duplicate spend and the count diff --git a/server/internal/api/middleware/middleware.go b/server/internal/api/middleware/middleware.go index aecce6d..88153ba 100644 --- a/server/internal/api/middleware/middleware.go +++ b/server/internal/api/middleware/middleware.go @@ -164,7 +164,10 @@ func PoolGate(q *store.Queries, poolAvailable func(ctx context.Context) bool, er } // 2. Per-user budget check. - day := time.Now().UTC().Unix() / 86400 + now := time.Now().UTC() + day := now.Unix() / 86400 + dayStart := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC).Unix() + allowRow, err := q.GetUserDailyAllowance(r.Context(), store.GetUserDailyAllowanceParams{ UserID: u.ID, Day: day, @@ -175,38 +178,43 @@ func PoolGate(q *store.Queries, poolAvailable func(ctx context.Context) bool, er return } - spendRow, err := q.GetUserDailySpend(r.Context(), store.GetUserDailySpendParams{ - UserID: u.ID, - Day: day, + // Read live spend directly from billing rows — accurate to within + // the last reconciler tick, not a stale cached aggregate. + liveSpend, err := q.GetUserLiveSpendToday(r.Context(), store.GetUserLiveSpendTodayParams{ + AttributedUserID: sql.NullString{String: u.ID, Valid: true}, + PioneerCreatedAt: dayStart, }) if err != nil { - // No spend yet — definitely allow. + // Can't read spend — allow rather than incorrectly block. next.ServeHTTP(w, r) return } - sharedRemaining := allowRow.SharedAllowanceMicros - spendRow.SharedSpentMicros - if sharedRemaining > 0 { - next.ServeHTTP(w, r) - return - } + sharedSpent := toI64(liveSpend.SharedSpentMicros) + privateSpent := toI64(liveSpend.PrivateSpentMicros) - // No shared remaining — check private reservation. - // Private = user's own keys' (max - shared) summed. - // For now we look this up from pool_keys directly. + // Compute private reservation from the user's own pool keys. + var privateReserved int64 privateRows, err := q.ListPoolKeysForUser(r.Context(), u.ID) if err == nil { - var privateReserved int64 for _, k := range privateRows { - if k.Active == 1 && k.RevokedAt.Valid == false { + if k.Active == 1 && !k.RevokedAt.Valid { privateReserved += k.MaxMicros - k.SharedMicros } } - privateRemaining := privateReserved - spendRow.PrivateSpentMicros - if privateRemaining > 0 { - next.ServeHTTP(w, r) - return - } + } + + privateRemaining := privateReserved - privateSpent + // Private overspend counts against the shared allowance too. + privateOverflow := int64(0) + if privateSpent > privateReserved { + privateOverflow = privateSpent - privateReserved + } + sharedRemaining := allowRow.SharedAllowanceMicros - sharedSpent - privateOverflow + + if sharedRemaining > 0 || privateRemaining > 0 { + next.ServeHTTP(w, r) + return } errResp(w, http.StatusPaymentRequired, "insufficient_funds", @@ -218,6 +226,19 @@ func PoolGate(q *store.Queries, poolAvailable func(ctx context.Context) bool, er // ---- per-user rate limiter -------------------------------------------- +// toI64 unboxes the interface{} sqlc emits for aggregate columns. +func toI64(v any) int64 { + switch x := v.(type) { + case int64: + return x + case int: + return int64(x) + case float64: + return int64(x) + } + return 0 +} + // RateLimit is a per-user token bucket. The map is bounded by the active // user set — keys are GC'd by the periodic sweep. func RateLimit(rps float64, burst int, errResp ErrorResponder) func(http.Handler) http.Handler { diff --git a/server/internal/api/web/allocations.go b/server/internal/api/web/allocations.go index de0114b..2fa3611 100644 --- a/server/internal/api/web/allocations.go +++ b/server/internal/api/web/allocations.go @@ -49,13 +49,15 @@ func (s *Server) RunSmartAllocation(ctx context.Context, setByUserID string) err return fmt.Errorf("list pool allocations: %w", err) } - spendRows, err := s.Q.ListUserDailySpendForDay(ctx, day) + spendRows, err := s.Q.ListAllUsersLiveSpendToday(ctx, dayStart.Unix()) if err != nil { - return fmt.Errorf("list daily spend: %w", err) + return fmt.Errorf("list live spend: %w", err) } spentByUser := map[string]int64{} for _, sp := range spendRows { - spentByUser[sp.UserID] = sp.SharedSpentMicros + if sp.UserID.Valid { + spentByUser[sp.UserID.String] = toInt64(sp.SharedSpentMicros) + } } // Historical profile per user. @@ -104,6 +106,13 @@ func (s *Server) RunSmartAllocation(ctx context.Context, setByUserID string) err allocations := smartAllocate(totalShared, spends, histories, dayFraction) + // "system" is not a real user ID — resolve to the first user in the pool + // so the FK constraint on set_by_user_id is satisfied. + resolvedBy := setByUserID + if resolvedBy == "system" && len(userIDs) > 0 { + resolvedBy = userIDs[0] + } + nowUnix := now.Unix() for i, id := range userIDs { a := allocations[i] @@ -116,7 +125,7 @@ func (s *Server) RunSmartAllocation(ctx context.Context, setByUserID string) err PredictedTotalMicros: a.PredictedTotal, HistoryDaysUsed: a.HistoryDays, SetAt: nowUnix, - SetByUserID: setByUserID, + SetByUserID: resolvedBy, }); err != nil { return fmt.Errorf("upsert allowance for %s: %w", id, err) } @@ -194,15 +203,22 @@ func fairShareAllocate(pool int64, spends []int64) []int64 { // buildAllocations computes the full allocations payload. func (s *Server) buildAllocations(r *http.Request) map[string]any { - day := time.Now().UTC().Unix() / 86400 + nowUtc := time.Now().UTC() + day := nowUtc.Unix() / 86400 + dayStartUnix := time.Date(nowUtc.Year(), nowUtc.Month(), nowUtc.Day(), 0, 0, 0, 0, time.UTC).Unix() rows, _ := s.Q.ListPoolAllocations(r.Context()) - spendRows, _ := s.Q.ListUserDailySpendForDay(r.Context(), day) + liveSpendRows, _ := s.Q.ListAllUsersLiveSpendToday(r.Context(), dayStartUnix) allowRows, _ := s.Q.ListUserDailyAllowancesForDay(r.Context(), day) spendByUser := map[string]struct{ shared, private int64 }{} - for _, sp := range spendRows { - spendByUser[sp.UserID] = struct{ shared, private int64 }{sp.SharedSpentMicros, sp.PrivateSpentMicros} + for _, sp := range liveSpendRows { + if sp.UserID.Valid { + spendByUser[sp.UserID.String] = struct{ shared, private int64 }{ + toInt64(sp.SharedSpentMicros), + toInt64(sp.PrivateSpentMicros), + } + } } allowByUser := map[string]store.UserDailyAllowance{} for _, a := range allowRows { diff --git a/server/internal/pool/billing_ingest.go b/server/internal/pool/billing_ingest.go index 22d50f3..88284d8 100644 --- a/server/internal/pool/billing_ingest.go +++ b/server/internal/pool/billing_ingest.go @@ -253,12 +253,40 @@ func (r *Reconciler) ingestBillingRows(ctx context.Context, key store.PoolKey, p }) // Accumulate spend for today only, non-duplicate rows. - // Only goes to private if the user owns the key AND the key has - // a non-zero private reservation (max > shared). + // Route to private bucket only up to the key's private reservation + // (max - shared). Any overflow spills into shared. if isDup == 0 && attrUserID.Valid && ts >= dayStart { - ownsKeyAndHasPrivate := attrUserID.String == key.UserID && - (key.MaxMicros-key.SharedMicros) > 0 - addSpend(attrUserID.String, costMicros, ownsKeyAndHasPrivate) + privateReservation := key.MaxMicros - key.SharedMicros + ownsKeyAndHasPrivate := attrUserID.String == key.UserID && privateReservation > 0 + if ownsKeyAndHasPrivate { + // Cap private at reservation; spill remainder to shared. + // How much private spend has already been attributed in this + // batch for this user (earlier rows in the same reconciler tick). + alreadyPrivate := int64(0) + if d := spendDeltas[attrUserID.String]; d != nil { + alreadyPrivate = d.private + } + // NOTE: this only caps within a single reconciler run. + // Across ticks, the gate uses GetUserLiveSpendToday which + // reads already-written rows and applies the same cap logic. + canPrivate := privateReservation - alreadyPrivate + if canPrivate < 0 { + canPrivate = 0 + } + privateAmount := costMicros + if privateAmount > canPrivate { + privateAmount = canPrivate + } + spillAmount := costMicros - privateAmount + if privateAmount > 0 { + addSpend(attrUserID.String, privateAmount, true) + } + if spillAmount > 0 { + addSpend(attrUserID.String, spillAmount, false) + } + } else { + addSpend(attrUserID.String, costMicros, false) + } } } diff --git a/server/internal/store/billing_rows.sql.go b/server/internal/store/billing_rows.sql.go index 7ae2bff..5830282 100644 --- a/server/internal/store/billing_rows.sql.go +++ b/server/internal/store/billing_rows.sql.go @@ -10,6 +10,45 @@ import ( "database/sql" ) +const getUserLiveSpendToday = `-- name: GetUserLiveSpendToday :one +SELECT + COALESCE(SUM(CASE + WHEN pk.user_id = b.attributed_user_id AND pk.max_micros > pk.shared_micros + THEN b.cost_micros ELSE 0 + END), 0) AS private_spent_micros, + COALESCE(SUM(CASE + WHEN NOT (pk.user_id = b.attributed_user_id AND pk.max_micros > pk.shared_micros) + THEN b.cost_micros ELSE 0 + END), 0) AS shared_spent_micros +FROM pool_key_billing_rows b +JOIN pool_keys pk ON pk.id = b.pool_key_id +WHERE b.attributed_user_id = ? + AND b.is_duplicate = 0 + AND b.pioneer_created_at >= ? +` + +type GetUserLiveSpendTodayParams struct { + AttributedUserID sql.NullString `json:"attributed_user_id"` + PioneerCreatedAt int64 `json:"pioneer_created_at"` +} + +type GetUserLiveSpendTodayRow struct { + PrivateSpentMicros interface{} `json:"private_spent_micros"` + SharedSpentMicros interface{} `json:"shared_spent_micros"` +} + +// Live spend for a user since the start of the current UTC day, read +// directly from billing rows. Used by PoolGate for accurate pre-flight +// checks without waiting for the reconciler to flush user_daily_spend. +// Splits into shared vs private: private = rows where the user owns the +// key and the key has a non-zero private reservation (max > shared). +func (q *Queries) GetUserLiveSpendToday(ctx context.Context, arg GetUserLiveSpendTodayParams) (GetUserLiveSpendTodayRow, error) { + row := q.db.QueryRowContext(ctx, getUserLiveSpendToday, arg.AttributedUserID, arg.PioneerCreatedAt) + var i GetUserLiveSpendTodayRow + err := row.Scan(&i.PrivateSpentMicros, &i.SharedSpentMicros) + return i, err +} + const latestBillingRowTime = `-- name: LatestBillingRowTime :one SELECT COALESCE(SUM(pioneer_created_at * 0) + MAX(pioneer_created_at), 0) FROM pool_key_billing_rows @@ -26,6 +65,57 @@ func (q *Queries) LatestBillingRowTime(ctx context.Context, poolKeyID string) (i return coalesce, err } +const listAllUsersLiveSpendToday = `-- name: ListAllUsersLiveSpendToday :many +SELECT + b.attributed_user_id AS user_id, + COALESCE(SUM(CASE + WHEN pk.user_id = b.attributed_user_id AND pk.max_micros > pk.shared_micros + THEN b.cost_micros ELSE 0 + END), 0) AS private_spent_micros, + COALESCE(SUM(CASE + WHEN NOT (pk.user_id = b.attributed_user_id AND pk.max_micros > pk.shared_micros) + THEN b.cost_micros ELSE 0 + END), 0) AS shared_spent_micros +FROM pool_key_billing_rows b +JOIN pool_keys pk ON pk.id = b.pool_key_id +WHERE b.is_duplicate = 0 + AND b.pioneer_created_at >= ? + AND b.attributed_user_id IS NOT NULL +GROUP BY b.attributed_user_id +` + +type ListAllUsersLiveSpendTodayRow struct { + UserID sql.NullString `json:"user_id"` + PrivateSpentMicros interface{} `json:"private_spent_micros"` + SharedSpentMicros interface{} `json:"shared_spent_micros"` +} + +// Live spend for ALL users since dayStart, grouped by user. +// Used by RunSmartAllocation so it doesn't depend on the stale +// user_daily_spend cache. +func (q *Queries) ListAllUsersLiveSpendToday(ctx context.Context, pioneerCreatedAt int64) ([]ListAllUsersLiveSpendTodayRow, error) { + rows, err := q.db.QueryContext(ctx, listAllUsersLiveSpendToday, pioneerCreatedAt) + if err != nil { + return nil, err + } + defer rows.Close() + items := []ListAllUsersLiveSpendTodayRow{} + for rows.Next() { + var i ListAllUsersLiveSpendTodayRow + if err := rows.Scan(&i.UserID, &i.PrivateSpentMicros, &i.SharedSpentMicros); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const listBillingRowsForKeyAfter = `-- name: ListBillingRowsForKeyAfter :many SELECT id, pool_key_id, pioneer_created_at, credit_micros, cost_micros, token_usage, model, endpoint, attributed_user_id, attribution, is_duplicate, matched_request_id, ingested_at FROM pool_key_billing_rows WHERE pool_key_id = ? diff --git a/server/internal/store/querier.go b/server/internal/store/querier.go index 10edd41..1dcbf37 100644 --- a/server/internal/store/querier.go +++ b/server/internal/store/querier.go @@ -58,11 +58,21 @@ type Querier interface { GetUserByID(ctx context.Context, id string) (User, error) GetUserDailyAllowance(ctx context.Context, arg GetUserDailyAllowanceParams) (UserDailyAllowance, error) GetUserDailySpend(ctx context.Context, arg GetUserDailySpendParams) (UserDailySpend, error) + // Live spend for a user since the start of the current UTC day, read + // directly from billing rows. Used by PoolGate for accurate pre-flight + // checks without waiting for the reconciler to flush user_daily_spend. + // Splits into shared vs private: private = rows where the user owns the + // key and the key has a non-zero private reservation (max > shared). + GetUserLiveSpendToday(ctx context.Context, arg GetUserLiveSpendTodayParams) (GetUserLiveSpendTodayRow, error) // Most recent pioneer_created_at for a key. // COALESCE(SUM(pioneer_created_at*0) + MAX(...)) is a workaround for // sqlc's sqlite parser rejecting bare COALESCE(MAX(...), 0). LatestBillingRowTime(ctx context.Context, poolKeyID string) (interface{}, error) ListAPIKeysForUser(ctx context.Context, userID string) ([]ApiKey, error) + // Live spend for ALL users since dayStart, grouped by user. + // Used by RunSmartAllocation so it doesn't depend on the stale + // user_daily_spend cache. + ListAllUsersLiveSpendToday(ctx context.Context, pioneerCreatedAt int64) ([]ListAllUsersLiveSpendTodayRow, error) // All billing rows for a key after a given timestamp. ListBillingRowsForKeyAfter(ctx context.Context, arg ListBillingRowsForKeyAfterParams) ([]PoolKeyBillingRow, error) // Billing rows attributed to a user in a time window. diff --git a/web/src/routes/dashboard/+page.svelte b/web/src/routes/dashboard/+page.svelte index 915185d..017360a 100644 --- a/web/src/routes/dashboard/+page.svelte +++ b/web/src/routes/dashboard/+page.svelte @@ -161,6 +161,7 @@ {@const _overW = (Math.max(0, _sharedUsed - _sharedCap) / _total) * 100} {@const _privateUsedW = (Math.min(_privateUsed, _private) / _total) * 100} {@const _privateRemW = (Math.max(0, _private - _privateUsed) / _total) * 100} + {@const _privateOverW = (Math.max(0, _privateUsed - _private) / _total) * 100}