fix(meeting-document): hold quota for in-flight generation claims
This commit is contained in:
Yun Chan
parent
f4724ddf53
commit
b35676c75c
5 changed files with 994 additions and 179 deletions
|
|
@ -0,0 +1,162 @@
|
|||
\set ON_ERROR_STOP on
|
||||
|
||||
-- Regression (red-team r1-17): claim_meeting_document_generation_v1 used to
|
||||
-- read daily_usage without counting in-flight claims, so N parallel requests
|
||||
-- with fresh idempotency keys all passed the quota check and each paid for a
|
||||
-- provider call before commit rejected N-1 of them. A 'processing' request now
|
||||
-- holds one unit of the allowance until it is failed, committed or its lease
|
||||
-- expires.
|
||||
|
||||
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;
|
||||
$$;
|
||||
|
||||
-- Returns the claim payload, or {"error": SQLERRM} when the claim raises.
|
||||
CREATE OR REPLACE FUNCTION pg_temp.try_claim(p_actor uuid, p_key uuid, p_meeting uuid, p_template uuid, p_model text)
|
||||
RETURNS jsonb
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
BEGIN
|
||||
RETURN public.claim_meeting_document_generation_v1(
|
||||
p_actor, p_key, p_meeting, p_template, 'Quota fixture document', p_model
|
||||
);
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
RETURN jsonb_build_object('error', SQLERRM);
|
||||
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
|
||||
(
|
||||
'37000000-0000-4000-8000-000000000001', 'authenticated', 'authenticated',
|
||||
'meeting-doc-quota-free@example.invalid', crypt('fixture-password', gen_salt('bf')), now(),
|
||||
'{"provider":"email","providers":["email"]}'::jsonb, '{}'::jsonb, now(), now()
|
||||
),
|
||||
(
|
||||
'37000000-0000-4000-8000-000000000002', 'authenticated', 'authenticated',
|
||||
'meeting-doc-quota-unlimited@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 = 0
|
||||
WHERE user_id = '37000000-0000-4000-8000-000000000001';
|
||||
UPDATE public.subscriptions SET tier = 'pro_plus', status = 'active', overage_credits = 0
|
||||
WHERE user_id = '37000000-0000-4000-8000-000000000002';
|
||||
|
||||
INSERT INTO public.meetings (id, user_id, title, status, raw_transcript)
|
||||
VALUES
|
||||
('37100000-0000-4000-8000-000000000001', '37000000-0000-4000-8000-000000000001',
|
||||
'Quota fixture meeting', 'completed', 'Speaker one talked about the roadmap.'),
|
||||
('37100000-0000-4000-8000-000000000002', '37000000-0000-4000-8000-000000000002',
|
||||
'Unlimited fixture meeting', 'completed', 'Speaker two talked about hiring.');
|
||||
|
||||
INSERT INTO public.user_templates (
|
||||
id, user_id, template_kind, name, template_type, system_prompt, is_builtin
|
||||
) VALUES
|
||||
('37200000-0000-4000-8000-000000000001', '37000000-0000-4000-8000-000000000001',
|
||||
'meeting_document', 'Quota fixture template', 'custom', 'Summarize the meeting.', false),
|
||||
('37200000-0000-4000-8000-000000000002', '37000000-0000-4000-8000-000000000002',
|
||||
'meeting_document', 'Unlimited fixture template', 'custom', 'Summarize the meeting.', false);
|
||||
|
||||
-- One weekly Haiku unit left for the free user.
|
||||
INSERT INTO public.daily_usage (user_id, date, feature, count)
|
||||
VALUES ('37000000-0000-4000-8000-000000000001', CURRENT_DATE, 'llm_haiku', 249);
|
||||
|
||||
DO $$
|
||||
DECLARE
|
||||
actor constant uuid := '37000000-0000-4000-8000-000000000001';
|
||||
meeting constant uuid := '37100000-0000-4000-8000-000000000001';
|
||||
template constant uuid := '37200000-0000-4000-8000-000000000001';
|
||||
haiku constant text := 'claude-haiku-4-5-20251001';
|
||||
first jsonb;
|
||||
parallel jsonb;
|
||||
replay jsonb;
|
||||
after_release jsonb;
|
||||
after_commit jsonb;
|
||||
overage_claim jsonb;
|
||||
overage_parallel jsonb;
|
||||
after_lease jsonb;
|
||||
usage_count integer;
|
||||
BEGIN
|
||||
first := pg_temp.try_claim(actor, '37300000-0000-4000-8000-000000000001', meeting, template, haiku);
|
||||
PERFORM pg_temp.assert_true((first->>'claimed')::boolean, 'last weekly unit can be claimed');
|
||||
|
||||
-- The bug: a second fresh key used to pass because only daily_usage was read.
|
||||
parallel := pg_temp.try_claim(actor, '37300000-0000-4000-8000-000000000002', meeting, template, haiku);
|
||||
PERFORM pg_temp.assert_true(
|
||||
parallel->>'error' = 'generation_quota_exceeded',
|
||||
'an in-flight claim holds the last unit, so a parallel claim is rejected before provider work: ' || parallel::text
|
||||
);
|
||||
|
||||
-- Replaying the in-flight key is idempotent, not a quota error.
|
||||
replay := pg_temp.try_claim(actor, '37300000-0000-4000-8000-000000000001', meeting, template, haiku);
|
||||
PERFORM pg_temp.assert_true(
|
||||
replay->>'error' IS NULL AND NOT (replay->>'claimed')::boolean AND replay->>'status' = 'processing',
|
||||
'replaying the in-flight key reports processing: ' || replay::text
|
||||
);
|
||||
|
||||
-- A failed request releases the unit it held.
|
||||
PERFORM public.fail_meeting_document_generation_v1(actor, '37300000-0000-4000-8000-000000000001', 'provider_timeout');
|
||||
after_release := pg_temp.try_claim(actor, '37300000-0000-4000-8000-000000000003', meeting, template, haiku);
|
||||
PERFORM pg_temp.assert_true((after_release->>'claimed')::boolean,
|
||||
'failure releases the in-flight unit: ' || after_release::text);
|
||||
|
||||
-- Commit converts the held unit into recorded usage; the allowance is spent.
|
||||
PERFORM public.commit_meeting_document_generation_v1(
|
||||
actor, '37300000-0000-4000-8000-000000000003', 'Generated body', 10, 1, 1
|
||||
);
|
||||
SELECT count INTO usage_count FROM public.daily_usage
|
||||
WHERE user_id = actor AND date = CURRENT_DATE AND feature = 'llm_haiku';
|
||||
PERFORM pg_temp.assert_true(usage_count = 250, 'commit records exactly one unit');
|
||||
after_commit := pg_temp.try_claim(actor, '37300000-0000-4000-8000-000000000004', meeting, template, haiku);
|
||||
PERFORM pg_temp.assert_true(after_commit->>'error' = 'generation_quota_exceeded',
|
||||
'exhausted allowance rejects new claims: ' || after_commit::text);
|
||||
|
||||
-- One overage credit covers exactly one in-flight request.
|
||||
UPDATE public.subscriptions SET overage_credits = 1 WHERE user_id = actor;
|
||||
overage_claim := pg_temp.try_claim(actor, '37300000-0000-4000-8000-000000000005', meeting, template, haiku);
|
||||
PERFORM pg_temp.assert_true((overage_claim->>'claimed')::boolean,
|
||||
'an overage credit allows one claim past the base limit: ' || overage_claim::text);
|
||||
overage_parallel := pg_temp.try_claim(actor, '37300000-0000-4000-8000-000000000006', meeting, template, haiku);
|
||||
PERFORM pg_temp.assert_true(overage_parallel->>'error' = 'generation_quota_exceeded',
|
||||
'a single overage credit is not spent twice by parallel claims: ' || overage_parallel::text);
|
||||
|
||||
-- A crashed request stops holding its unit once its lease expires.
|
||||
UPDATE public.meeting_document_generation_requests
|
||||
SET created_at = now() - interval '11 minutes'
|
||||
WHERE user_id = actor AND idempotency_key = '37300000-0000-4000-8000-000000000005';
|
||||
after_lease := pg_temp.try_claim(actor, '37300000-0000-4000-8000-000000000007', meeting, template, haiku);
|
||||
PERFORM pg_temp.assert_true((after_lease->>'claimed')::boolean,
|
||||
'an expired in-flight lease no longer holds a unit: ' || after_lease::text);
|
||||
END;
|
||||
$$;
|
||||
|
||||
DO $$
|
||||
DECLARE
|
||||
actor constant uuid := '37000000-0000-4000-8000-000000000002';
|
||||
meeting constant uuid := '37100000-0000-4000-8000-000000000002';
|
||||
template constant uuid := '37200000-0000-4000-8000-000000000002';
|
||||
first jsonb;
|
||||
second jsonb;
|
||||
BEGIN
|
||||
first := pg_temp.try_claim(actor, '37400000-0000-4000-8000-000000000001', meeting, template, 'claude-haiku-4-5-20251001');
|
||||
second := pg_temp.try_claim(actor, '37400000-0000-4000-8000-000000000002', meeting, template, 'claude-haiku-4-5-20251001');
|
||||
PERFORM pg_temp.assert_true(
|
||||
(first->>'claimed')::boolean AND (second->>'claimed')::boolean,
|
||||
'unlimited allowances are not throttled by in-flight claims'
|
||||
);
|
||||
END;
|
||||
$$;
|
||||
|
||||
ROLLBACK;
|
||||
Loading…
Add table
Add a link
Reference in a new issue