296 lines
10 KiB
PL/PgSQL
296 lines
10 KiB
PL/PgSQL
-- llm-proxy: reserve LLM quota before calling the provider.
|
|
--
|
|
-- llm-proxy used to run a read-only checkQuota (a SELECT on daily_usage),
|
|
-- call Anthropic, and only then charge with consume_quota. Nothing counted
|
|
-- requests that were still in flight, so N parallel requests from one user all
|
|
-- passed the pre-check and each paid for a full Anthropic generation before
|
|
-- consume_quota rejected every request past the allowance and threw the paid
|
|
-- reply away (the same flaw 20260928000037 fixed for meeting documents).
|
|
--
|
|
-- reserve_llm_quota now takes one unit up front, under the same
|
|
-- per-user/per-feature advisory lock that consume_quota uses, and records it
|
|
-- in daily_usage immediately, so a concurrent request (or a consume_quota
|
|
-- call for the same feature) sees it. The edge function calls the provider
|
|
-- only while it holds a reservation, then settles it:
|
|
-- * finalize_llm_quota(id, true) -> 'completed' (the unit stays spent)
|
|
-- * finalize_llm_quota(id, false) -> 'released' (daily_usage -1, and the
|
|
-- overage credit is returned when the unit came from overage)
|
|
-- A crashed worker cannot hold a unit forever: reservations whose 10 minute
|
|
-- lease has expired are released by the next reserve call for that
|
|
-- user/feature, well beyond the proxy's 45 s time-to-first-byte and 180 s
|
|
-- stream deadlines.
|
|
--
|
|
-- p_base_limit / p_period come from PLAN_QUOTA (packages/core plan catalog),
|
|
-- exactly like consume_quota. Only service_role may call these functions.
|
|
-- Apply this migration before deploying the llm-proxy that calls it.
|
|
|
|
CREATE TABLE IF NOT EXISTS public.llm_quota_reservations (
|
|
id uuid PRIMARY KEY,
|
|
user_id uuid NOT NULL REFERENCES auth.users(id) ON DELETE CASCADE,
|
|
feature text NOT NULL CHECK (feature IN ('llm_haiku', 'llm_sonnet', 'llm_opus')),
|
|
usage_date date NOT NULL DEFAULT CURRENT_DATE,
|
|
consumed_from text NOT NULL CHECK (consumed_from IN ('base', 'overage', 'unlimited')),
|
|
status text NOT NULL DEFAULT 'reserved' CHECK (status IN ('reserved', 'completed', 'released')),
|
|
tier text NOT NULL,
|
|
quota_period text NOT NULL CHECK (quota_period IN ('daily', 'weekly')),
|
|
quota_limit integer NOT NULL,
|
|
current_count integer NOT NULL,
|
|
overage_after integer NOT NULL,
|
|
lease_expires_at timestamptz NOT NULL,
|
|
created_at timestamptz NOT NULL DEFAULT now(),
|
|
finalized_at timestamptz,
|
|
release_reason text
|
|
);
|
|
|
|
ALTER TABLE public.llm_quota_reservations ENABLE ROW LEVEL SECURITY;
|
|
REVOKE ALL ON TABLE public.llm_quota_reservations FROM PUBLIC, anon, authenticated;
|
|
GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE public.llm_quota_reservations TO service_role;
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_llm_quota_reservations_reclaim
|
|
ON public.llm_quota_reservations(user_id, feature, lease_expires_at)
|
|
WHERE status = 'reserved';
|
|
|
|
CREATE OR REPLACE FUNCTION public.reserve_llm_quota(
|
|
p_user_id uuid,
|
|
p_reservation_id uuid,
|
|
p_feature text,
|
|
p_base_limit integer,
|
|
p_period text
|
|
) RETURNS jsonb
|
|
LANGUAGE plpgsql
|
|
SECURITY DEFINER
|
|
SET search_path = public, pg_temp
|
|
AS $$
|
|
DECLARE
|
|
existing public.llm_quota_reservations%ROWTYPE;
|
|
expired public.llm_quota_reservations%ROWTYPE;
|
|
subscription_tier text := 'free';
|
|
overage integer := 0;
|
|
new_overage integer;
|
|
current_count integer := 0;
|
|
consumed_from text;
|
|
BEGIN
|
|
IF p_user_id IS NULL
|
|
OR p_reservation_id IS NULL
|
|
OR p_feature IS NULL
|
|
OR p_feature NOT IN ('llm_haiku', 'llm_sonnet', 'llm_opus')
|
|
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_llm_quota_reservation' USING ERRCODE = '22023';
|
|
END IF;
|
|
|
|
-- Same lock key as consume_quota so both paths serialise per user/feature.
|
|
PERFORM pg_advisory_xact_lock(hashtextextended(p_user_id::text || ':' || p_feature, 20260928));
|
|
|
|
SELECT * INTO existing
|
|
FROM public.llm_quota_reservations
|
|
WHERE id = p_reservation_id
|
|
FOR UPDATE;
|
|
|
|
IF FOUND THEN
|
|
IF existing.user_id <> p_user_id OR existing.feature <> p_feature THEN
|
|
RAISE EXCEPTION 'llm_quota_reservation_conflict' USING ERRCODE = 'PT409';
|
|
END IF;
|
|
RETURN jsonb_build_object(
|
|
'allowed', existing.status IN ('reserved', 'completed'),
|
|
'reservation_id', existing.id,
|
|
'status', existing.status,
|
|
'current', existing.current_count,
|
|
'limit', existing.quota_limit,
|
|
'period', existing.quota_period,
|
|
'tier', existing.tier,
|
|
'overage_credits', existing.overage_after,
|
|
'consumed_from', existing.consumed_from
|
|
);
|
|
END IF;
|
|
|
|
-- Reclaim crashed requests before calculating the next allowance.
|
|
FOR expired IN
|
|
SELECT *
|
|
FROM public.llm_quota_reservations
|
|
WHERE user_id = p_user_id
|
|
AND feature = p_feature
|
|
AND status = 'reserved'
|
|
AND lease_expires_at <= now()
|
|
FOR UPDATE
|
|
LOOP
|
|
UPDATE public.daily_usage
|
|
SET count = greatest(count - 1, 0)
|
|
WHERE user_id = expired.user_id
|
|
AND date = expired.usage_date
|
|
AND feature = expired.feature;
|
|
|
|
IF expired.consumed_from = 'overage' THEN
|
|
UPDATE public.subscriptions
|
|
SET overage_credits = overage_credits + 1,
|
|
updated_at = now()
|
|
WHERE user_id = expired.user_id;
|
|
END IF;
|
|
|
|
UPDATE public.llm_quota_reservations
|
|
SET status = 'released', finalized_at = now(), release_reason = 'lease_expired'
|
|
WHERE id = expired.id;
|
|
END LOOP;
|
|
|
|
SELECT coalesce(tier, 'free'), coalesce(overage_credits, 0)
|
|
INTO subscription_tier, overage
|
|
FROM public.subscriptions
|
|
WHERE user_id = p_user_id
|
|
FOR UPDATE;
|
|
|
|
IF NOT FOUND THEN
|
|
subscription_tier := 'free';
|
|
overage := 0;
|
|
END IF;
|
|
|
|
-- Not available: never spend credits on a model the tier does not include.
|
|
IF p_base_limit = 0 THEN
|
|
RETURN jsonb_build_object(
|
|
'allowed', false,
|
|
'reservation_id', NULL,
|
|
'status', 'denied',
|
|
'current', 0,
|
|
'limit', 0,
|
|
'period', p_period,
|
|
'tier', subscription_tier,
|
|
'overage_credits', overage,
|
|
'consumed_from', 'none'
|
|
);
|
|
END IF;
|
|
|
|
-- daily_usage already includes units held by in-flight reservations.
|
|
SELECT coalesce(sum(count), 0)::integer INTO current_count
|
|
FROM public.daily_usage
|
|
WHERE user_id = p_user_id
|
|
AND feature = p_feature
|
|
AND date >= CASE WHEN p_period = 'weekly' THEN CURRENT_DATE - 6 ELSE CURRENT_DATE END
|
|
AND date <= CURRENT_DATE;
|
|
|
|
IF p_base_limit = -1 THEN
|
|
consumed_from := 'unlimited';
|
|
ELSIF current_count < p_base_limit THEN
|
|
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 new_overage;
|
|
|
|
IF NOT FOUND THEN
|
|
RETURN jsonb_build_object(
|
|
'allowed', false,
|
|
'reservation_id', NULL,
|
|
'status', 'denied',
|
|
'current', current_count,
|
|
'limit', p_base_limit,
|
|
'period', p_period,
|
|
'tier', subscription_tier,
|
|
'overage_credits', 0,
|
|
'consumed_from', 'none'
|
|
);
|
|
END IF;
|
|
|
|
overage := new_overage;
|
|
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;
|
|
|
|
current_count := current_count + 1;
|
|
INSERT INTO public.llm_quota_reservations(
|
|
id, user_id, feature, consumed_from, tier, quota_period, quota_limit,
|
|
current_count, overage_after, lease_expires_at
|
|
) VALUES (
|
|
p_reservation_id, p_user_id, p_feature, consumed_from, subscription_tier, p_period, p_base_limit,
|
|
current_count, overage, now() + interval '10 minutes'
|
|
);
|
|
|
|
RETURN jsonb_build_object(
|
|
'allowed', true,
|
|
'reservation_id', p_reservation_id,
|
|
'status', 'reserved',
|
|
'current', current_count,
|
|
'limit', p_base_limit,
|
|
'period', p_period,
|
|
'tier', subscription_tier,
|
|
'overage_credits', overage,
|
|
'consumed_from', consumed_from
|
|
);
|
|
END;
|
|
$$;
|
|
|
|
CREATE OR REPLACE FUNCTION public.finalize_llm_quota(
|
|
p_reservation_id uuid,
|
|
p_succeeded boolean
|
|
) RETURNS jsonb
|
|
LANGUAGE plpgsql
|
|
SECURITY DEFINER
|
|
SET search_path = public, pg_temp
|
|
AS $$
|
|
DECLARE
|
|
reservation public.llm_quota_reservations%ROWTYPE;
|
|
final_status text;
|
|
BEGIN
|
|
IF p_reservation_id IS NULL OR p_succeeded IS NULL THEN
|
|
RAISE EXCEPTION 'invalid_llm_quota_finalize' USING ERRCODE = '22023';
|
|
END IF;
|
|
|
|
SELECT * INTO reservation
|
|
FROM public.llm_quota_reservations
|
|
WHERE id = p_reservation_id;
|
|
IF NOT FOUND THEN
|
|
RAISE EXCEPTION 'llm_quota_reservation_not_found' USING ERRCODE = 'P0002';
|
|
END IF;
|
|
|
|
PERFORM pg_advisory_xact_lock(hashtextextended(reservation.user_id::text || ':' || reservation.feature, 20260928));
|
|
SELECT * INTO reservation
|
|
FROM public.llm_quota_reservations
|
|
WHERE id = p_reservation_id
|
|
FOR UPDATE;
|
|
|
|
IF reservation.status <> 'reserved' THEN
|
|
RETURN jsonb_build_object('reservation_id', reservation.id, 'status', reservation.status);
|
|
END IF;
|
|
|
|
IF p_succeeded THEN
|
|
final_status := 'completed';
|
|
ELSE
|
|
UPDATE public.daily_usage
|
|
SET count = greatest(count - 1, 0)
|
|
WHERE user_id = reservation.user_id
|
|
AND date = reservation.usage_date
|
|
AND feature = reservation.feature;
|
|
|
|
IF reservation.consumed_from = 'overage' THEN
|
|
UPDATE public.subscriptions
|
|
SET overage_credits = overage_credits + 1,
|
|
updated_at = now()
|
|
WHERE user_id = reservation.user_id;
|
|
END IF;
|
|
final_status := 'released';
|
|
END IF;
|
|
|
|
UPDATE public.llm_quota_reservations
|
|
SET status = final_status,
|
|
finalized_at = now(),
|
|
release_reason = CASE WHEN p_succeeded THEN NULL ELSE 'provider_failed' END
|
|
WHERE id = reservation.id;
|
|
|
|
RETURN jsonb_build_object('reservation_id', reservation.id, 'status', final_status);
|
|
END;
|
|
$$;
|
|
|
|
REVOKE ALL ON FUNCTION public.reserve_llm_quota(uuid, uuid, text, integer, text) FROM PUBLIC, anon, authenticated;
|
|
REVOKE ALL ON FUNCTION public.finalize_llm_quota(uuid, boolean) FROM PUBLIC, anon, authenticated;
|
|
GRANT EXECUTE ON FUNCTION public.reserve_llm_quota(uuid, uuid, text, integer, text) TO service_role;
|
|
GRANT EXECUTE ON FUNCTION public.finalize_llm_quota(uuid, boolean) TO service_role;
|
|
|
|
COMMENT ON TABLE public.llm_quota_reservations IS
|
|
'Service-only leases that reserve LLM quota before the Anthropic call in llm-proxy and refund failed or expired generations.';
|