fix(quota): spend overage credits against the weekly window and count 7 days, not 8
This commit is contained in:
parent
f4f9653361
commit
102f3ef934
5 changed files with 601 additions and 65 deletions
84
server/supabase/functions/_shared/quota-policy.ts
Normal file
84
server/supabase/functions/_shared/quota-policy.ts
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
// server/supabase/functions/_shared/quota-policy.ts
|
||||
// Pure quota policy shared by the quota adapter (quota.ts) and its tests.
|
||||
//
|
||||
// Policy (pure): window arithmetic, allow/deny decision, RPC response parsing.
|
||||
// Adapter: quota.ts (subscriptions + daily_usage + consume_quota RPC).
|
||||
//
|
||||
// The weekly window is the rolling 7 calendar days ending today
|
||||
// (today-6 .. today, UTC). The SQL quota functions (`consume_quota`,
|
||||
// `reserve_stt_quota`, the meeting-document RPC) use `CURRENT_DATE - 6`;
|
||||
// this module must stay on the same window so a pre-check never blocks usage
|
||||
// that the atomic SQL consumption has already rolled off.
|
||||
|
||||
import type { PlanQuotaPeriod } from './core-contract.generated.ts'
|
||||
|
||||
/** Days before today that still count toward a weekly quota. */
|
||||
export const WEEKLY_WINDOW_EXTRA_DAYS = 6
|
||||
|
||||
export type QuotaConsumedFrom = 'base' | 'overage' | 'unlimited' | 'none'
|
||||
|
||||
export interface QuotaConsumeResponse {
|
||||
allowed: boolean
|
||||
current: number
|
||||
limit: number
|
||||
overageCredits: number
|
||||
consumedFrom: QuotaConsumedFrom
|
||||
}
|
||||
|
||||
function isoDate(date: Date): string {
|
||||
return date.toISOString().slice(0, 10)
|
||||
}
|
||||
|
||||
/** Today's usage date (UTC, matches the database `CURRENT_DATE` on Supabase). */
|
||||
export function quotaUsageDate(now: Date): string {
|
||||
return isoDate(now)
|
||||
}
|
||||
|
||||
/**
|
||||
* First usage date (inclusive) that counts toward the period.
|
||||
* daily → today; weekly → today-6 (7 calendar days including today).
|
||||
*/
|
||||
export function quotaWindowStart(period: PlanQuotaPeriod, now: Date): string {
|
||||
if (period !== 'weekly') return isoDate(now)
|
||||
const start = new Date(now.getTime())
|
||||
start.setUTCDate(start.getUTCDate() - WEEKLY_WINDOW_EXTRA_DAYS)
|
||||
return isoDate(start)
|
||||
}
|
||||
|
||||
/**
|
||||
* Whether one more call may go through.
|
||||
* limit: -1 = unlimited, 0 = not available, >0 = base allowance in the window.
|
||||
* Past the base allowance a positive overage balance lets the call through;
|
||||
* the atomic `consume_quota` RPC then spends one credit for it.
|
||||
*/
|
||||
export function isQuotaAllowed(limit: number, current: number, overageCredits: number): boolean {
|
||||
if (limit === 0) return false
|
||||
if (limit === -1) return true
|
||||
return current < limit || overageCredits > 0
|
||||
}
|
||||
|
||||
const CONSUMED_FROM: readonly QuotaConsumedFrom[] = ['base', 'overage', 'unlimited', 'none']
|
||||
|
||||
/** Validate the `consume_quota` RPC payload; throws on any unexpected shape. */
|
||||
export function parseQuotaConsumeResponse(data: unknown): QuotaConsumeResponse {
|
||||
if (!data || typeof data !== 'object' || Array.isArray(data)) {
|
||||
throw new Error('Invalid quota consumption response.')
|
||||
}
|
||||
const result = data as Record<string, unknown>
|
||||
if (
|
||||
typeof result.allowed !== 'boolean'
|
||||
|| typeof result.current !== 'number'
|
||||
|| typeof result.limit !== 'number'
|
||||
|| typeof result.overage_credits !== 'number'
|
||||
|| !(CONSUMED_FROM as readonly unknown[]).includes(result.consumed_from)
|
||||
) {
|
||||
throw new Error('Invalid quota consumption response.')
|
||||
}
|
||||
return {
|
||||
allowed: result.allowed,
|
||||
current: result.current,
|
||||
limit: result.limit,
|
||||
overageCredits: result.overage_credits,
|
||||
consumedFrom: result.consumed_from as QuotaConsumedFrom,
|
||||
}
|
||||
}
|
||||
151
server/supabase/functions/_shared/quota.test.ts
Normal file
151
server/supabase/functions/_shared/quota.test.ts
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
import type { createClient } from '@supabase/supabase-js'
|
||||
import { checkQuota, consumeQuota } from './quota.ts'
|
||||
import { isQuotaAllowed, parseQuotaConsumeResponse, quotaWindowStart } from './quota-policy.ts'
|
||||
|
||||
function assert(condition: boolean, message: string): asserts condition {
|
||||
if (!condition) throw new Error(message)
|
||||
}
|
||||
|
||||
type Row = Record<string, unknown>
|
||||
type ServiceClient = ReturnType<typeof createClient>
|
||||
|
||||
/** Minimal in-memory stand-in for the supabase-js query builder used by quota.ts. */
|
||||
class FakeQuery implements PromiseLike<{ data: unknown; error: null }> {
|
||||
private filters: Array<(row: Row) => boolean> = []
|
||||
private singleRow = false
|
||||
constructor(private rows: Row[]) {}
|
||||
select(_columns: string): this {
|
||||
return this
|
||||
}
|
||||
eq(column: string, value: unknown): this {
|
||||
this.filters.push((row) => row[column] === value)
|
||||
return this
|
||||
}
|
||||
gte(column: string, value: string): this {
|
||||
this.filters.push((row) => String(row[column]) >= value)
|
||||
return this
|
||||
}
|
||||
lte(column: string, value: string): this {
|
||||
this.filters.push((row) => String(row[column]) <= value)
|
||||
return this
|
||||
}
|
||||
single(): this {
|
||||
this.singleRow = true
|
||||
return this
|
||||
}
|
||||
then<T1 = { data: unknown; error: null }, T2 = never>(
|
||||
onfulfilled?: ((value: { data: unknown; error: null }) => T1 | PromiseLike<T1>) | null,
|
||||
onrejected?: ((reason: unknown) => T2 | PromiseLike<T2>) | null,
|
||||
): PromiseLike<T1 | T2> {
|
||||
const matched = this.rows.filter((row) => this.filters.every((f) => f(row)))
|
||||
const data = this.singleRow ? (matched[0] ?? null) : matched
|
||||
return Promise.resolve({ data, error: null }).then(onfulfilled, onrejected)
|
||||
}
|
||||
}
|
||||
|
||||
class FakeQuotaClient {
|
||||
tables: Record<string, Row[]> = { subscriptions: [], daily_usage: [] }
|
||||
rpcCalls: Array<{ name: string; args: Record<string, unknown> }> = []
|
||||
rpcResult: unknown = {
|
||||
allowed: true,
|
||||
current: 1,
|
||||
limit: 250,
|
||||
overage_credits: 0,
|
||||
consumed_from: 'overage',
|
||||
}
|
||||
from(table: string): FakeQuery {
|
||||
return new FakeQuery(this.tables[table] ?? [])
|
||||
}
|
||||
rpc(name: string, args: Record<string, unknown>): Promise<{ data: unknown; error: null }> {
|
||||
this.rpcCalls.push({ name, args })
|
||||
return Promise.resolve({ data: this.rpcResult, error: null })
|
||||
}
|
||||
asClient(): ServiceClient {
|
||||
return this as unknown as ServiceClient
|
||||
}
|
||||
}
|
||||
|
||||
const NOW = new Date('2026-09-28T12:00:00Z')
|
||||
const USER = 'user-1'
|
||||
|
||||
function freeUserWithUsage(rows: Array<[string, number]>, overageCredits = 0): FakeQuotaClient {
|
||||
const client = new FakeQuotaClient()
|
||||
client.tables.subscriptions.push({ user_id: USER, tier: 'free', overage_credits: overageCredits })
|
||||
for (const [date, count] of rows) {
|
||||
client.tables.daily_usage.push({ user_id: USER, feature: 'llm_haiku', date, count })
|
||||
}
|
||||
return client
|
||||
}
|
||||
|
||||
Deno.test('weekly window is today-6..today (7 calendar days), same as the SQL quota functions', () => {
|
||||
assert(quotaWindowStart('weekly', NOW) === '2026-09-22', `weekly start ${quotaWindowStart('weekly', NOW)}`)
|
||||
assert(quotaWindowStart('daily', NOW) === '2026-09-28', 'daily start is today')
|
||||
// Month boundary in UTC.
|
||||
assert(quotaWindowStart('weekly', new Date('2026-10-03T00:30:00Z')) === '2026-09-27', 'month rollover')
|
||||
})
|
||||
|
||||
Deno.test('usage from 7 days ago no longer blocks a free user (8-day window regression)', async () => {
|
||||
const client = freeUserWithUsage([['2026-09-21', 250], ['2026-09-27', 10]])
|
||||
const check = await checkQuota(USER, 'llm_haiku', client.asClient(), NOW)
|
||||
assert(check.current === 10, `stale usage counted: current=${check.current}`)
|
||||
assert(check.allowed, 'free user blocked by usage that already rolled off')
|
||||
assert(check.period === 'weekly' && check.limit === 250, 'free haiku policy')
|
||||
})
|
||||
|
||||
Deno.test('250 uses spread across 6 days reach the weekly limit; a credit keeps the call allowed', async () => {
|
||||
const days: Array<[string, number]> = [
|
||||
['2026-09-22', 45], ['2026-09-23', 41], ['2026-09-24', 41],
|
||||
['2026-09-25', 41], ['2026-09-26', 41], ['2026-09-27', 41],
|
||||
]
|
||||
const noCredit = await checkQuota(USER, 'llm_haiku', freeUserWithUsage(days).asClient(), NOW)
|
||||
assert(noCredit.current === 250 && !noCredit.allowed, `limit not enforced: ${JSON.stringify(noCredit)}`)
|
||||
const withCredit = await checkQuota(USER, 'llm_haiku', freeUserWithUsage(days, 1).asClient(), NOW)
|
||||
assert(withCredit.allowed && withCredit.overageCredits === 1, 'credit should allow the call')
|
||||
})
|
||||
|
||||
Deno.test('consumeQuota asks the RPC to consume over the policy period', async () => {
|
||||
const client = freeUserWithUsage([])
|
||||
const result = await consumeQuota(USER, 'llm_haiku', client.asClient(), 250)
|
||||
const call = client.rpcCalls[0]
|
||||
assert(call?.name === 'consume_quota', 'consume_quota not called')
|
||||
assert(call.args.p_period === 'weekly', `free haiku must consume weekly, got ${String(call.args.p_period)}`)
|
||||
assert(call.args.p_base_limit === 250 && call.args.p_feature === 'llm_haiku', 'limit/feature passed through')
|
||||
assert(result.consumedFrom === 'overage' && result.overageCredits === 0, 'response mapped')
|
||||
|
||||
const explicit = new FakeQuotaClient()
|
||||
await consumeQuota(USER, 'realtime_session', explicit.asClient(), 30, 'daily')
|
||||
assert(explicit.rpcCalls[0]?.args.p_period === 'daily', 'explicit period passed as-is')
|
||||
})
|
||||
|
||||
Deno.test('paid daily policies consume over a daily window', async () => {
|
||||
const client = new FakeQuotaClient()
|
||||
client.tables.subscriptions.push({ user_id: USER, tier: 'pro', overage_credits: 0 })
|
||||
await consumeQuota(USER, 'llm_sonnet', client.asClient(), 300)
|
||||
assert(client.rpcCalls[0]?.args.p_period === 'daily', 'pro sonnet is daily')
|
||||
})
|
||||
|
||||
Deno.test('allow decision: 0 = unavailable even with credits, -1 = unlimited, credits extend the base', () => {
|
||||
assert(!isQuotaAllowed(0, 0, 5), 'unavailable feature allowed by credits')
|
||||
assert(isQuotaAllowed(-1, 10_000, 0), 'unlimited denied')
|
||||
assert(isQuotaAllowed(250, 249, 0), 'under limit denied')
|
||||
assert(!isQuotaAllowed(250, 250, 0), 'at limit without credit allowed')
|
||||
assert(isQuotaAllowed(250, 250, 1), 'at limit with credit denied')
|
||||
})
|
||||
|
||||
Deno.test('RPC response parser rejects malformed payloads', () => {
|
||||
const ok = parseQuotaConsumeResponse({
|
||||
allowed: false, current: 250, limit: 250, overage_credits: 0, consumed_from: 'none',
|
||||
})
|
||||
assert(!ok.allowed && ok.consumedFrom === 'none' && ok.current === 250, 'valid payload mapped')
|
||||
for (const bad of [null, [], { allowed: true }, {
|
||||
allowed: true, current: 1, limit: 1, overage_credits: 0, consumed_from: 'bogus',
|
||||
}]) {
|
||||
let threw = false
|
||||
try {
|
||||
parseQuotaConsumeResponse(bad)
|
||||
} catch {
|
||||
threw = true
|
||||
}
|
||||
assert(threw, `accepted ${JSON.stringify(bad)}`)
|
||||
}
|
||||
})
|
||||
|
|
@ -10,6 +10,13 @@ import {
|
|||
type PlanQuotaPeriod,
|
||||
type PlanQuotaTier,
|
||||
} from './core-contract.generated.ts'
|
||||
import {
|
||||
isQuotaAllowed,
|
||||
parseQuotaConsumeResponse,
|
||||
type QuotaConsumedFrom,
|
||||
quotaUsageDate,
|
||||
quotaWindowStart,
|
||||
} from './quota-policy.ts'
|
||||
|
||||
export type Tier = PlanQuotaTier
|
||||
|
||||
|
|
@ -50,7 +57,7 @@ export interface QuotaConsumeResult {
|
|||
current: number
|
||||
limit: number
|
||||
overageCredits: number
|
||||
consumedFrom: 'base' | 'overage' | 'unlimited' | 'none'
|
||||
consumedFrom: QuotaConsumedFrom
|
||||
}
|
||||
|
||||
export interface SttQuotaReservation {
|
||||
|
|
@ -123,65 +130,75 @@ export async function finalizeSttQuota(
|
|||
return status
|
||||
}
|
||||
|
||||
/**
|
||||
* 쿼터 확인 — 모델별, 기간별(daily/weekly).
|
||||
* weekly인 경우 최근 7일 daily_usage를 합산.
|
||||
*/
|
||||
export async function checkQuota(
|
||||
interface SubscriptionQuotaState {
|
||||
tier: Tier
|
||||
overageCredits: number
|
||||
}
|
||||
|
||||
/** 구독 티어 + overage 잔액. 행이 없으면 free/0 (기존 checkQuota 동작 유지). */
|
||||
async function readSubscriptionQuotaState(
|
||||
userId: string,
|
||||
feature: QuotaFeature,
|
||||
serviceRoleClient: ReturnType<typeof createClient>,
|
||||
): Promise<QuotaCheck> {
|
||||
// 티어 + overage 조회
|
||||
): Promise<SubscriptionQuotaState> {
|
||||
const { data: sub } = await serviceRoleClient
|
||||
.from('subscriptions')
|
||||
.select('tier, overage_credits')
|
||||
.eq('user_id', userId)
|
||||
.single()
|
||||
|
||||
const tier: Tier = (sub?.tier as Tier) ?? 'free'
|
||||
const overageCredits = (sub?.overage_credits as number) ?? 0
|
||||
const policy = getQuotaPolicy(tier, feature)
|
||||
|
||||
// 사용불가 (limit=0)
|
||||
if (policy.limit === 0) {
|
||||
return { allowed: false, current: 0, limit: 0, period: policy.period, tier, overageCredits }
|
||||
return {
|
||||
tier: (sub?.tier as Tier) ?? 'free',
|
||||
overageCredits: (sub?.overage_credits as number) ?? 0,
|
||||
}
|
||||
}
|
||||
|
||||
// 무제한
|
||||
if (policy.limit === -1) {
|
||||
return { allowed: true, current: 0, limit: -1, period: policy.period, tier, overageCredits }
|
||||
}
|
||||
|
||||
// 사용량 조회 (daily vs weekly)
|
||||
let current: number
|
||||
if (policy.period === 'weekly') {
|
||||
// 최근 7일 합산
|
||||
const weekAgo = new Date()
|
||||
weekAgo.setDate(weekAgo.getDate() - 7)
|
||||
/** 기간 창(daily=오늘, weekly=오늘-6..오늘) 안의 daily_usage 합계. */
|
||||
async function readWindowUsage(
|
||||
userId: string,
|
||||
feature: QuotaFeature,
|
||||
period: QuotaPeriod,
|
||||
serviceRoleClient: ReturnType<typeof createClient>,
|
||||
now: Date,
|
||||
): Promise<number> {
|
||||
const { data: rows } = await serviceRoleClient
|
||||
.from('daily_usage')
|
||||
.select('count')
|
||||
.eq('user_id', userId)
|
||||
.eq('feature', feature)
|
||||
.gte('date', weekAgo.toISOString().slice(0, 10))
|
||||
.gte('date', quotaWindowStart(period, now))
|
||||
.lte('date', quotaUsageDate(now))
|
||||
const usageRows = (rows ?? []) as Array<{ count: number | null }>
|
||||
current = usageRows.reduce((sum, row) => sum + (row.count ?? 0), 0)
|
||||
} else {
|
||||
// 오늘만
|
||||
const today = new Date().toISOString().slice(0, 10)
|
||||
const { data: usage } = await serviceRoleClient
|
||||
.from('daily_usage')
|
||||
.select('count')
|
||||
.eq('user_id', userId)
|
||||
.eq('date', today)
|
||||
.eq('feature', feature)
|
||||
.maybeSingle()
|
||||
current = (usage?.count as number) ?? 0
|
||||
return usageRows.reduce((sum, row) => sum + (row.count ?? 0), 0)
|
||||
}
|
||||
|
||||
/**
|
||||
* 쿼터 확인 — 모델별, 기간별(daily/weekly).
|
||||
* weekly인 경우 최근 7일(오늘 포함, 오늘-6..오늘) daily_usage를 합산 —
|
||||
* SQL 쪽 `consume_quota`·`reserve_stt_quota`의 `CURRENT_DATE - 6` 창과 같다.
|
||||
*/
|
||||
export async function checkQuota(
|
||||
userId: string,
|
||||
feature: QuotaFeature,
|
||||
serviceRoleClient: ReturnType<typeof createClient>,
|
||||
now: Date = new Date(),
|
||||
): Promise<QuotaCheck> {
|
||||
const { tier, overageCredits } = await readSubscriptionQuotaState(userId, serviceRoleClient)
|
||||
const policy = getQuotaPolicy(tier, feature)
|
||||
|
||||
// 사용불가 (limit=0) / 무제한 (limit=-1) 은 사용량 조회가 필요 없다.
|
||||
if (policy.limit === 0 || policy.limit === -1) {
|
||||
return {
|
||||
allowed: current < policy.limit || overageCredits > 0,
|
||||
allowed: isQuotaAllowed(policy.limit, 0, overageCredits),
|
||||
current: 0,
|
||||
limit: policy.limit,
|
||||
period: policy.period,
|
||||
tier,
|
||||
overageCredits,
|
||||
}
|
||||
}
|
||||
|
||||
const current = await readWindowUsage(userId, feature, policy.period, serviceRoleClient, now)
|
||||
return {
|
||||
allowed: isQuotaAllowed(policy.limit, current, overageCredits),
|
||||
current,
|
||||
limit: policy.limit,
|
||||
period: policy.period,
|
||||
|
|
@ -191,42 +208,36 @@ export async function checkQuota(
|
|||
}
|
||||
|
||||
/**
|
||||
* 쿼터 소비 — 항상 오늘 날짜의 daily_usage를 +1 증가.
|
||||
* (weekly 집계는 checkQuota에서 7일 합산으로 처리)
|
||||
* 쿼터 소비 — 오늘 날짜의 daily_usage를 +1 증가.
|
||||
* `consume_quota` RPC가 기간 창(weekly=최근 7일 합산) 사용량을 base 한도와 비교해
|
||||
* base 소진 + overage 있으면 overage 1 크레딧을 원자적으로 차감한다.
|
||||
* 무제한(-1)이면 카운터만 증가하고 allowed=true.
|
||||
* base 소진 + overage 있으면 overage 차감.
|
||||
*
|
||||
* `period`를 넘기지 않으면 구독 티어로 정책 기간을 다시 조회한다
|
||||
* (checkQuota 결과의 `period`를 넘기면 조회 1회를 아낀다).
|
||||
*/
|
||||
export async function consumeQuota(
|
||||
userId: string,
|
||||
feature: QuotaFeature,
|
||||
serviceRoleClient: ReturnType<typeof createClient>,
|
||||
baseLimit: number,
|
||||
period?: QuotaPeriod,
|
||||
): Promise<QuotaConsumeResult> {
|
||||
const resolvedPeriod = period
|
||||
?? getQuotaPolicy((await readSubscriptionQuotaState(userId, serviceRoleClient)).tier, feature).period
|
||||
|
||||
const { data, error } = await serviceRoleClient.rpc('consume_quota', {
|
||||
p_user_id: userId,
|
||||
p_feature: feature,
|
||||
p_base_limit: baseLimit,
|
||||
p_period: resolvedPeriod,
|
||||
})
|
||||
|
||||
if (error) {
|
||||
throw new Error(`Failed to consume quota: ${error.message}`)
|
||||
}
|
||||
|
||||
const result = data as {
|
||||
allowed: boolean
|
||||
current: number
|
||||
limit: number
|
||||
overage_credits: number
|
||||
consumed_from: 'base' | 'overage' | 'unlimited' | 'none'
|
||||
}
|
||||
|
||||
return {
|
||||
allowed: result.allowed,
|
||||
current: result.current,
|
||||
limit: result.limit,
|
||||
overageCredits: result.overage_credits,
|
||||
consumedFrom: result.consumed_from,
|
||||
}
|
||||
return parseQuotaConsumeResponse(data)
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -0,0 +1,132 @@
|
|||
-- ============================================================================
|
||||
-- consume_quota: period-aware base/overage consumption
|
||||
-- ============================================================================
|
||||
-- The original consume_quota (20260412000001) compared only TODAY's
|
||||
-- daily_usage row with the base limit. For weekly policies (free tier:
|
||||
-- llm_haiku 250/week) that meant a user whose 7-day total had already reached
|
||||
-- the limit kept taking the 'base' branch as long as today's own count stayed
|
||||
-- under 250, so an overage credit (e.g. one rewarded ad) let every call through
|
||||
-- without ever being deducted.
|
||||
--
|
||||
-- This version takes the policy period and sums the same window the other SQL
|
||||
-- quota functions use (reserve_stt_quota, meeting-document generation):
|
||||
-- weekly -> CURRENT_DATE - 6 .. CURRENT_DATE (7 calendar days)
|
||||
-- daily -> CURRENT_DATE
|
||||
-- and chooses base vs overage from that window total, under a per-user/feature
|
||||
-- advisory lock plus the subscription row lock so concurrent calls cannot both
|
||||
-- take the last base unit or the last credit.
|
||||
--
|
||||
-- Additional corrections:
|
||||
-- * limit 0 ("not available") is denied instead of spending overage credits.
|
||||
-- * the overage branch only succeeds when a credit is actually left
|
||||
-- (`overage_credits > 0` in the UPDATE), so a race can never go negative.
|
||||
-- * 'current' is the window total after this call (equal to today's count
|
||||
-- for daily policies), matching what checkQuota reports.
|
||||
--
|
||||
-- p_period defaults to 'daily' so an edge function still calling the old
|
||||
-- three-argument form keeps its previous (daily) behaviour until redeployed.
|
||||
-- Apply this migration before deploying the quota.ts that passes p_period.
|
||||
|
||||
DROP FUNCTION IF EXISTS public.consume_quota(uuid, text, integer);
|
||||
|
||||
CREATE OR REPLACE FUNCTION public.consume_quota(
|
||||
p_user_id uuid,
|
||||
p_feature text,
|
||||
p_base_limit integer,
|
||||
p_period text DEFAULT 'daily'
|
||||
) RETURNS jsonb
|
||||
LANGUAGE plpgsql
|
||||
SECURITY DEFINER
|
||||
SET search_path = public, pg_temp
|
||||
AS $$
|
||||
DECLARE
|
||||
v_window_start date;
|
||||
v_current integer := 0;
|
||||
v_overage integer := 0;
|
||||
v_new_overage integer;
|
||||
v_consumed_from text;
|
||||
BEGIN
|
||||
IF p_user_id IS NULL
|
||||
OR p_feature IS NULL
|
||||
OR p_base_limit IS NULL
|
||||
OR p_base_limit < -1
|
||||
OR p_period IS NULL
|
||||
OR p_period NOT IN ('daily', 'weekly') THEN
|
||||
RAISE EXCEPTION 'invalid_quota_consumption' USING ERRCODE = '22023';
|
||||
END IF;
|
||||
|
||||
-- Serialise read-then-increment for this user/feature.
|
||||
PERFORM pg_advisory_xact_lock(hashtextextended(p_user_id::text || ':' || p_feature, 20260928));
|
||||
|
||||
SELECT coalesce(overage_credits, 0) INTO v_overage
|
||||
FROM public.subscriptions
|
||||
WHERE user_id = p_user_id
|
||||
FOR UPDATE;
|
||||
v_overage := coalesce(v_overage, 0);
|
||||
|
||||
-- Not available: never spend credits on a feature the tier does not include.
|
||||
IF p_base_limit = 0 THEN
|
||||
RETURN jsonb_build_object(
|
||||
'allowed', false,
|
||||
'current', 0,
|
||||
'limit', 0,
|
||||
'overage_credits', v_overage,
|
||||
'consumed_from', 'none'
|
||||
);
|
||||
END IF;
|
||||
|
||||
v_window_start := CASE WHEN p_period = 'weekly' THEN CURRENT_DATE - 6 ELSE CURRENT_DATE END;
|
||||
|
||||
SELECT coalesce(sum(count), 0)::integer INTO v_current
|
||||
FROM public.daily_usage
|
||||
WHERE user_id = p_user_id
|
||||
AND feature = p_feature
|
||||
AND date >= v_window_start
|
||||
AND date <= CURRENT_DATE;
|
||||
|
||||
IF p_base_limit = -1 THEN
|
||||
v_consumed_from := 'unlimited';
|
||||
ELSIF v_current < p_base_limit THEN
|
||||
v_consumed_from := 'base';
|
||||
ELSE
|
||||
UPDATE public.subscriptions
|
||||
SET overage_credits = overage_credits - 1,
|
||||
updated_at = now()
|
||||
WHERE user_id = p_user_id
|
||||
AND overage_credits > 0
|
||||
RETURNING overage_credits INTO v_new_overage;
|
||||
|
||||
IF NOT FOUND THEN
|
||||
RETURN jsonb_build_object(
|
||||
'allowed', false,
|
||||
'current', v_current,
|
||||
'limit', p_base_limit,
|
||||
'overage_credits', 0,
|
||||
'consumed_from', 'none'
|
||||
);
|
||||
END IF;
|
||||
|
||||
v_overage := v_new_overage;
|
||||
v_consumed_from := 'overage';
|
||||
END IF;
|
||||
|
||||
INSERT INTO public.daily_usage (user_id, date, feature, count)
|
||||
VALUES (p_user_id, CURRENT_DATE, p_feature, 1)
|
||||
ON CONFLICT (user_id, date, feature) DO UPDATE
|
||||
SET count = public.daily_usage.count + 1;
|
||||
|
||||
RETURN jsonb_build_object(
|
||||
'allowed', true,
|
||||
'current', v_current + 1,
|
||||
'limit', p_base_limit,
|
||||
'overage_credits', v_overage,
|
||||
'consumed_from', v_consumed_from
|
||||
);
|
||||
END;
|
||||
$$;
|
||||
|
||||
COMMENT ON FUNCTION public.consume_quota(uuid, text, integer, text) IS
|
||||
'Atomic quota consumption over the policy window (daily = today, weekly = CURRENT_DATE-6..CURRENT_DATE): base first, then one overage credit. Used by llm-proxy and realtime-token.';
|
||||
|
||||
REVOKE ALL ON FUNCTION public.consume_quota(uuid, text, integer, text) FROM PUBLIC, anon, authenticated;
|
||||
GRANT EXECUTE ON FUNCTION public.consume_quota(uuid, text, integer, text) TO service_role;
|
||||
|
|
@ -0,0 +1,158 @@
|
|||
\set ON_ERROR_STOP on
|
||||
|
||||
-- Regression: weekly quota consumption must spend overage credits once the
|
||||
-- 7-day window total reaches the base limit, and usage from 7 days ago must
|
||||
-- no longer count (window = CURRENT_DATE-6 .. CURRENT_DATE).
|
||||
|
||||
BEGIN;
|
||||
|
||||
CREATE OR REPLACE FUNCTION pg_temp.assert_true(condition boolean, message text)
|
||||
RETURNS void
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
IF condition IS NOT TRUE THEN
|
||||
RAISE EXCEPTION 'assertion_failed: %', message;
|
||||
END IF;
|
||||
END;
|
||||
$$;
|
||||
|
||||
INSERT INTO auth.users (
|
||||
id, aud, role, email, encrypted_password, email_confirmed_at,
|
||||
raw_app_meta_data, raw_user_meta_data, created_at, updated_at
|
||||
) VALUES
|
||||
(
|
||||
'31000000-0000-4000-8000-000000000001', 'authenticated', 'authenticated',
|
||||
'consume-quota-weekly@example.invalid', crypt('fixture-password', gen_salt('bf')), now(),
|
||||
'{"provider":"email","providers":["email"]}'::jsonb, '{}'::jsonb, now(), now()
|
||||
),
|
||||
(
|
||||
'31000000-0000-4000-8000-000000000002', 'authenticated', 'authenticated',
|
||||
'consume-quota-daily@example.invalid', crypt('fixture-password', gen_salt('bf')), now(),
|
||||
'{"provider":"email","providers":["email"]}'::jsonb, '{}'::jsonb, now(), now()
|
||||
);
|
||||
|
||||
UPDATE public.subscriptions SET tier = 'free', status = 'active', overage_credits = 1
|
||||
WHERE user_id = '31000000-0000-4000-8000-000000000001';
|
||||
UPDATE public.subscriptions SET tier = 'pro', status = 'active', overage_credits = 0
|
||||
WHERE user_id = '31000000-0000-4000-8000-000000000002';
|
||||
|
||||
SELECT pg_temp.assert_true(
|
||||
NOT has_function_privilege('authenticated', 'public.consume_quota(uuid,text,integer,text)', 'EXECUTE'),
|
||||
'authenticated users cannot consume quota directly'
|
||||
);
|
||||
SELECT pg_temp.assert_true(
|
||||
has_function_privilege('service_role', 'public.consume_quota(uuid,text,integer,text)', 'EXECUTE'),
|
||||
'service role can consume quota'
|
||||
);
|
||||
|
||||
-- 250 Haiku calls spread over the six days before today (about 42 a day),
|
||||
-- plus a large row 7 days ago that has already left the weekly window.
|
||||
INSERT INTO public.daily_usage (user_id, date, feature, count)
|
||||
SELECT '31000000-0000-4000-8000-000000000001', CURRENT_DATE - d, 'llm_haiku',
|
||||
CASE WHEN d = 1 THEN 250 - 41 * 5 ELSE 41 END
|
||||
FROM generate_series(1, 6) AS d;
|
||||
INSERT INTO public.daily_usage (user_id, date, feature, count)
|
||||
VALUES ('31000000-0000-4000-8000-000000000001', CURRENT_DATE - 7, 'llm_haiku', 200);
|
||||
|
||||
DO $$
|
||||
DECLARE
|
||||
first_call jsonb;
|
||||
second_call jsonb;
|
||||
credits integer;
|
||||
today_count integer;
|
||||
BEGIN
|
||||
first_call := public.consume_quota('31000000-0000-4000-8000-000000000001', 'llm_haiku', 250, 'weekly');
|
||||
SELECT overage_credits INTO credits FROM public.subscriptions
|
||||
WHERE user_id = '31000000-0000-4000-8000-000000000001';
|
||||
PERFORM pg_temp.assert_true(
|
||||
(first_call->>'allowed')::boolean
|
||||
AND first_call->>'consumed_from' = 'overage'
|
||||
AND (first_call->>'overage_credits')::integer = 0
|
||||
AND (first_call->>'current')::integer = 251
|
||||
AND credits = 0,
|
||||
format('weekly total at the limit spends the overage credit: %s credits=%s', first_call, credits)
|
||||
);
|
||||
|
||||
second_call := public.consume_quota('31000000-0000-4000-8000-000000000001', 'llm_haiku', 250, 'weekly');
|
||||
SELECT count INTO today_count FROM public.daily_usage
|
||||
WHERE user_id = '31000000-0000-4000-8000-000000000001' AND feature = 'llm_haiku' AND date = CURRENT_DATE;
|
||||
PERFORM pg_temp.assert_true(
|
||||
NOT (second_call->>'allowed')::boolean
|
||||
AND second_call->>'consumed_from' = 'none'
|
||||
AND today_count = 1,
|
||||
format('no base and no credit left denies without counting: %s today=%s', second_call, today_count)
|
||||
);
|
||||
END;
|
||||
$$;
|
||||
|
||||
-- The row from 7 days ago rolls off: removing one day inside the window
|
||||
-- frees base allowance even though CURRENT_DATE-7 still holds 200 uses.
|
||||
DELETE FROM public.daily_usage
|
||||
WHERE user_id = '31000000-0000-4000-8000-000000000001' AND feature = 'llm_haiku' AND date = CURRENT_DATE - 2;
|
||||
|
||||
DO $$
|
||||
DECLARE
|
||||
call jsonb;
|
||||
BEGIN
|
||||
call := public.consume_quota('31000000-0000-4000-8000-000000000001', 'llm_haiku', 250, 'weekly');
|
||||
PERFORM pg_temp.assert_true(
|
||||
(call->>'allowed')::boolean AND call->>'consumed_from' = 'base',
|
||||
format('usage older than CURRENT_DATE-6 does not count: %s', call)
|
||||
);
|
||||
END;
|
||||
$$;
|
||||
|
||||
-- Not-available features never spend credits.
|
||||
UPDATE public.subscriptions SET overage_credits = 3
|
||||
WHERE user_id = '31000000-0000-4000-8000-000000000001';
|
||||
DO $$
|
||||
DECLARE
|
||||
call jsonb;
|
||||
credits integer;
|
||||
BEGIN
|
||||
call := public.consume_quota('31000000-0000-4000-8000-000000000001', 'llm_sonnet', 0, 'daily');
|
||||
SELECT overage_credits INTO credits FROM public.subscriptions
|
||||
WHERE user_id = '31000000-0000-4000-8000-000000000001';
|
||||
PERFORM pg_temp.assert_true(
|
||||
NOT (call->>'allowed')::boolean AND call->>'consumed_from' = 'none' AND credits = 3,
|
||||
format('limit 0 is denied without touching credits: %s credits=%s', call, credits)
|
||||
);
|
||||
END;
|
||||
$$;
|
||||
|
||||
-- Daily policies (and the legacy three-argument call) still count only today.
|
||||
INSERT INTO public.daily_usage (user_id, date, feature, count)
|
||||
VALUES ('31000000-0000-4000-8000-000000000002', CURRENT_DATE - 1, 'llm_sonnet', 300);
|
||||
|
||||
DO $$
|
||||
DECLARE
|
||||
daily_call jsonb;
|
||||
legacy_call jsonb;
|
||||
unlimited_call jsonb;
|
||||
BEGIN
|
||||
daily_call := public.consume_quota('31000000-0000-4000-8000-000000000002', 'llm_sonnet', 300, 'daily');
|
||||
legacy_call := public.consume_quota('31000000-0000-4000-8000-000000000002', 'llm_sonnet', 300);
|
||||
unlimited_call := public.consume_quota('31000000-0000-4000-8000-000000000002', 'stt_transcribe', -1, 'daily');
|
||||
PERFORM pg_temp.assert_true(
|
||||
daily_call->>'consumed_from' = 'base'
|
||||
AND (daily_call->>'current')::integer = 1
|
||||
AND legacy_call->>'consumed_from' = 'base'
|
||||
AND (legacy_call->>'current')::integer = 2
|
||||
AND unlimited_call->>'consumed_from' = 'unlimited'
|
||||
AND (unlimited_call->>'limit')::integer = -1,
|
||||
format('daily window ignores yesterday: %s / %s / %s', daily_call, legacy_call, unlimited_call)
|
||||
);
|
||||
END;
|
||||
$$;
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
PERFORM public.consume_quota('31000000-0000-4000-8000-000000000002', 'llm_sonnet', 300, 'monthly');
|
||||
RAISE EXCEPTION 'assertion_failed: unknown period was accepted';
|
||||
EXCEPTION
|
||||
WHEN invalid_parameter_value THEN NULL;
|
||||
END;
|
||||
$$;
|
||||
|
||||
ROLLBACK;
|
||||
Loading…
Add table
Add a link
Reference in a new issue