fix(quota): spend overage credits against the weekly window and count 7 days, not 8
This commit is contained in:
Yun Chan
parent
f4f9653361
commit
102f3ef934
5 changed files with 601 additions and 65 deletions
|
|
@ -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