diff --git a/internal/gateway/fallback.go b/internal/gateway/fallback.go index 3b68f92d..2a734a77 100644 --- a/internal/gateway/fallback.go +++ b/internal/gateway/fallback.go @@ -94,15 +94,26 @@ func (c candidate) isOpenRouterFree() bool { return c.p.Preset == "openrouter" && strings.HasSuffix(c.model, ":free") } -// restID is a free OpenRouter model's own rest key. A rate limit on that -// model is its free-tier limit, while an account-level rest stays on restKey. +// restID is the candidate's model's own rest key: a free OpenRouter +// model's, whose rate limit is its free-tier limit, and a subscription's, +// out of a pool of its allowance that counts some models only (pooled). +// An account-level rest stays on restKey. func (c candidate) restID() string { - if c.isOpenRouterFree() { + if c.isOpenRouterFree() || c.p.Account != nil { return c.restKey() + "/" + c.model } return c.restKey() } +// pooled says whether the candidate's subscription, refused for its +// allowance, is out for its model alone (provider.Allowance.Pooled). +func (c candidate) pooled(now time.Time) bool { + if c.p.Account == nil { + return false + } + return allowances(c.p.Account.UsageAgent())[c.p.Account.User].Pooled(c.model, provider.SpentShareOf(c.p.Routing), now) +} + // who is the key or account itself, however many the provider has on: a // provider's one key rests as the provider, and as itself once another // is added, and a conversation it answered stays with it all the same. diff --git a/internal/gateway/pool_rest_test.go b/internal/gateway/pool_rest_test.go new file mode 100644 index 00000000..5a0d39c7 --- /dev/null +++ b/internal/gateway/pool_rest_test.go @@ -0,0 +1,59 @@ +package gateway + +import ( + "net/http" + "testing" + "time" + + "github.com/yetone/magpie/internal/provider" +) + +// A subscription refused for a pool of its allowance that counts some +// models only rests for that model alone: the account still takes the +// models the pool doesn't count (Xiaopodev on X: Cursor's Other Models +// used up kept Auto off the account too). One of the whole account's +// windows used up rests the account. +func TestPoolRefusalRestsTheModelOnly(t *testing.T) { + old := allowances + defer func() { allowances = old }() + fresh(t) + now := time.Now() + week := now.Add(39 * time.Hour) + allowances = func(string) map[string]provider.Allowance { + return map[string]provider.Allowance{ + "me@x.com": { + {Used: 100, Resets: week, Span: 7 * 24 * time.Hour, Model: "opus"}, + {Used: 40, Resets: week, Span: 7 * 24 * time.Hour}, + }, + "out@x.com": {{Used: 100, Resets: week, Span: 7 * 24 * time.Hour}}, + } + } + acct := func(user, model string) candidate { + return candidate{p: provider.Provider{ID: "claude", Account: &provider.Account{Agent: "claude", User: user}}, model: model, rest: "claude"} + } + s := &Server{} + refused := []byte(`{"error":{"message":"usage limit reached for this model"}}`) + + opus, sonnet := acct("me@x.com", "claude-opus-4-7"), acct("me@x.com", "claude-sonnet-4-6") + r := s.restAfter(opus, 429, http.Header{}, refused) + if r.Why != failQuota || r.Key != opus.restID() || r.By != "window" { + t.Fatalf("Opus out of its week rests %+v", r) + } + if _, ok := restOf(sonnet.restKey()); ok { + t.Fatal("Opus out of its own week rests the account") + } + if _, ok := restOf(sonnet.restID()); ok { + t.Fatal("Opus out of its own week rests Sonnet") + } + if !spentAfter([]candidate{opus}) || spentAfter([]candidate{sonnet}) { + t.Fatal("spentAfter doesn't go by the model's rest") + } + + whole := acct("out@x.com", "claude-sonnet-4-6") + if r := s.restAfter(whole, 429, http.Header{}, refused); r.Key != whole.restKey() { + t.Fatalf("the account's week used up rests %s, not the account", r.Key) + } + if _, ok := restOf(acct("out@x.com", "claude-opus-4-7").restKey()); !ok { + t.Fatal("the account's week used up leaves its other models in") + } +} diff --git a/internal/gateway/routing.go b/internal/gateway/routing.go index 6ebff614..44ee919b 100644 --- a/internal/gateway/routing.go +++ b/internal/gateway/routing.go @@ -332,6 +332,10 @@ func (s *Server) restAfterMarked(c candidate, status int, header http.Header, bo now := time.Now() d := fallbackCooldown why := failureOf(c, status, body) + // read before the allowance is asked again below: out of a pool that + // counts some models only, the account's others are still in theirs + // (Cursor's Auto with Other Models used up) + pooled := why == failQuota && c.pooled(now) r := Rest{Why: why, Status: status, By: "cooldown"} if why == failProxy { // the account is as good as it was; the proxy is the user's to start @@ -402,6 +406,9 @@ func (s *Server) restAfterMarked(c candidate, status int, header http.Header, bo if why == failRate && c.isOpenRouterFree() && sharedPool { id = c.restID() } + if pooled { + id = c.restID() + } r.Key = id restingUntil.Lock() restingUntil.m[id] = r.Until diff --git a/internal/provider/allowance_pooled_test.go b/internal/provider/allowance_pooled_test.go new file mode 100644 index 00000000..895ff074 --- /dev/null +++ b/internal/provider/allowance_pooled_test.go @@ -0,0 +1,50 @@ +package provider + +import ( + "testing" + "time" + + "github.com/yetone/magpie/internal/plugin" +) + +// A refusal for a model counted by a pool of some models only is that +// model's, not the account's: Cursor's Other Models used up leaves Auto in +// Cursor Models (Xiaopodev on X), Opus's own week leaves Sonnet; one of +// the whole account's windows used up is the account's. +func TestAllowancePooled(t *testing.T) { + now := time.Now() + month := now.Add(20 * 24 * time.Hour) + pool := map[string]bool{"auto": true, "composer-2": true} + q := quotaOfPlugin(SubscriptionQuota{}, plugin.Usage{Windows: []plugin.UsageWindow{ + {Name: "Cursor Models", Used: 12, Models: []string{"auto", "composer-2"}}, + {Name: "Other Models", Used: 100, NotModels: []string{"auto", "composer-2"}}, + {Name: "Total", Used: 60, Aside: true}, + }}) + cursor := allowanceOf(q.Windows, now) + for _, m := range []string{"auto", "composer-2", "claude-4.6-sonnet", "gpt-5.5"} { + if !cursor.Pooled(m, 100, now) { + t.Errorf("Cursor: %s refused isn't the model's alone", m) + } + if used, _ := cursor.For(m, now); (used == 100) == pool[m] { + t.Errorf("Cursor: %s used %v", m, used) + } + } + + claude := Allowance{ + {Used: 100, Resets: month, Span: 7 * 24 * time.Hour, Model: "opus"}, + {Used: 40, Resets: month, Span: 7 * 24 * time.Hour}, + } + if !claude.Pooled("claude-opus-4-7", 100, now) { + t.Error("Opus's own week used up rests the account") + } + if claude.Pooled("claude-sonnet-4-6", 100, now) { + t.Error("Sonnet, refused with no pool of its full, counts as pooled") + } + claude[1].Used = 100 + if claude.Pooled("claude-opus-4-7", 100, now) { + t.Error("the account's week used up rests only Opus") + } + if (Allowance{}).Pooled("auto", 100, now) { + t.Error("an allowance not known counts as pooled") + } +} diff --git a/internal/provider/routing.go b/internal/provider/routing.go index f8f734c3..64d58db1 100644 --- a/internal/provider/routing.go +++ b/internal/provider/routing.go @@ -265,6 +265,33 @@ func (a Allowance) Full(model string, share float64, now time.Time) time.Time { return t } +// Pooled says whether a refusal of model for its allowance leaves the +// account's other models alone: the windows that count model and are full +// at share are each of some models only (Opus's own week, Cursor's Other +// Models pool) — or, none known full, no window of the whole account +// counts it, only pools. Then it is that model which is out, not the +// account: Cursor's Other Models used up leaves Auto and Composer in +// theirs (Xiaopodev on X). +func (a Allowance) Pooled(model string, share float64, now time.Time) bool { + model = strings.ToLower(model) + full, pooled, whole := false, false, false + for _, l := range a { + if !l.applies(model) { + continue + } + scoped := l.Model != "" || l.matches != nil + if l.Used >= share && (l.Resets.IsZero() || l.Resets.After(now)) { + if !scoped { + return false + } + full = true + } + pooled = pooled || scoped + whole = whole || !scoped + } + return full || pooled && !whole +} + // budgetSpan is the shortest window that is a budget rather than a rate // cap: the week (Kiro's month too), not the five hours in it — what the // five hours leave at their reset is nothing lost.