From ce0673929c2520d2944f22a864de1dbf4ac5bc57 Mon Sep 17 00:00:00 2001 From: Aditya kumar singh <143548997+Adityakk9031@users.noreply.github.com> Date: Thu, 27 Aug 2026 13:50:44 +0530 Subject: [PATCH] middleware: exempt subscription turns from spend caps without requiring UsageBypassEnabled --- internal/server/middleware/api_key_spend_cap.go | 11 ++++------- internal/server/middleware/api_key_spend_cap_test.go | 11 +++++++---- internal/server/middleware/org_monthly_spend_cap.go | 7 +++---- .../server/middleware/org_monthly_spend_cap_test.go | 11 +++++++---- 4 files changed, 21 insertions(+), 19 deletions(-) diff --git a/internal/server/middleware/api_key_spend_cap.go b/internal/server/middleware/api_key_spend_cap.go index ad501abc2..7832e884d 100644 --- a/internal/server/middleware/api_key_spend_cap.go +++ b/internal/server/middleware/api_key_spend_cap.go @@ -38,14 +38,11 @@ func WithAPIKeySpendCap(svc *billing.Service) gin.HandlerFunc { return } - // The cap bounds PAID spend, not free subscription usage: a usage-bypass - // org presenting a Claude/Codex credential that covers this route serves at + // The cap bounds PAID spend, not free subscription usage: a request + // presenting a Claude/Codex credential that covers this route serves at // $0 on the caller's own plan, so exempt it from the cap-reached 402 below - // (mirrors WithBalanceCheck). This gate keys off the api key, so read the - // installation from context to check UsageBypassEnabled. - installation := InstallationFrom(c) - subscriptionExempt := installation != nil && installation.UsageBypassEnabled && - proxy.RequestPresentsCoveringSubscription(c.Request.Context(), c.Request.Header, c.FullPath()) + // (mirrors WithBalanceCheck). + subscriptionExempt := proxy.RequestPresentsCoveringSubscription(c.Request.Context(), c.Request.Header, c.FullPath()) result, err := svc.CheckAPIKeySpendCap(c.Request.Context(), apiKey.ID) if err != nil { diff --git a/internal/server/middleware/api_key_spend_cap_test.go b/internal/server/middleware/api_key_spend_cap_test.go index 39a0556aa..c9e8874fc 100644 --- a/internal/server/middleware/api_key_spend_cap_test.go +++ b/internal/server/middleware/api_key_spend_cap_test.go @@ -141,10 +141,13 @@ func TestAPIKeySpendCap_CapReachedNoSubscriptionStillRejected(t *testing.T) { assert.Equal(t, http.StatusPaymentRequired, w.Code) } -func TestAPIKeySpendCap_CapReachedSubscriptionWithoutBypassRejected(t *testing.T) { +func TestAPIKeySpendCap_CapReachedSubscriptionWithoutBypassServesSubscriptionOnly(t *testing.T) { + // Exemption depends only on whether the request presents a covering subscription, + // not on UsageBypassEnabled (matching WithBalanceCheck). repo := &stubBillingRepo{spendFound: true, capMicros: capPtr(1_000_000), spendMicros: 1_000_000} setInstall := func(c *gin.Context) { withInstallation(c, "org_prepaid") } - w, reached, _ := runSpendCapSub(t, "/v1/messages", "k9", setInstall, "Bearer sk-ant-oat-abc123", repo) - assert.False(t, reached, "exemption must not apply without the usage-bypass gate") - assert.Equal(t, http.StatusPaymentRequired, w.Code) + w, reached, subOnly := runSpendCapSub(t, "/v1/messages", "k9", setInstall, "Bearer sk-ant-oat-abc123", repo) + assert.True(t, reached, "a covered turn must pass even when the org lacks the usage-bypass toggle") + assert.Equal(t, http.StatusOK, w.Code) + assert.True(t, subOnly, "the request must be flagged subscription-only") } diff --git a/internal/server/middleware/org_monthly_spend_cap.go b/internal/server/middleware/org_monthly_spend_cap.go index 38ded00e2..b4846476b 100644 --- a/internal/server/middleware/org_monthly_spend_cap.go +++ b/internal/server/middleware/org_monthly_spend_cap.go @@ -36,12 +36,11 @@ func WithOrgMonthlySpendCap(svc *billing.Service) gin.HandlerFunc { return } - // The cap bounds PAID spend, not free subscription usage: a usage-bypass - // org presenting a Claude/Codex credential that covers this route serves at + // The cap bounds PAID spend, not free subscription usage: a request + // presenting a Claude/Codex credential that covers this route serves at // $0 on the caller's own plan, so exempt it from the cap-reached 402 below // (mirrors WithBalanceCheck). - subscriptionExempt := installation.UsageBypassEnabled && - proxy.RequestPresentsCoveringSubscription(c.Request.Context(), c.Request.Header, c.FullPath()) + subscriptionExempt := proxy.RequestPresentsCoveringSubscription(c.Request.Context(), c.Request.Header, c.FullPath()) result, err := svc.CheckOrgMonthlySpend(c.Request.Context(), orgID) if err != nil { diff --git a/internal/server/middleware/org_monthly_spend_cap_test.go b/internal/server/middleware/org_monthly_spend_cap_test.go index d33eef219..1462e4454 100644 --- a/internal/server/middleware/org_monthly_spend_cap_test.go +++ b/internal/server/middleware/org_monthly_spend_cap_test.go @@ -128,12 +128,15 @@ func TestOrgMonthlySpendCap_CapReachedNoSubscriptionStillRejected(t *testing.T) assert.Equal(t, http.StatusPaymentRequired, w.Code) } -func TestOrgMonthlySpendCap_CapReachedSubscriptionWithoutBypassRejected(t *testing.T) { +func TestOrgMonthlySpendCap_CapReachedSubscriptionWithoutBypassServesSubscriptionOnly(t *testing.T) { + // Exemption depends only on whether the request presents a covering subscription, + // not on UsageBypassEnabled (matching WithBalanceCheck). repo := &stubBillingRepo{orgMonthSpent: 1_000_000, orgMonthLimit: capPtr(1_000_000)} setInstall := func(c *gin.Context) { withInstallation(c, "org_prepaid") } - w, reached, _ := runOrgMonthlyCapSub(t, "/v1/messages", setInstall, "Bearer sk-ant-oat-abc123", repo) - assert.False(t, reached, "exemption must not apply without the usage-bypass gate") - assert.Equal(t, http.StatusPaymentRequired, w.Code) + w, reached, subOnly := runOrgMonthlyCapSub(t, "/v1/messages", setInstall, "Bearer sk-ant-oat-abc123", repo) + assert.True(t, reached, "a covered turn must pass even when the org lacks the usage-bypass toggle") + assert.Equal(t, http.StatusOK, w.Code) + assert.True(t, subOnly, "the request must be flagged subscription-only") } func TestOrgMonthlySpendCap_OverridePassesThrough(t *testing.T) {