From 2889f9d35a56ccee750b5aad0c5345e2f3f07881 Mon Sep 17 00:00:00 2001 From: Jon Durbin Date: Mon, 30 Mar 2026 16:20:51 -0400 Subject: [PATCH 1/3] Control quota usage (or strictly paygo) with headers. --- api/invocation/util.py | 34 ++++++++++++++++++++++++++++++++++ api/main.py | 2 ++ 2 files changed, 36 insertions(+) diff --git a/api/invocation/util.py b/api/invocation/util.py index f6690d6c..991bc032 100644 --- a/api/invocation/util.py +++ b/api/invocation/util.py @@ -452,6 +452,19 @@ async def check_quota_and_balance(request, current_user, chute): from api.user.schemas import InvocationQuota from api.user.service import chutes_user_id + # Optional client headers to control quota/paygo behavior. + # X-Paygo-Only: skip quota, charge directly from balance. + # X-Quota-Only: never fall back to paygo if quota/caps are exceeded. + paygo_only = request.headers.get("x-paygo-only", "").lower() == "true" + quota_only = request.headers.get("x-quota-only", "").lower() == "true" + if paygo_only and quota_only: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Cannot set both X-Paygo-Only and X-Quota-Only headers.", + ) + request.state.paygo_only = paygo_only + request.state.quota_only = quota_only + quota_date = date.today() # Fully discounted chutes are free but have usage caps for unprivileged users. @@ -527,6 +540,13 @@ async def check_quota_and_balance(request, current_user, chute): ) request.state.free_invocation = True + # X-Paygo-Only: clear free_invocation so the request enters the quota/paygo + # path below. 100% discounted chutes (discount == 1.0) are inherently $0 so + # there is nothing to charge; private chutes are zeroed out in the billing + # layer regardless, so we only override for public, non-free chutes. + if paygo_only and request.state.free_invocation and chute.discount < 1.0 and chute.public: + request.state.free_invocation = False + # Check account quotas if not free/invoiced. quota_date = date.today() if not ( @@ -557,6 +577,10 @@ async def check_quota_and_balance(request, current_user, chute): ): quota = 0 + # X-Paygo-Only: treat quota as 0 so the user always falls through to paygo. + if paygo_only: + quota = 0 + # Quota-200 users (one-time $5 payment) cannot use TEE models without balance. # $3/mo sub users (quota 300 or 301) cannot use premium chutes without balance. # In both cases, if the user has balance, force paygo (never free_invocation). @@ -579,6 +603,11 @@ async def check_quota_and_balance(request, current_user, chute): # Automatically switch to paygo when the quota is exceeded. if request_count >= quota: + if quota_only and not request.state.free_invocation: + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail="Quota exceeded and X-Quota-Only is set, paygo fallback disabled.", + ) if effective_balance <= 0 and not request.state.free_invocation: logger.warning( f"Payment required: attempted invocation of {chute.name} " @@ -665,6 +694,11 @@ async def check_quota_and_balance(request, current_user, chute): f"Subscription cap exceeded for {current_user.user_id} " f"[{current_user.username}]: {', '.join(exceeded)}" ) + if quota_only: + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail="Subscription cap exceeded and X-Quota-Only is set, paygo fallback disabled.", + ) if effective_balance <= 0: raise HTTPException( status_code=status.HTTP_402_PAYMENT_REQUIRED, diff --git a/api/main.py b/api/main.py index 161f54be..3d5df095 100644 --- a/api/main.py +++ b/api/main.py @@ -282,6 +282,8 @@ async def host_router_middleware(request: Request, call_next): request.state.chute_id = None request.state.squad_request = False request.state.free_invocation = False + request.state.paygo_only = False + request.state.quota_only = False host = request.headers.get("host", "") host_parts = re.search(r"^([a-z0-9-]+)\.[a-z0-9-]+", host.lower()) From 9664bcd8664977073ecdc39cca3186b80aefdfbd Mon Sep 17 00:00:00 2001 From: Jon Durbin Date: Wed, 1 Apr 2026 05:12:21 -0400 Subject: [PATCH 2/3] More paths for paygo only/etc. --- api/invocation/util.py | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/api/invocation/util.py b/api/invocation/util.py index 991bc032..7c12fe96 100644 --- a/api/invocation/util.py +++ b/api/invocation/util.py @@ -577,15 +577,16 @@ async def check_quota_and_balance(request, current_user, chute): ): quota = 0 - # X-Paygo-Only: treat quota as 0 so the user always falls through to paygo. - if paygo_only: - quota = 0 - # Quota-200 users (one-time $5 payment) cannot use TEE models without balance. # $3/mo sub users (quota 300 or 301) cannot use premium chutes without balance. # In both cases, if the user has balance, force paygo (never free_invocation). force_paygo = False if quota == 200 and chute.tee: + if quota_only: + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail="TEE models require paygo; incompatible with X-Quota-Only.", + ) if effective_balance <= 0: raise HTTPException( status_code=status.HTTP_402_PAYMENT_REQUIRED, @@ -594,6 +595,11 @@ async def check_quota_and_balance(request, current_user, chute): force_paygo = True if get_subscription_tier(quota) == 3.0 and chute.chute_id in settings.premium_chute_ids: + if quota_only: + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail="This premium model requires paygo; incompatible with X-Quota-Only.", + ) if effective_balance <= 0: raise HTTPException( status_code=status.HTTP_402_PAYMENT_REQUIRED, @@ -601,6 +607,11 @@ async def check_quota_and_balance(request, current_user, chute): ) force_paygo = True + # X-Paygo-Only: treat quota as 0 so the user always falls through to paygo. + # Placed after force_paygo checks so they can evaluate the real quota value. + if paygo_only: + quota = 0 + # Automatically switch to paygo when the quota is exceeded. if request_count >= quota: if quota_only and not request.state.free_invocation: @@ -634,6 +645,15 @@ async def check_quota_and_balance(request, current_user, chute): error_kwargs["detail"]["quota_reset_timestamp"] = quota_reset raise HTTPException(**error_kwargs) + # X-Paygo-Only subscribers should still receive their paygo discount. + if paygo_only: + ( + po_sub_quota, _, _, _, + ) = await InvocationQuota.get_subscription_record(current_user.user_id) + if (po_price := get_subscription_tier(po_sub_quota)) is not None: + request.state.subscriber_paygo_discount = SUBSCRIPTION_PAYGO_DISCOUNTS.get( + po_price, 0.0 + ) else: # When within the quota, check subscription caps before marking as free. # force_paygo skips free_invocation entirely (TEE/premium restrictions). From 11220aa526d55f70c5adc60c4fb0ed9ac2b36a81 Mon Sep 17 00:00:00 2001 From: Jon Durbin Date: Wed, 1 Apr 2026 05:16:55 -0400 Subject: [PATCH 3/3] Lint. --- api/invocation/util.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/api/invocation/util.py b/api/invocation/util.py index 7c12fe96..050f8cbb 100644 --- a/api/invocation/util.py +++ b/api/invocation/util.py @@ -648,7 +648,10 @@ async def check_quota_and_balance(request, current_user, chute): # X-Paygo-Only subscribers should still receive their paygo discount. if paygo_only: ( - po_sub_quota, _, _, _, + po_sub_quota, + _, + _, + _, ) = await InvocationQuota.get_subscription_record(current_user.user_id) if (po_price := get_subscription_tier(po_sub_quota)) is not None: request.state.subscriber_paygo_discount = SUBSCRIPTION_PAYGO_DISCOUNTS.get(