diff --git a/server/supabase/functions/_shared/quota-policy.ts b/server/supabase/functions/_shared/quota-policy.ts new file mode 100644 index 0000000..c455ec9 --- /dev/null +++ b/server/supabase/functions/_shared/quota-policy.ts @@ -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 + 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, + } +} diff --git a/server/supabase/functions/_shared/quota.test.ts b/server/supabase/functions/_shared/quota.test.ts new file mode 100644 index 0000000..2884e1f --- /dev/null +++ b/server/supabase/functions/_shared/quota.test.ts @@ -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 +type ServiceClient = ReturnType + +/** 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( + onfulfilled?: ((value: { data: unknown; error: null }) => T1 | PromiseLike) | null, + onrejected?: ((reason: unknown) => T2 | PromiseLike) | null, + ): PromiseLike { + 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 = { subscriptions: [], daily_usage: [] } + rpcCalls: Array<{ name: string; args: Record }> = [] + 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): 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)}`) + } +}) diff --git a/server/supabase/functions/_shared/quota.ts b/server/supabase/functions/_shared/quota.ts index a4d1333..54ed7db 100644 --- a/server/supabase/functions/_shared/quota.ts +++ b/server/supabase/functions/_shared/quota.ts @@ -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, -): Promise { - // 티어 + overage 조회 +): Promise { const { data: sub } = await serviceRoleClient .from('subscriptions') .select('tier, overage_credits') .eq('user_id', userId) .single() + return { + tier: (sub?.tier as Tier) ?? 'free', + overageCredits: (sub?.overage_credits as number) ?? 0, + } +} - const tier: Tier = (sub?.tier as Tier) ?? 'free' - const overageCredits = (sub?.overage_credits as number) ?? 0 +/** 기간 창(daily=오늘, weekly=오늘-6..오늘) 안의 daily_usage 합계. */ +async function readWindowUsage( + userId: string, + feature: QuotaFeature, + period: QuotaPeriod, + serviceRoleClient: ReturnType, + now: Date, +): Promise { + const { data: rows } = await serviceRoleClient + .from('daily_usage') + .select('count') + .eq('user_id', userId) + .eq('feature', feature) + .gte('date', quotaWindowStart(period, now)) + .lte('date', quotaUsageDate(now)) + const usageRows = (rows ?? []) as Array<{ count: number | null }> + 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, + now: Date = new Date(), +): Promise { + const { tier, overageCredits } = await readSubscriptionQuotaState(userId, serviceRoleClient) const policy = getQuotaPolicy(tier, feature) - // 사용불가 (limit=0) - if (policy.limit === 0) { - return { allowed: false, current: 0, limit: 0, period: policy.period, tier, overageCredits } - } - - // 무제한 - 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) - const { data: rows } = await serviceRoleClient - .from('daily_usage') - .select('count') - .eq('user_id', userId) - .eq('feature', feature) - .gte('date', weekAgo.toISOString().slice(0, 10)) - 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 + // 사용불가 (limit=0) / 무제한 (limit=-1) 은 사용량 조회가 필요 없다. + if (policy.limit === 0 || policy.limit === -1) { + return { + 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: current < policy.limit || overageCredits > 0, + 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, baseLimit: number, + period?: QuotaPeriod, ): Promise { + 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) } /** diff --git a/server/supabase/migrations/20260928000131_consume_quota_period_window.sql b/server/supabase/migrations/20260928000131_consume_quota_period_window.sql new file mode 100644 index 0000000..119ddcb --- /dev/null +++ b/server/supabase/migrations/20260928000131_consume_quota_period_window.sql @@ -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; diff --git a/server/supabase/tests/consume-quota-period-window.integration.sql b/server/supabase/tests/consume-quota-period-window.integration.sql new file mode 100644 index 0000000..ecbe53f --- /dev/null +++ b/server/supabase/tests/consume-quota-period-window.integration.sql @@ -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;