Skip to content

middleware: exempt subscription turns from spend caps without requiring UsageBypassEnabled #1074

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Adityakk9031 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

middleware: exempt subscription turns from spend caps without requiring UsageBypassEnabled #1074

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 4 additions & 7 deletions internal/server/middleware/api_key_spend_cap.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
11 changes: 7 additions & 4 deletions internal/server/middleware/api_key_spend_cap_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
}
7 changes: 3 additions & 4 deletions internal/server/middleware/org_monthly_spend_cap.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
11 changes: 7 additions & 4 deletions internal/server/middleware/org_monthly_spend_cap_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down