From b0a96cbbd7b30f6ad1696b85693bcda5b4eaca3e Mon Sep 17 00:00:00 2001 From: forever Date: Tue, 6 Oct 2026 03:57:47 -0700 Subject: [PATCH] gateway, provider, gui: a Sub2API key given 5h/1d/7d limits is weighed and rested by its own windows, like an account (#1004) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * gateway: a Sub2API key out of its 5h/1d/7d limit ("…限额已用完") rests as out of quota, not a 1-minute rate limit Sub2API refuses a key that has spent the 5-hour, day or 7-day limit its owner gave it with 429 "api key 7天限额已用完" (or 日限额, 5小时限额), typed rate_limit_exceeded. It matched rateWords and not plannedWords ("限额" is not "额度"), so failure() said rate: the key rested a minute, was tried again, refused again, and sank, for the days its 7-day window had left. plannedWords now has 限额已用完, so it is failQuota. The relay's other refusals keep their classes, fixed by the test: the pool out of accounts (v0.1.149's 503 "No available accounts": other; later versions' 429 "All available accounts are currently rate-limited": rate), its upstream's rate limit (rate), the key's own total quota API_KEY_QUOTA_EXHAUSTED (quota) and the wallet's 403 INSUFFICIENT_BALANCE (credit). Verified: TestSub2APIKeyLimitIsQuota, with Sub2API v0.1.149's message bodies. Without the change it fails: "429 {…api key 7天限额已用完…}: rate, want quota", and the same for 日限额 and 5小时限额. * provider: a Sub2API key's /v1/usage shows its 5 hours / 1 day / 7 days windows instead of nothing at "remaining" A Sub2API key its owner gave a 5-hour, day or 7-day spending limit answers /v1/usage in quota_limited mode with rate_limits — each window's limit, used and reset_at in dollars — and, without a total quota of its own, no "remaining". With the Balance URL set to it, the key's card said `nothing at "remaining" in the reply`. readSub2APIKeyLimits reads rate_limits into the windows the cards already name 5 hours, 1 day and 7 days, in USD. A window not started (used 0, no reset_at) is taken to renew a span from now; a window with a limit of 0 is none. The card shows them, beside a balance field the user wrote; with none, Check balance and `magpie provider show` give what the fullest window has left. A wallet key's reply is unchanged. KeyAllowance gives routing the same windows, per provider and key: read behind the request, never waited for, at most once a minute (once for any number of asks at a time); a failed read keeps the last; after a restart it is seeded from the card's last reading in quotas.json; StaleKeyAllowance has it read again at the next ask. ForgetBalances (a key or Balance URL changed) has every key read again, a read begun before dropped, the last reading standing until then. Only keys whose Balance URL is a Sub2API /v1/usage are asked: any other key, and every account, costs nothing here. Verified with real bytes of a limited key's reply, its key removed (testdata/sub2api_usage_key_limits*.json): TestReadSub2APIKeyLimits, TestSub2APIKeyLimitsOnItsCard, TestKeyAllowance*. Each fails without its part, e.g. the card without the windows: `Error:nothing at "remaining" in the reply`; KeyAllowance waiting for the read: "waited 30.0s"; asking every time: "asked 11 times"; a failed read asked again at once: "asked 22 times after failing"; not seeded: "seeded = [], false"; any key asked: "asked /api/usage/token"; ForgetBalances dropping what was read: "after ForgetBalances, and a read begun before it: []"; leaving a read under way loading: "the key's windows were never read"; a stale-generation read landing: it replaced 95%; a limit of 0 taken: "a day's limit of 0 = [{… Used:100 …}]"; Balance from the first window only: "$300.00" for $80 left. * gateway, gui: keys whose windows magpie reads are weighed and rested by them, like accounts Smart, least-used and pace routing weighed only subscriptions by their windows: weighRouted skipped every key (c.p.Account == nil), so a Sub2API key beside Codex accounts in a group went by its place in the order, and, refused for its 7-day limit, rested a quarter of an hour at a time while the window had days to run. A key whose own windows magpie reads (KeyAllowance: a Sub2API key given a 5h/1d/7d limit, its Balance URL /v1/usage) is now weighed as an account is: fine, low (90%) or spent (98%), the one renewing soonest first, in a provider and in a group alike, and its trace row carries them. Out of quota it rests till the window it filled renews, by candidate.full; when its windows as last read aren't full it rests quotaRest, and either way they are read again (staleKeyAllowance). Every other failure of a key backs off as before, never to the window: the relay out of accounts for everyone (503/429) says nothing of the key's own spend. capHeld and pooled stay accounts' only. Keys magpie reads no windows of are ordered exactly as before. The Routing page tells such a key as it does an account: its row's bar, used share and renewal, why it went first by its allowance, "Every key and account here is at 90% or more", and the Smart header. A key not read yet beside one read is told as weighed with it, not "keys go in their order" or by its tokens. New strings in en, zh, ja and de. Other users of the relay draining its accounts don't show in a key's windows, which count only its spend: they show as the pool's 503/429, which back off and fall over to the next key, account or fallback. Verified: TestKeysWeighedByTheirWindows, TestGroupWeighsKeyWindows, TestKeyOutOfItsWindowRestsUntilReset, TestKeyPoolErrorDoesNotRestForWindow; and routing-key-windows.test.cjs in Chromium and WebKit (en/zh/ja/de). Without the change they fail: "smart: kw#a kw#b kw#c"; "the key renewing first: codex sub2api"; "full: rests 14m59.99s by quota"; "read again: []"; backing off to the window: "pool empty: rests 71h59m59.99s by window (other)"; and on the page "relay-a goes first: keys go in their order, those that suit the request first." where all are at 90% or more, or a key not read yet beside one read. * gateway, provider: a Sub2API key out of its limit is back once its windows read short of full, not only when the window renews A key refused for its 5h/1d/7d limit rests till the window it filled renews: days, for the week. Its windows are read again behind the next requests (a minute apart) and by its card, but nothing lifted the rest when they came back short of full, so the owner raising the key's limit in the relay's panel, or resetting its usage, left the key out till the week turned, though the relay would take it at once. Accounts have this already: a Codex reset spent tells OnRenewed, and the gateway lifts the account's rests (renewed). A key's reading, routing's or its card's, now tells OnRenewed the same way, as agent "" and the key's KeyAllowanceID, when a window it was full in (Allowance.Full at SpentShareOf its routing, as candidate.full counts it) is full no more, or full till sooner. Its first reading, a failed read, or one that finds it fuller tells nothing. The gateway notes the key's KeyAllowanceID on its rest out of quota only, so renewed lifts that rest and no other: the relay out of accounts for everyone (503) still backs off, and another key still out of its week stays out. Where nothing weighs the key by its windows, nothing read them again either: an ordered group weighs each member's keys alone, so a member with one key, like a provider with one key, never asks for them. After the one read its 429 starts, the key sat out its week on that reading. restingOf, which plans a resting key behind the rest, now asks for the windows of a key resting out of them (agent "" and a user on its rest), so the next request planned past a minute reads them, behind it, and a reading short of full brings the key back there too. No card is read for it, and a key resting for any other reason asks nothing. Verified: TestKeyAllowanceRenewed (provider) and TestKeyBackOnceItsWindowIsNotFull (gateway, end to end through the key cards' reads and restAfter), each failing with a piece of the change undone: never telling, "its week's limit raised to $1,000: told []" and "a, its limit raised: still rests by window till …"; telling on any change, "its first reading, out of its week: told [/sub2api#…]"; a failed read telling, "a read that failed: told [… …]"; the card not telling, "its usage reset, read on its card: told [/sub2api#…]"; the rest not noted, "a, its limit raised: still rests by window"; every rest of a key noted, "c, the relay out of accounts: …, false"; renewed not matching the key, "b, still out of its week: …, false". And TestRestingKeyReadAgainInAnOrderedGroup (gateway: an ordered group, its first member one key on a fake /v1/usage, no card read after the 429), failing without restingOf's ask: "its limit raised: still rests by window till …, its windows read 1 times". Tried on a real Sub2API v0.1.149 relay, in a sandbox HOME, its key's 7d limit lowered under its spend for the test and then restored, with magpie serve on a group ordered: the key, then a local stand-in. The key answered 429 "api key 7天限额已用完"; it rested by window till its 7d reset, 2026-10-13 00:00 +08:00, and the stand-in served. Planned again a minute later, still over, it still rested. With the limit back at $800, the next request was planned with it resting, the stand-in serving, and read its windows; the one after went to the key, 200. --- docs/reference.md | 8 +- docs/subsystems/gateway-routing.md | 8 +- docs/subsystems/providers-accounts.md | 6 +- internal/gateway/fallback.go | 26 +- internal/gateway/key_limits_test.go | 345 ++++++++++++ internal/gateway/routing.go | 86 ++- internal/gateway/routing_test.go | 8 +- internal/gui/assets/i18n.js | 3 + internal/gui/assets/routing.js | 33 +- .../gui/tests/routing-key-windows.test.cjs | 168 ++++++ internal/provider/balance.go | 64 ++- internal/provider/key_allowance.go | 202 +++++++ internal/provider/routing.go | 10 +- internal/provider/sub2api_usage.go | 61 ++ internal/provider/sub2api_usage_test.go | 533 ++++++++++++++++++ .../testdata/sub2api_usage_key_limits.json | 1 + .../sub2api_usage_key_limits_not_started.json | 1 + 17 files changed, 1493 insertions(+), 70 deletions(-) create mode 100644 internal/gateway/key_limits_test.go create mode 100644 internal/gui/tests/routing-key-windows.test.cjs create mode 100644 internal/provider/key_allowance.go create mode 100644 internal/provider/sub2api_usage.go create mode 100644 internal/provider/sub2api_usage_test.go create mode 100644 internal/provider/testdata/sub2api_usage_key_limits.json create mode 100644 internal/provider/testdata/sub2api_usage_key_limits_not_started.json diff --git a/docs/reference.md b/docs/reference.md index c9dfe9db..e00772df 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -674,7 +674,13 @@ spare, the one whose allowance renews soonest first), `order` (the first model until it can't answer, then the next), `rotate` (each turn to the next member), `usage` (least used first) or `pace` (weekly pace: the account with the most of its week left per hour until it renews first, so less of a week -is lost at its reset). `stays=` is how long a conversation +is lost at its reset). A key has no allowance to weigh, so it goes by its +order or by what magpie sent it lately — except a sub2api key its owner gave +a 5-hour, day or 7-day limit: with the provider's Balance URL set to the +relay's `/v1/usage`, its card shows those windows and routing weighs it by +them as it does a subscription, resting it till the window renews once the +relay says `api key 7天限额已用完` — or till its windows are next read short +of full, its limit raised or its usage reset. `stays=` is how long a conversation stays with the key or account that answered it: `auto` (the default, while the vendor's cache of it is worth keeping), `session`, `turn` or `off`. `models=` replaces the whole list, in order; a bare model id works when only diff --git a/docs/subsystems/gateway-routing.md b/docs/subsystems/gateway-routing.md index d9e5804f..8740c82b 100644 --- a/docs/subsystems/gateway-routing.md +++ b/docs/subsystems/gateway-routing.md @@ -19,7 +19,7 @@ served by a plugin are in [Provider and plugin ownership](provider-plugins.md). | HTTP surface | `Handler` routes `/v1/chat/completions`, `/v1/responses`, `/v1/messages` (each also without `/v1`), `/v1/models`, Gemini's `/v1beta/models/…`, the Codex backend path, `/mcp/{name}` and the `/v1/magpie/*` status endpoints. A GET on `/v1/responses` (a client opening Responses over a WebSocket) is answered 426, a plain one 405: Responses is relayed over HTTP only. Every route is wrapped in `counted`, `callerGuard`, `withCaller` and `keyLimited`. | [`gateway.go`](../../internal/gateway/gateway.go) | | Model resolution | `serve` resolves the model the agent asked for: Claude tier stand-ins, a routing group (`provider.GroupFor`), auto stand-ins, then `provider.Resolve`. An unknown model gets a 404 naming why: a switched-off provider, a disabled group or an empty group. | [`gateway.go`](../../internal/gateway/gateway.go), [`internal/provider/group.go`](../../internal/provider/group.go) | | Candidates | `plan` lists one provider's candidates: one per key or account on, then each `Fallback` model, with resting ones last (`restLast`). `planGroup` and `planLevel` weigh a group's members together. A nested group plans by its own routing (#576). A member's own fallbacks are not the group's. | [`fallback.go`](../../internal/gateway/fallback.go) | -| Order of keys and accounts | `Routing`: smart (""), `order`, `rotate`, `usage` (least used, with an `usageHalfLife` of 1h), `pace` or `weight` (`weightedFirst`). A failure rests a candidate (`failure`, `failureOf`). | [`routing.go`](../../internal/gateway/routing.go), [`weighted.go`](../../internal/gateway/weighted.go), [`internal/provider/routing.go`](../../internal/provider/routing.go) | +| Order of keys and accounts | `Routing`: smart (""), `order`, `rotate`, `usage` (least used, with an `usageHalfLife` of 1h), `pace` or `weight` (`weightedFirst`). Smart, `usage` and `pace` weigh a subscription by its allowance (`allowances`), and a key by its own usage windows when magpie reads them (`keyAllowance`: a sub2api key given a 5-hour, day or 7-day limit); any other key has none. A failure rests a candidate (`failure`, `failureOf`). | [`routing.go`](../../internal/gateway/routing.go), [`weighted.go`](../../internal/gateway/weighted.go), [`internal/provider/routing.go`](../../internal/provider/routing.go) | | Sink | With `Sink` on, a key or account that answers 429 while it still has quota goes to the back, behind every one not rate limited since. Held in memory until restart. | [`sink.go`](../../internal/gateway/sink.go) | | Affinity | A conversation stays with the key or account that answered it, so the vendor's prompt cache is read again. Modes: auto (""), `session`, `turn`, `off`. Stored in `affinity.json` next to `providers.json`. | [`affinity.go`](../../internal/gateway/affinity.go) | | One try | `attempt` sends one candidate the request. It goes through the provider's proxy (`p.Via`). A subscription with its own client is served specially, such as Claude through the genuine Claude Code. The request is relayed as is (`passthrough`) when the model is usable on the agent's API, except that a Claude model where the provider has an Anthropic URL is translated to Messages from Chat or Responses too (`OnMessages`: OpenCode speaks Chat for every model, and a relay drops cache_control there), unless the user set the model's API or the vendor's list names its APIs. A request offering web search to a provider that doesn't search by itself is translated too, so magpie can search for it (`search.go`), but on the agent's own API where the provider serves it and no `ResponsesFirst`/`MessagesFirst` rule prefers another: Codex offers web_search on every turn, and each went out as Chat (#997). Otherwise it is translated (`translate`) to the first API `usable` gives: Chat, except Responses for OpenAI's models where their maker serves them (`ResponsesFirst`) and Anthropic's Messages for a Claude model where the provider has an Anthropic URL (`MessagesFirst`, which keeps its cache breakpoints, #997). `Provider.Native`, the API agents with a per-model API (Pi, omp, droid) are wired to ask on, follows the same order. | [`gateway.go`](../../internal/gateway/gateway.go) | @@ -47,6 +47,8 @@ served by a plugin are in [Provider and plugin ownership](provider-plugins.md). - In turn (`rotate`), a conversation nobody has answered yet goes to the key, account or model the fewest other conversations of the provider's model or the group answered in the last `heldFor` (30 minutes) are on (`leastHeld`, #946). Ties, resting ones and a group with a group in it keep routing's order. With session affinity, two sessions of one group so stay on different models. - A ChatGPT account that holds credits answers past its windows on them; the vendor spends them. That is the default. Set not to spend them (Usage card's *Use credits*, `magpie quota credits off`), the account is held at 100% of a window till it renews, by magpie's latest reading of its windows (`provider.Allowances`, refreshed about every minute), so the other accounts, groups and fallbacks take the request. A cap set on the account (1–99%) takes precedence. Unlike a cap, the hold is at the vendor's own 100%, so an account allowed to spend its resets spends one there, both when every candidate is held and after the last one is refused. A reset made forgets the account's cached allowance (`renewedNow`), so neither hold outlives the windows it started again. - Different failures are classified separately: out of credit, out of quota, rate limit and other errors rest for different lengths. A fix must not merge them into one bucket (see #873). +- A sub2api key out of its own limit answers 429 `api key 7天限额已用完` (or `日限额`, `5小时限额`) as `rate_limit_exceeded`. It is out of quota (`plannedWords`), not rate limited: it rests till the window it filled renews when its windows as last read say so, else `quotaRest`, and its windows are read again (`staleKeyAllowance`). A reading that finds that window full no more (the limit raised in the relay's panel, or the key's usage reset) lifts the rest (`renewed`, told through `provider.OnRenewed`), as a Codex reset spent does an account's; the key's other rests stay. While it rests out of its windows, they are read again each time it is planned (`restingOf`, once a minute at most), so this holds where nothing weighs it by them: an ordered group, or a provider with one key. The relay out of accounts for everyone (503 `No available accounts`, or a later version's 429 `All available accounts are currently rate-limited`) says nothing of the key's windows: it backs off as any failure or rate limit does, never to the window. A key's windows count only its own spend, so other users draining the relay's accounts show up as those errors, not in the windows. +- Only an account is held at a usage cap (`capHeld`) or rested out of one model's pool (`pooled`); a key's windows order it and rest it, nothing more. - Codex offers its image tool (the `image_gen` namespace, #870) to every magpie provider. When a provider other than a ChatGPT sign-in or openai.com answers a bare 400, `attempt` asks it once more without the tool (`image_tool.go`). This happens both passed through and translated. If that works, the provider is marked unfit for the tool (`tool:image_gen`) and isn't offered it again. If it doesn't, the tool is restored and the 400 stands (#949). - A translated Responses reply gives Codex reasoning items magpie made: an `rs_` id and a summary, nothing sealed. Codex hands them back on every turn. A request relayed as is to OpenAI's API or Azure OpenAI goes without them (`withoutBareReasoning`) unless `store` is true, because both look the id up among stored items and answer 400 "Item with id 'rs_…' not found" (#1008). Codex's compaction offers no tools, so it is relayed even when its turns were translated. When a provider answers its summary request 404, or 400 for an item it didn't find (`itemNotFound`, a relay in front of either), `codexCompact` asks once more with the conversation as text (`plainCompact`), then compacts locally (#866, #1008). - Group members are `provider/model[:effort]`. A member with a fixed effort is asked for that effort, bounded by the group's levels (#671). @@ -55,8 +57,8 @@ served by a plugin are in [Provider and plugin ownership](provider-plugins.md). ## Verification ```sh -go test -tags nogui ./internal/gateway -run 'TestQuiet|TestSilentHeld|TestRotateSpreadsSessions|TestRouting|TestSmartRouting|TestFallback|TestNoFallbackForOtherErrors|TestLastFallbackErrorReachesTheAgent|TestRateLimitedSinksToTheBack|TestGroupRateLimitedSinks|TestAffinity|TestSeveralKeysOnTakeOverFromEachOther|TestSubscriptionAccountsTakeOver|TestGroup|TestTrace|TestImageTool|TestCodexNoCredits|TestClaudeGoesOnMessagesWhereTheRelayHasThem|TestAutoRelaysOnTheClientsOwnAPI|TestSearchOfferedStaysOnTheClientsAPI|TestCodexSearchAnsweredOnResponses|TestAgentEffortOnCursorLocal|TestAzureCompactionWithoutMagpiesReasoning|TestCodexCompact|TestNotAnAPIReply|TestAPIReplyUntouched|TestCodexBackendWebPage' -go test -tags nogui ./internal/provider -run 'TestGroup|TestCodexCreditsSwitch|TestRenewedAccountForgetsItsAllowance|TestNative' +go test -tags nogui ./internal/gateway -run 'TestQuiet|TestSilentHeld|TestRotateSpreadsSessions|TestRouting|TestSmartRouting|TestFallback|TestNoFallbackForOtherErrors|TestLastFallbackErrorReachesTheAgent|TestRateLimitedSinksToTheBack|TestGroupRateLimitedSinks|TestAffinity|TestSeveralKeysOnTakeOverFromEachOther|TestSubscriptionAccountsTakeOver|TestGroup|TestTrace|TestImageTool|TestCodexNoCredits|TestClaudeGoesOnMessagesWhereTheRelayHasThem|TestAutoRelaysOnTheClientsOwnAPI|TestSearchOfferedStaysOnTheClientsAPI|TestCodexSearchAnsweredOnResponses|TestAgentEffortOnCursorLocal|TestAzureCompactionWithoutMagpiesReasoning|TestCodexCompact|TestNotAnAPIReply|TestAPIReplyUntouched|TestCodexBackendWebPage|TestSub2APIKeyLimitIsQuota|TestKeysWeighedByTheirWindows|TestGroupWeighsKeyWindows|TestKeyOutOfItsWindowRestsUntilReset|TestKeyPoolErrorDoesNotRestForWindow|TestKeyBackOnceItsWindowIsNotFull|TestRestingKeyReadAgainInAnOrderedGroup' +go test -tags nogui ./internal/provider -run 'TestGroup|TestCodexCreditsSwitch|TestRenewedAccountForgetsItsAllowance|TestNative|TestKeyAllowance' ``` The gateway's `TestMain` gives the package a home of its own, so these tests diff --git a/docs/subsystems/providers-accounts.md b/docs/subsystems/providers-accounts.md index ecb5b59c..96e9900e 100644 --- a/docs/subsystems/providers-accounts.md +++ b/docs/subsystems/providers-accounts.md @@ -28,6 +28,7 @@ keys and accounts per request is in [Gateway routing and fallback](gateway-routi | Used-up switch | When the account Codex or Claude Code is signed in to has used its allowance, the agent is signed in to the next account that is on and has room (#208). | [`codex_switch.go`](../../internal/provider/codex_switch.go) | | Routing groups | User-made groups from the Routing view, and auto-found groups of the same model across providers (derived, never stored). Members are `provider/model[:effort]`. A group can nest another (`group/`); `SaveGroup` refuses loops. `Match` takes a glob or `re:` pattern (#766). | [`group.go`](../../internal/provider/group.go) | | Allowance | `Quotas` asks every vendor at once: subscription windows (`SubscriptionUsage`), key-bought plans (`PlanQuotas`) and key balances (`KeyBalances`). The Usage page, `magpie quota --json` and `GET /v1/magpie/quotas` read it. A single card can be read again on its own (#840). A remote magpie's cards come from what it has kept (`CachedCards`, `GET /v1/magpie/quotas/cards`) and are never read there for this computer except by a card's refresh (`POST /v1/magpie/quotas/refresh`, throttled per card). | [`quotas.go`](../../internal/provider/quotas.go), [`subscription_usage.go`](../../internal/provider/subscription_usage.go), [`planquota.go`](../../internal/provider/planquota.go), [`balance.go`](../../internal/provider/balance.go), [`quota_refresh.go`](../../internal/provider/quota_refresh.go), [`remote_quotas.go`](../../internal/provider/remote_quotas.go) | +| A key's own windows | A sub2api key given a spending limit per 5 hours, day or 7 days answers its `/v1/usage` with `rate_limits` and no `remaining` (`readSub2APIKeyLimits`). With the Balance URL set to it, the key's card shows them as the windows *5 hours*, *1 day*, *7 days*, in USD, beside a balance field the user wrote; *Check balance* and `magpie provider show` give what its fullest window has left (`Balance`). `KeyAllowance` gives routing the same windows: never waited for, read in the background at most once a minute, a failed read keeping the last, and seeded from the card's last reading (`quotas.json`) after a restart. | [`sub2api_usage.go`](../../internal/provider/sub2api_usage.go), [`key_allowance.go`](../../internal/provider/key_allowance.go) | | Last served | `served.json` records when each account, plan or key last answered, shown as `lastServedAt` (#570). | [`served.go`](../../internal/provider/served.go) | | Protocol detection | Works out which APIs a base URL speaks when a provider is added. | [`detect.go`](../../internal/provider/detect.go) | @@ -36,7 +37,7 @@ keys and accounts per request is in [Gateway routing and fallback](gateway-routi 1. **Add.** The Providers page adds a key provider through `Add` or `Save`, or an account through `StartSignIn`. An agent signed in on its own is found by `Accounts` with no step from the user. 2. **Resolve.** The gateway calls `provider.Resolve` or `GroupFor` with the model name, then plans candidates over the provider's keys or accounts. 3. **Sign.** For each try, `Sign` gives the request its auth. A token about to expire is refreshed first, and the rotated tokens go back to their one holder: the agent's store, or `logins.json` for an account standing behind it. -4. **Allowance.** Usage endpoints are rate limited, so `SubscriptionUsage` caches its results: a cached copy comes back at once and a stale one is refreshed in the background. When an account's windows start again, `OnRenewed` tells the gateway so a resting account can come back. +4. **Allowance.** Usage endpoints are rate limited, so `SubscriptionUsage` caches its results: a cached copy comes back at once and a stale one is refreshed in the background. When an account's windows start again, `OnRenewed` tells the gateway so a resting account can come back. A key's reading that finds a window it was full in (at `SpentShareOf` its routing) full no more, or full till sooner, tells it too, as agent `""` and the key's `KeyAllowanceID`, so a key out of its limit is back once the limit is raised or its usage reset; a first reading, a failed one, or one that finds it fuller tells nothing. ## Constraints and failure behavior @@ -47,12 +48,13 @@ keys and accounts per request is in [Gateway routing and fallback](gateway-routi - A switched-off provider (`Off`, #163) stays saved. The gateway says it is off instead of saying the model is unknown. - A provider removed or switched off takes its models out of the catalog. [Agent wiring](agent-wiring.md)'s `Reseat` moves agents off them. - A gateway key held to some accounts (#905) is held to those of the providers its list names: a provider it names no account of, the key uses as it always did. The list keeps an account by its stable id — `logins.json`'s `id`, set once written and kept through renames; the agents that keep no logins (and plugins) are known by a hash of the account's name, so an account renamed there drops out of a key's list, holding it closer, never wider. Storing an id for every agent is a larger change, not this one's. +- A sub2api key's windows are fixed: each starts at the key's first request after it was empty, aligned to the relay's day, and a window not started yet says `used 0` with no `reset_at`, so it is taken to renew a span from now. They count this key's spend only, not the relay's accounts behind it. `ForgetBalances` (a key or Balance URL changed) has every key's windows read again at its next ask, a read begun before dropped; until then the last reading stands, fresher than the card's, which routing's reads never write. A key changed is another key (`keyID`), with windows of its own. - Whether a Codex account spends its credits (`CodexNoCredits` in settings, opt-out, by lower-cased email) is the gateway's routing only. It never changes which account Codex is signed in to: `NextLogin` moves off an account at its routing's spent share, 100% at most, credits or not. ## Verification ```sh -go test -tags nogui ./internal/provider -run 'TestAccount|TestCodex|TestGroup|TestQuota|TestLogin|TestSignIn|TestSeveralKeys|TestKeyProtocol|TestSetKeyWeight' +go test -tags nogui ./internal/provider -run 'TestAccount|TestCodex|TestGroup|TestQuota|TestLogin|TestSignIn|TestSeveralKeys|TestKeyProtocol|TestSetKeyWeight|TestReadSub2APIKeyLimits|TestSub2APIKeyLimitsOnItsCard|TestKeyAllowance' go test -tags nogui ./internal/provider ``` diff --git a/internal/gateway/fallback.go b/internal/gateway/fallback.go index bd61472e..08c269d2 100644 --- a/internal/gateway/fallback.go +++ b/internal/gateway/fallback.go @@ -627,9 +627,7 @@ func asideOf(cs []candidate, q provider.Provider, fallback bool, from provider.P // restLast moves those resting after a recent failure behind the rest. func restLast(out []candidate, pl planned) ([]candidate, planned) { if len(out) == 1 { - if r, ok := restOf(out[0].restKey()); ok { - pl.order[0].Rest = &r // tried all the same: there is no other - } else if r, ok := restOf(out[0].restID()); ok { + if r, ok := restingOf(out[0]); ok { pl.order[0].Rest = &r // tried all the same: there is no other } return out, pl @@ -637,11 +635,7 @@ func restLast(out []candidate, pl planned) ([]candidate, planned) { var ready, resting []candidate var wReady, wResting []Weighed for i, c := range out { - r, ok := restOf(c.restKey()) - if !ok && c.restID() != c.restKey() { - r, ok = restOf(c.restID()) - } - if ok { + if r, ok := restingOf(c); ok { pl.order[i].Rest = &r resting, wResting = append(resting, c), append(wResting, pl.order[i]) } else { @@ -652,6 +646,22 @@ func restLast(out []candidate, pl planned) ([]candidate, planned) { return append(ready, resting...), pl } +// restingOf is why c rests, by its own key or else its provider's, while +// it does. A key resting out of its own windows has them read again as it +// is planned, behind the request and once a minute at most: an ordered +// group, or a provider with one key, weighs nothing by them, and a +// reading that finds them full no more is what brings it back (renewed). +func restingOf(c candidate) (Rest, bool) { + r, ok := restOf(c.restKey()) + if !ok && c.restID() != c.restKey() { + r, ok = restOf(c.restID()) + } + if ok && r.agent == "" && r.user != "" { + keyAllowance(c.p) + } + return r, ok +} + // spentAfter says whether every candidate in cs rests with its allowance // run out: none of them is likely to answer. func spentAfter(cs []candidate) bool { diff --git a/internal/gateway/key_limits_test.go b/internal/gateway/key_limits_test.go new file mode 100644 index 00000000..5715b787 --- /dev/null +++ b/internal/gateway/key_limits_test.go @@ -0,0 +1,345 @@ +package gateway + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/yetone/magpie/internal/provider" +) + +// sub2api's refusals, as v0.1.149 sends them on the OpenAI and Responses +// paths: a key out of the 5-hour, day or 7-day limit its owner gave it is +// out of quota till that window starts again, though it says +// rate_limit_exceeded; the relay out of accounts for everyone, or its +// upstream rate limiting, is not the key's quota. +func TestSub2APIKeyLimitIsQuota(t *testing.T) { + for _, c := range []struct { + status int + body string + want string + }{ + {429, `{"error":{"message":"api key 7天限额已用完","type":"rate_limit_exceeded"}}`, failQuota}, + {429, `{"error":{"message":"api key 日限额已用完","type":"rate_limit_exceeded"}}`, failQuota}, + {429, `{"error":{"message":"api key 5小时限额已用完","type":"rate_limit_exceeded"}}`, failQuota}, + // the key's own total quota (its middleware's words) + {429, `{"code":"API_KEY_QUOTA_EXHAUSTED","message":"API key 额度已用完"}`, failQuota}, + // the pool: v0.1.149's 503, later versions' 429 + {503, `{"error":{"message":"No available accounts","type":"api_error"}}`, failOther}, + {429, `{"error":{"message":"All available accounts are currently rate-limited. Please retry later.","type":"rate_limit_error"}}`, failRate}, + {429, `{"error":{"message":"Upstream rate limit exceeded, please retry later","type":"rate_limit_error"}}`, failRate}, + // the wallet + {403, `{"code":"INSUFFICIENT_BALANCE","message":"Insufficient account balance"}`, failCredit}, + } { + if got := failure(c.status, []byte(c.body)); got != c.want { + t.Errorf("%d %s: %s, want %s", c.status, c.body, got, c.want) + } + } +} + +// keyWindows swaps in each key's windows, by the key, for a test. +func keyWindows(t *testing.T, of map[string]provider.Allowance) *[]string { + t.Helper() + old, oldStale := keyAllowance, staleKeyAllowance + var mu sync.Mutex + var staled []string + keyAllowance = func(p provider.Provider) (provider.Allowance, bool) { + a, ok := of[p.Key] + return a, ok + } + staleKeyAllowance = func(p provider.Provider) { + mu.Lock() + staled = append(staled, p.Key) + mu.Unlock() + } + t.Cleanup(func() { keyAllowance, staleKeyAllowance = old, oldStale }) + return &staled +} + +func week(used float64, renews time.Time) provider.Allowance { + return provider.Allowance{{Used: used, Resets: renews, Span: 7 * 24 * time.Hour, Amount: used * 8, Of: 800, Unit: "USD"}} +} + +// A provider's keys whose windows magpie reads (sub2api keys given a +// limit) are weighed by them as its accounts are: smart spends first the +// one whose week renews soonest and leaves one used up last, least used +// goes by the share, weekly pace by what is left per hour. +func TestKeysWeighedByTheirWindows(t *testing.T) { + now := time.Now() + of := map[string]provider.Allowance{} + keyWindows(t, of) + key := func(k string) candidate { + return candidate{p: provider.Provider{ID: "kw", Key: k}, model: "gpt-6-astra", rest: "kw#" + k} + } + cs := []candidate{key("a"), key("b"), key("c")} + p := provider.Provider{ID: "kw"} + + // smart: b's week renews in a day, a's in six; c used up + of["a"], of["b"], of["c"] = week(10, now.Add(6*24*time.Hour)), week(40, now.Add(24*time.Hour)), week(99, now.Add(2*24*time.Hour)) + got, wg := weigh(p, cs, "gpt-6-astra", provider.Responses) + if restsOf(got) != "kw#b kw#a kw#c " { + t.Fatalf("smart: %s", restsOf(got)) + } + if w := weighed(got[0], p, wg, false, provider.Responses); w.Kind != "key" || !w.Known || w.Used != 40 || w.Limit != 800 || w.Unit != "USD" || len(w.Renews) != 1 { + t.Fatalf("the trace of b: %+v", w) + } + + // least used: by the share + p.Routing = provider.LeastUsed + of["a"], of["b"], of["c"] = week(60, now.Add(6*24*time.Hour)), week(70, now.Add(24*time.Hour)), week(20, now.Add(2*24*time.Hour)) + if got := restsOf(route(p, cs, "gpt-6-astra", provider.Responses)); got != "kw#c kw#a kw#b " { + t.Fatalf("least used: %s", got) + } + + // weekly pace: a has 90% left for 20 hours, b 100% for a week + p.Routing = provider.Pace + of["a"], of["b"], of["c"] = week(10, now.Add(20*time.Hour)), week(0, now.Add(7*24*time.Hour)), week(99, now.Add(time.Hour)) + if got := restsOf(route(p, []candidate{key("b"), key("a"), key("c")}, "gpt-6-astra", provider.Responses)); got != "kw#a kw#b kw#c " { + t.Fatalf("pace: %s", got) + } +} + +// In a group, a sub2api key given a limit stands beside subscriptions by +// its windows: first while its week renews sooner than theirs, after +// them once it is used up — rather than after every subscription with +// quota, whatever it has left. +func TestGroupWeighsKeyWindows(t *testing.T) { + now := time.Now() + old := allowances + t.Cleanup(func() { allowances = old }) + allowances = func(string) map[string]provider.Allowance { + return map[string]provider.Allowance{"me@example.com": week(30, now.Add(5*24*time.Hour))} + } + of := map[string]provider.Allowance{"sk-relay": week(10, now.Add(24*time.Hour))} + keyWindows(t, of) + codex := provider.Provider{ID: "codex", Account: &provider.Account{Agent: "codex", User: "me@example.com"}} + relay := provider.Provider{ID: "sub2api", Key: "sk-relay", Responses: "https://relay.example/v1"} + heads := []candidate{ + {p: codex, model: "gpt-6-astra", rest: "codex"}, + {p: relay, model: "gpt-6-astra", rest: "sub2api"}, + } + g := provider.Provider{ID: provider.GroupPrefix + "g"} + got, wg := weigh(g, heads, "", provider.Responses) + if restsOf(got) != "sub2api codex " { + t.Fatalf("the key renewing first: %s", restsOf(got)) + } + if w := weighed(got[0], relay, wg, false, provider.Responses); w.Kind != "provider" || !w.Known || w.Used != 10 { + t.Fatalf("the key's trace: %+v", w) + } + of["sk-relay"] = week(99, now.Add(24*time.Hour)) + if got, _ := weigh(g, heads, "", provider.Responses); restsOf(got) != "codex sub2api " { + t.Fatalf("the key used up: %s", restsOf(got)) + } +} + +// A key refused for its own limit rests until the window it filled starts +// again, as an account does, and its windows are read again; when they +// were last read short of full, it rests a quarter of an hour. +func TestKeyOutOfItsWindowRestsUntilReset(t *testing.T) { + now := time.Now() + of := map[string]provider.Allowance{ + "sk-full": week(100, now.Add(3*24*time.Hour)), + "sk-stale": week(40, now.Add(3*24*time.Hour)), + } + staled := keyWindows(t, of) + key := func(k string) candidate { + return candidate{p: provider.Provider{ID: "kr", Key: k}, model: "gpt-6-astra", rest: "kr#" + k} + } + s := &Server{} + body := []byte(`{"error":{"message":"api key 7天限额已用完","type":"rate_limit_exceeded"}}`) + near := func(d, want time.Duration) bool { return d > want-time.Minute && d <= want } + + r := s.restAfter(key("sk-full"), 429, http.Header{}, body) + if r.Why != failQuota || r.By != "window" || !near(time.Until(r.Until), 3*24*time.Hour) { + t.Fatalf("full: rests %v by %s (%s)", time.Until(r.Until), r.By, r.Why) + } + r = s.restAfter(key("sk-stale"), 429, http.Header{}, body) + if r.Why != failQuota || r.By != "quota" || !near(time.Until(r.Until), quotaRest) { + t.Fatalf("stale: rests %v by %s (%s)", time.Until(r.Until), r.By, r.Why) + } + if len(*staled) != 2 || (*staled)[0] != "sk-full" || (*staled)[1] != "sk-stale" { + t.Fatalf("read again: %v", *staled) + } +} + +// The relay out of accounts for everyone (sub2api's 503) says nothing of +// the key's own windows: it backs off a minute, full or not, and they +// aren't read again for it. +func TestKeyPoolErrorDoesNotRestForWindow(t *testing.T) { + of := map[string]provider.Allowance{"sk-pool": week(100, time.Now().Add(3*24*time.Hour))} + staled := keyWindows(t, of) + c := candidate{p: provider.Provider{ID: "kp", Key: "sk-pool"}, model: "gpt-6-astra", rest: "kp#sk-pool"} + forget := func() { routed.Lock(); delete(routed.failures, c.restKey()); routed.Unlock() } + forget() // its first failure + t.Cleanup(forget) + r := (&Server{}).restAfter(c, 503, http.Header{}, []byte(`{"error":{"message":"No available accounts","type":"api_error"}}`)) + if r.Why != failOther || r.By != "backoff" || time.Until(r.Until) > fallbackCooldown { + t.Fatalf("pool empty: rests %v by %s (%s)", time.Until(r.Until), r.By, r.Why) + } + if len(*staled) != 0 { + t.Fatalf("read again: %v", *staled) + } +} + +// A key out of its own window is back as soon as a reading finds that +// window full no more — its owner raised its limit, or reset its usage — +// rather than days later as the window starts again. Only its own rest +// out of quota is lifted: another key still out of its week stays out, +// and the relay out of accounts for everyone says nothing of a key's +// windows. +func TestKeyBackOnceItsWindowIsNotFull(t *testing.T) { + fresh(t) + provider.ForgetBalances() + t.Cleanup(provider.ForgetBalances) + old := keyAllowance + keyAllowance = provider.KeyAllowance // the keys' windows as read, here + t.Cleanup(func() { keyAllowance = old }) + week := time.Now().Add(72 * time.Hour).Truncate(time.Second).Format(time.RFC3339) + var mu sync.Mutex + limit := map[string]int{"Bearer sk-a": 800, "Bearer sk-b": 800, "Bearer sk-c": 800} + relay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + l := limit[r.Header.Get("Authorization")] + mu.Unlock() + fmt.Fprintf(w, `{"isValid":true,"mode":"quota_limited","status":"active","rate_limits":[{"window":"7d","limit":%d,"used":800,"remaining":%d,"reset_at":%q}]}`, l, l-800, week) + })) + defer relay.Close() + if err := provider.Save(provider.Provider{ID: "kb", Name: "KB", Key: "sk-a", Keys: []provider.KeyAccount{{Key: "sk-b"}, {Key: "sk-c"}}, + Responses: relay.URL + "/v1", BalanceURL: relay.URL + "/v1/usage"}); err != nil { + t.Fatal(err) + } + saved, err := provider.Find("kb") + if err != nil { + t.Fatal(err) + } + keys := map[string]candidate{} + for _, c := range perKey(*saved, "gpt-6-astra", provider.Responses) { + keys[c.p.Key] = c + } + a, b, c := keys["sk-a"], keys["sk-b"], keys["sk-c"] + forget := func() { + routed.Lock() + delete(routed.failures, c.restKey()) + routed.Unlock() + for _, k := range keys { + clearRest(k.restKey()) + } + } + forget() + t.Cleanup(forget) + provider.KeyBalances(context.Background()) // each key's week read: used up + s := &Server{} + quota := []byte(`{"error":{"message":"api key 7天限额已用完","type":"rate_limit_exceeded"}}`) + for _, k := range []candidate{a, b} { + if r := s.restAfter(k, 429, http.Header{}, quota); r.Why != failQuota || r.By != "window" { + t.Fatalf("%s out of its week: rests by %s (%s)", k.p.Key, r.By, r.Why) + } + } + if r := s.restAfter(c, 503, http.Header{}, []byte(`{"error":{"message":"No available accounts","type":"api_error"}}`)); r.By != "backoff" { + t.Fatalf("the relay out of accounts: rests by %s (%s)", r.By, r.Why) + } + + // a and c given $1,000 a week, b not; the cards read them again + mu.Lock() + limit["Bearer sk-a"], limit["Bearer sk-c"] = 1000, 1000 + mu.Unlock() + provider.ForgetBalances() + provider.KeyBalances(context.Background()) + if r, ok := restOf(a.restKey()); ok { + t.Fatalf("a, its limit raised: still rests by %s till %v", r.By, r.Until) + } + if r, ok := restOf(b.restKey()); !ok || r.By != "window" { + t.Fatalf("b, still out of its week: %+v, %v", r, ok) + } + if r, ok := restOf(c.restKey()); !ok || r.By != "backoff" { + t.Fatalf("c, the relay out of accounts: %+v, %v", r, ok) + } +} + +// A key resting out of its own window has it read again as it is +// planned, so a limit raised brings it back where nothing weighs it by +// its windows: an ordered group weighs each member's keys alone, and a +// member with one key not at all. No card is read for it. +func TestRestingKeyReadAgainInAnOrderedGroup(t *testing.T) { + fresh(t) + provider.ForgetBalances() + t.Cleanup(provider.ForgetBalances) + old := keyAllowance + keyAllowance = provider.KeyAllowance // the key's windows as read, here + t.Cleanup(func() { keyAllowance = old }) + week := time.Now().Add(72 * time.Hour).Truncate(time.Second).Format(time.RFC3339) + var mu sync.Mutex + limit, reads := 800, 0 + relay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + l := limit + reads++ + mu.Unlock() + fmt.Fprintf(w, `{"isValid":true,"mode":"quota_limited","status":"active","rate_limits":[{"window":"7d","limit":%d,"used":800,"remaining":%d,"reset_at":%q}]}`, l, l-800, week) + })) + defer relay.Close() + for _, p := range []provider.Provider{ + {ID: "ko", Name: "KO", Key: "sk-one", Models: []string{"gpt-6-astra"}, Responses: relay.URL + "/v1", BalanceURL: relay.URL + "/v1/usage"}, + {ID: "kx", Name: "KX", Key: "sk-x", Models: []string{"gpt-6-astra"}, Responses: "http://127.0.0.1:1/v1"}, + } { + if err := provider.Save(p); err != nil { + t.Fatal(err) + } + } + if err := provider.SaveGroup(provider.Group{ID: "in-order", Name: "In order", Members: []string{"ko/gpt-6-astra", "kx/gpt-6-astra"}, Routing: provider.Ordered}); err != nil { + t.Fatal(err) + } + g, ms, ok := provider.FindGroup(provider.GroupPrefix + "in-order") + if !ok { + t.Fatal("no group") + } + s := &Server{} + first := func() (string, *Rest) { + out, pl := s.planGroup(g, ms, provider.Responses) + for i, c := range out { + if c.p.ID == "ko" { + return out[0].p.ID, pl.order[i].Rest + } + } + t.Fatalf("ko not planned: %d planned", len(out)) + return "", nil + } + cs, _ := s.planGroup(g, ms, provider.Responses) + if len(cs) != 2 || cs[0].p.ID != "ko" { + t.Fatalf("planned %d, ko first: %v", len(cs), len(cs) > 0 && cs[0].p.ID == "ko") + } + a := cs[0] + t.Cleanup(func() { clearRest(a.restKey()) }) + provider.KeyBalances(context.Background()) // its card: its week used up + quota := []byte(`{"error":{"message":"api key 7天限额已用完","type":"rate_limit_exceeded"}}`) + if r := s.restAfter(a, 429, http.Header{}, quota); r.Why != failQuota || r.By != "window" { + t.Fatalf("out of its week: rests by %s (%s)", r.By, r.Why) + } + + // given $1,000 a week, and no card read again + mu.Lock() + limit = 1000 + mu.Unlock() + if at, r := first(); at != "kx" || r == nil { + t.Fatalf("resting: %s first, its rest %v", at, r) + } + for deadline := time.Now().Add(5 * time.Second); ; time.Sleep(10 * time.Millisecond) { + r, ok := restOf(a.restKey()) + if !ok { + break + } + if time.Now().After(deadline) { + mu.Lock() + n := reads + mu.Unlock() + t.Fatalf("its limit raised: still rests by %s till %v, its windows read %d times", r.By, r.Until.Format(time.RFC3339), n) + } + } + if at, r := first(); at != "ko" || r != nil { + t.Fatalf("back: %s first, its rest %v", at, r) + } +} diff --git a/internal/gateway/routing.go b/internal/gateway/routing.go index 5e359361..862cccb8 100644 --- a/internal/gateway/routing.go +++ b/internal/gateway/routing.go @@ -121,8 +121,10 @@ var ( // quarter of an hour rather than a minute (#153). rateWords = regexp.MustCompile(`(?i)rate.?limit|too many requests|per.?(second|sec|minute|min)\b|\b[rt]pm\b|频率|太频繁`) // plannedWords: a 429 that says the plan's own allowance is used, rate - // words or not — a day's free requests, say. - plannedWords = regexp.MustCompile(`(?i)quota|usage.?limit|hit your .*limit|limit.{0,24}resets|per.?(day|week|month)|daily|weekly|monthly|额度|用量|套餐`) + // words or not — a day's free requests, say, or a sub2api key's own + // 5-hour, day or 7-day limit ("api key 7天限额已用完", sent as + // rate_limit_exceeded). + plannedWords = regexp.MustCompile(`(?i)quota|usage.?limit|hit your .*limit|limit.{0,24}resets|per.?(day|week|month)|daily|weekly|monthly|额度|用量|套餐|限额已用完`) // resetsWords: Claude Code's "usage limit reached|". resetsWords = regexp.MustCompile(`(?i)limit reached\|(\d{10})\b`) ) @@ -245,13 +247,18 @@ type Rest struct { // Link: where the vendor said to verify the account, for failVerify Link string `json:"link,omitempty"` - agent, user string // the subscription account resting, when it is one + // the subscription account resting, when it is one; agent "" and its + // provider.KeyAllowanceID for a key out of its own windows + agent, user string said string // the error it gave, for a failVerify held hold time.Time // until when a failVerify is answered without asking } // renewed lifts the rests of a subscription account whose windows were -// just started again (a Codex reset spent): out of quota no longer. +// just started again (a Codex reset spent): out of quota no longer. Told +// agent "", it is a key (provider.KeyAllowanceID) whose windows a reading +// found full no more — its limit raised, or its usage reset — and its +// rest out of them is lifted; a key's other rests aren't noted by it. func renewed(agent, user string) { restingUntil.Lock() defer restingUntil.Unlock() @@ -268,6 +275,13 @@ func init() { provider.OnRenewed(renewed) } // staleAllowance is provider.StaleAllowance, swapped in tests. var staleAllowance = provider.StaleAllowance +// keyAllowance and staleKeyAllowance are provider.KeyAllowance and +// provider.StaleKeyAllowance, swapped in tests. +var ( + keyAllowance = provider.KeyAllowance + staleKeyAllowance = provider.StaleKeyAllowance +) + // RestOf is why what rests by key — a provider's id, or its id#key — is // passed over now, while it is. func RestOf(key string) (Rest, bool) { return restOf(key) } @@ -390,8 +404,9 @@ func (s *Server) restAfterMarked(c candidate, status int, header http.Header, bo routed.Unlock() d, r.By, r.Failures = min(fallbackCooldown< x !== f && !x.fallback && !x.rest); + const someKnown = r.order.some((x) => x.known); const rested = r.order.filter((x) => x !== f && !x.fallback && x.rest).map(who); const restedTo = (s) => rested.length ? t("With {rested} resting after a failure, {who} goes first: ", { rested: rested.join(", "), who: w }) + s : null; switch (f.routing) { @@ -414,11 +415,11 @@ case "rotate": return t("In turn: it's {who}'s turn — each request starts one further along.", { who: w }); case "weight": return t("By weight: it's {who}'s share — each key takes requests as its weight says.", { who: w }); case "usage": - if (f.kind === "account" && f.known) return t("Least used first: {who} has the most of its allowance left — {n} used.", { who: w, n: pct(f.used) }); - if (f.kind === "key") return t("Least used first: {who} served the fewest tokens lately — {n}.", { who: w, n: tokens(f.tokens || 0) }); + if (f.known) return t("Least used first: {who} has the most of its allowance left — {n} used.", { who: w, n: pct(f.used) }); + if (f.kind === "key" && !someKnown) return t("Least used first: {who} served the fewest tokens lately — {n}.", { who: w, n: tokens(f.tokens || 0) }); return t("Least used first: {who} goes first.", { who: w }); case "pace": - if (f.kind === "account" && f.known) { + if (f.known) { // The allowance window used for pace, which may be shorter than a week. const left = known0(f.due) ? Math.min(100, Math.round((f.pace || 0) * Math.max(1, (at(f.due) - at(r.time)) / 36e5))) : null; if (left !== null && f.dueBy === "reset") return t("Weekly pace: {who} has the most remaining allowance per hour until one of its Codex resets about to run out is used by itself — {n} left, used in {d}.", { who: w, n: pct(left), d: dur(at(f.due) - at(r.time)) }); @@ -426,17 +427,22 @@ ? t("Weekly pace: {who} has the most remaining allowance per hour until reset — {n} left, resets in {d}.", { who: w, n: pct(left), d: dur(at(f.due) - at(r.time)) }) : t("Weekly pace: {who} has the most remaining allowance per hour until reset — {n} used.", { who: w, n: pct(f.used) }); } - if (f.kind === "key") return t("Weekly pace: {who} served the fewest tokens lately — {n}.", { who: w, n: tokens(f.tokens || 0) }); + if (f.kind === "key" && !someKnown) return t("Weekly pace: {who} served the fewest tokens lately — {n}.", { who: w, n: tokens(f.tokens || 0) }); return t("Weekly pace: {who} goes first.", { who: w }); } - if (f.kind === "key") { + // keys that read their own windows (a sub2api key's limits) are + // weighed by them, as accounts are, and one not read yet beside them + // too: told below as one + if (f.kind === "key" && !someKnown) { if (!f.fit && peers.some((p) => p.fit > 0)) return t("{who} goes first: it's made for {api}, the API {model} is at home in, so nothing is translated.", { who: w, api: API[f.speaks] || f.speaks, model: f.model }); return restedTo(t("the others go in their order.")) || t("{who} goes first: keys go in their order, those that suit the request first.", { who: w }); } - if (f.kind !== "account") return t("{who} goes first.", { who: w }); + if (f.kind !== "account" && !f.known) return t("{who} goes first.", { who: w }); if (f.learns && peers.some((p) => p.known)) return t("{who} goes first: what it has left isn't known yet, and its answer tells — kept behind those known, it would never answer and never be known.", { who: w }); if (!f.known && !peers.some((p) => p.known)) return t("The vendor hasn't said yet what these accounts have left, so they go in their order: {who} first.", { who: w }); - if (group(f) !== "fine") return t("Every account is at 90% or more of its allowance, so the one with the most left goes first: {who}, at {n}.", { who: w, n: pct(f.used) }); + if (group(f) !== "fine") return r.order.some((x) => x.kind !== "account" && x.known) + ? t("Every key and account here is at 90% or more of its allowance, so the one with the most left goes first: {who}, at {n}.", { who: w, n: pct(f.used) }) + : t("Every account is at 90% or more of its allowance, so the one with the most left goes first: {who}, at {n}.", { who: w, n: pct(f.used) }); const next = peers.find((p) => p.known && group(p) === "fine"); const soon = renewsBy(f).find(Boolean); if (!next) return t("{who} goes first: it has quota to spare, and the others are kept for last.", { who: w }); @@ -461,7 +467,7 @@ } const rule = ruleWhy(r, false); if (rule) out.push(rule); - const smart = (x) => !x.routing && x.kind === "account"; + const smart = (x) => !x.routing && (x.kind === "account" || x.known); const someKnown = r.order.some((x) => x.known); for (const x of r.order.slice(1)) { if (x.sunk && !x.rest) out.push(t("{who} was rate limited at {time} while it had quota left, so it went to the back: it comes round again once those ahead of it are rate limited in turn.", { who: who(x), time: clock(x.sunk) })); @@ -1150,7 +1156,8 @@ what.replaceChildren(el("b", "", g ? g.name : f?.name || r.provider), el("span", "", (g ? " · " + t("routing group") : "") + " · " + t(on === 1 ? "one on" : "{n} on", { n: on }) + (many > 1 ? " · " + t("{n} agents at once", { n: many }) : ""))); - mode.textContent = g && routing === "order" ? t(GROUP_ORDER) : t(!routing && f?.kind === "key" && !g ? KEYS_SMART : m[1]); + // keys that read their own windows are weighed as accounts are + mode.textContent = g && routing === "order" ? t(GROUP_ORDER) : t(!routing && f?.kind === "key" && !r.order.some((x) => x.known) && !g ? KEYS_SMART : m[1]); } // what a row says now: resting, answering, or what routing weighed it by @@ -1182,7 +1189,7 @@ else if (answered.has(id)) s = agents.size > 1 ? t("answered {agent}", { agent: agentName(r.agent) }) : t("answered this request"); else if (gave.has(id)) s = t("{status} · {fail} · passed to {agent}", { status: gave.get(id).status, fail: failWord(gave.get(id).fail), agent: agentName(r.agent) }); else if (w.sunk) s = t("rate limited at {time} · at the back", { time: clock(w.sunk) }); - else if (w.kind === "account" && w.known) { + else if (w.known) { const soon = renews(w)[0]; s = !w.routing && w.used >= 98 ? quota(w, "{n} used · all but used up", "{n} left · all but used up") : !w.routing && w.used >= 90 ? quota(w, "{n} used · kept for last", "{n} left · kept for last") @@ -1193,7 +1200,7 @@ else if (w.speaks) s = t("{api} only", { api: API[w.speaks] || w.speaks }); else s = w.kind === "key" ? t("API key") : t("one key"); if (row.st.textContent !== s) row.st.textContent = s; - const bar = w.kind === "account" && w.known; + const bar = w.known; row.li.classList.toggle("nobar", !bar); row.bi.style.width = bar ? fill(w) : "0"; const on = !resting && (trying.has(id) || answered.has(id) || onWire.has(id)); @@ -1866,7 +1873,7 @@ const resting = a.rest && at(a.rest.until) > n; if (resting) { st = `${failWord(a.rest.why)} · ${restWhen(a.rest)}`; cls = "rest"; } else if (w.unlisted) { st = unlistedWord(w); cls = "left"; } - else if (w.kind === "account" && w.known) { + else if (w.known) { const soon = renews(w)[0]; st = soon && soon <= n ? quota(w, "{n} used at {time}; it has renewed since", "{n} left at {time}; it has renewed since", { time: clock(a.at) }) : (soon ? quota(w, "{n} used · renews in {d}", "{n} left · renews in {d}", { d: dur(soon - n) }) : quota(w, "{n} used", "{n} left")) + " · " + t("as of {time}", { time: clock(a.at) }); @@ -1900,7 +1907,7 @@ fix.append(go); row.append(fix); } - if (w.kind === "account" && w.known) { + if (w.known) { const bar = el("div", "bar"), bi = el("i"); bi.style.width = fill(w); bar.append(bi); diff --git a/internal/gui/tests/routing-key-windows.test.cjs b/internal/gui/tests/routing-key-windows.test.cjs new file mode 100644 index 00000000..ea0466b5 --- /dev/null +++ b/internal/gui/tests/routing-key-windows.test.cjs @@ -0,0 +1,168 @@ +// Run with Node's test runner and Playwright on the module path; see README.md. +// A key whose own usage windows magpie reads (a sub2api key given a 5-hour, +// day or 7-day limit) is told on the Routing page as an account is: its +// row says how much of its window is used and when it renews, with a bar, +// and the story says why it went first by its allowance — beside an +// account in a group, or among a provider's keys, all at 90% or more. A +// key not known is told as before. In English, Chinese, Japanese and +// German. +const assert = require("node:assert/strict"); +const fs = require("node:fs/promises"); +const path = require("node:path"); +const { test } = require("node:test"); +const { chromium, webkit } = require("playwright"); + +const assets = path.resolve(__dirname, "../assets"); +const now = new Date(); +const ago = (minutes) => new Date(now.getTime() - minutes * 60e3).toISOString(); +const ahead = (hours) => new Date(now.getTime() + hours * 36e5).toISOString(); +const week = (used, hours) => ({ known: true, used, renews: [ahead(hours)], amount: used * 8, limit: 800, unit: "USD" }); + +const scenes = { + // a group: the sub2api key's week renews in a day, the account's in five + group: { + group: { id: "g", name: "Daily", routing: "" }, + order: [ + { id: "sub2api", provider: "sub2api", name: "Sub2API", kind: "provider", model: "gpt-6-astra", ...week(10, 24) }, + { id: "codex", provider: "codex", name: "Codex", who: "me@example.com", kind: "account", agent: "codex", model: "gpt-6-astra", known: true, used: 30, renews: [ahead(120)] }, + ], + }, + // a provider's two keys, both at 90% or more + low: { + order: [ + { id: "relay#a", provider: "relay", name: "Relay", who: "relay-a", kind: "key", model: "gpt-6-astra", ...week(92, 30) }, + { id: "relay#b", provider: "relay", name: "Relay", who: "relay-b", kind: "key", model: "gpt-6-astra", ...week(95, 50) }, + ], + }, + // a key not read yet beside one read: weighed with it, not in order + mixed: { + order: [ + { id: "relay#b", provider: "relay", name: "Relay", who: "relay-b", kind: "key", model: "gpt-6-astra" }, + { id: "relay#a", provider: "relay", name: "Relay", who: "relay-a", kind: "key", model: "gpt-6-astra", ...week(95, 30) }, + ], + }, + mixedUsage: { + order: [ + { id: "relay#b", provider: "relay", name: "Relay", who: "relay-b", kind: "key", model: "gpt-6-astra", routing: "usage", tokens: 900 }, + { id: "relay#a", provider: "relay", name: "Relay", who: "relay-a", kind: "key", model: "gpt-6-astra", routing: "usage", tokens: 100, ...week(95, 30) }, + ], + }, + // keys whose windows aren't read: as before + plain: { + order: [ + { id: "relay#a", provider: "relay", name: "Relay", who: "relay-a", kind: "key", model: "gpt-6-astra" }, + { id: "relay#b", provider: "relay", name: "Relay", who: "relay-b", kind: "key", model: "gpt-6-astra" }, + ], + }, +}; + +function serve(lang, scene) { + const s = scenes[scene]; + const route = { + id: 1, seq: 1, time: ago(1), agent: "codex", model: "gpt-6-astra", provider: s.order[0].provider, done: true, status: 200, ms: 900, tokens: 1200, + ...(s.group ? { group: s.group } : {}), order: s.order, + tries: [{ id: s.order[0].id, model: "gpt-6-astra", start: ago(1), done: true, status: 200, ms: 900 }], + }; + return async (request) => { + const url = new URL(request.request().url()); + const json = (data) => request.fulfill({ json: data }); + if (url.pathname === "/boot.js") return request.fulfill({ contentType: "text/javascript", body: `window.bootPrefs = {lang:"${lang}",theme:"light",web:true};` }); + if (url.pathname === "/wails/runtime.js") return request.fulfill({ contentType: "text/javascript", body: "export const Window = {};" }); + if (url.pathname === "/api/state") return json({ agents: [{ id: "codex", name: "Codex", path: "/test/config.toml", fields: [] }], profiles: [], settings: { lang, theme: "light" } }); + if (url.pathname === "/api/gateway/trace") { + if (url.searchParams.get("wait")) return new Promise(() => {}); // nothing more comes + return json({ mine: true, now: now.toISOString(), seq: 1, totals: { requests: 1 }, routes: [route] }); + } + if (url.pathname === "/api/gateway/history") return json({ cut: false, days: [], routes: [] }); + if (url.pathname === "/api/groups") return json({ groups: [] }); + if (url.pathname === "/api/providers") return json({ providers: [], excluded: [], gateway: { running: true, window: true } }); + if (url.pathname.startsWith("/api/")) return json({}); + const file = path.join(assets, url.pathname === "/" ? "index.html" : url.pathname); + const contentType = { ".html": "text/html", ".js": "text/javascript", ".css": "text/css", ".svg": "image/svg+xml", ".png": "image/png" }[path.extname(file)]; + await request.fulfill({ body: await fs.readFile(file), contentType }); + }; +} + +async function open(t, engine, lang, scene) { + const browser = await (engine === "webkit" ? webkit.launch() : chromium.launch({ channel: "chromium" })); + t.after(() => browser.close()); + const page = await browser.newPage({ viewport: { width: 1100, height: 900 }, reducedMotion: "reduce" }); + page.setDefaultTimeout(5000); + const errors = []; + page.on("pageerror", (error) => errors.push(error.message)); + await page.route("**/*", serve(lang, scene)); + await page.goto("http://magpie.test/?view=routing"); + await page.locator(".rt-req").first().click(); + await page.locator(".rt-steps li.why").first().waitFor(); + // each row's state is filled on the next frame, after the story + await page.waitForFunction(() => [...document.querySelectorAll("li")].filter((li) => li.querySelector(".who")).every((li) => li.querySelector("em")?.textContent.trim())); + // what a string says in lang, its {names} filled in + const say = (key, vars = {}) => page.evaluate(([lang, key, vars]) => { + const s = (lang !== "en" && I18N[lang]?.[key]) || key; + return s.replace(/\{(\w+)\}/g, (_, k) => vars[k] ?? `{${k}}`); + }, [lang, key, vars]); + const row = (who) => page.locator("li", { has: page.locator(".who", { hasText: who }) }).first(); + return { page, errors, say, row, why: () => page.locator(".rt-steps li.why").first().textContent() }; +} + +for (const engine of (process.env.BROWSER ? [process.env.BROWSER] : ["chromium", "webkit"])) { + for (const lang of ["en", "zh", "ja", "de"]) { + test(`${engine} ${lang}: a key that reads its windows goes first in a group by them`, async (t) => { + const { page, errors, say, row, why } = await open(t, engine, lang, "group"); + const head = (await say("{who} goes first: of those with quota to spare, its allowance renews soonest — in {d} — and what it has left then is lost. {other} renews later and keeps its own.", { who: "Sub2API" })).split("{d}")[0]; + assert.ok((await why()).startsWith(head), `${await why()}\nwant ${head}…`); + // its row has a bar, as an account's; what it did says its window + const key = row("Sub2API"); + assert.equal(await key.locator(".bar i").evaluate((i) => i.style.width), "10%"); + assert.equal(await key.evaluate((li) => li.classList.contains("nobar")), false); + // a provider's own row is under its name, named by its model + const act = page.locator(".rt-prov", { hasText: "Sub2API" }).locator("xpath=following-sibling::div[contains(@class, 'rt-act')][1]"); + await act.locator(".st", { hasText: "80 / 800" }).waitFor(); + const st = (await act.locator(".st").textContent()).trim(); + assert.ok(st.includes("10%") && st.includes("80 / 800"), st); + assert.equal(await act.locator(".bar i").evaluate((i) => i.style.width), "10%"); + assert.deepEqual(errors, []); + }); + + test(`${engine} ${lang}: keys all at 90% or more say so, keys and all`, async (t) => { + const { page, errors, say, row, why } = await open(t, engine, lang, "low"); + const want = await say("Every key and account here is at 90% or more of its allowance, so the one with the most left goes first: {who}, at {n}.", { who: "relay-a", n: "92%" }); + assert.equal((await why()).trim(), want); + // the other key's place is told as an account's would be + const aside = await say("{who} is at {n} — kept for when the others can't.", { who: "relay-b", n: "95%" }); + const steps = await page.locator(".rt-steps li").allTextContents(); + assert.ok(steps.some((s) => s.includes(aside)), steps.join("\n")); + assert.match(await row("relay-b").locator("em").textContent(), /95%/); + // the header tells Smart as it weighs them, not as keys in their order + assert.equal((await page.locator(".rt-mode").textContent()).trim(), + await say("Smart: of the accounts with quota to spare, the one whose allowance renews soonest goes first — what it has left would be lost at the reset. The week decides; an account with five hours and no week goes by its five hours. One at 90% or more waits until the others can't answer; one resting after a failure goes last.")); + assert.deepEqual(errors, []); + }); + } + for (const lang of ["en", "zh"]) { + test(`${engine} ${lang}: a key not read yet beside one read is told as weighed with it`, async (t) => { + const { page, errors, say, why } = await open(t, engine, lang, "mixed"); + assert.equal((await why()).trim(), await say("{who} goes first.", { who: "relay-b" })); + const aside = await say("{who} is at {n} — kept for when the others can't.", { who: "relay-a", n: "95%" }); + const steps = await page.locator(".rt-steps li").allTextContents(); + assert.ok(steps.some((s) => s.includes(aside)), steps.join("\n")); + assert.equal((await page.locator(".rt-mode").textContent()).trim(), + await say("Smart: of the accounts with quota to spare, the one whose allowance renews soonest goes first — what it has left would be lost at the reset. The week decides; an account with five hours and no week goes by its five hours. One at 90% or more waits until the others can't answer; one resting after a failure goes last.")); + assert.deepEqual(errors, []); + }); + + test(`${engine} ${lang}: least used, a key not read yet beside one read isn't told by its tokens`, async (t) => { + const { errors, say, why } = await open(t, engine, lang, "mixedUsage"); + assert.equal((await why()).trim(), await say("Least used first: {who} goes first.", { who: "relay-b" })); + assert.deepEqual(errors, []); + }); + + test(`${engine} ${lang}: keys whose windows aren't read are told as before`, async (t) => { + const { errors, say, row, why } = await open(t, engine, lang, "plain"); + assert.equal((await why()).trim(), await say("{who} goes first: keys go in their order, those that suit the request first.", { who: "relay-a" })); + assert.equal((await row("relay-b").locator("em").textContent()).trim(), await say("API key")); + assert.equal(await row("relay-b").evaluate((li) => li.classList.contains("nobar")), true); + assert.deepEqual(errors, []); + }); + } +} diff --git a/internal/provider/balance.go b/internal/provider/balance.go index a5bd70ec..a01b5214 100644 --- a/internal/provider/balance.go +++ b/internal/provider/balance.go @@ -665,26 +665,36 @@ func (e *balanceExpr) at(path string) (any, error) { } // Balance asks the vendor what is left on the provider's key in use. ok is -// false when there is no way to ask it. +// false when there is no way to ask it. A key given limits and no quota of +// its own has what its fullest window has left: what it can spend now. func Balance(ctx context.Context, p Provider) (amount string, ok bool, err error) { - amount, _, ok, err = balanceParts(ctx, p) + amount, _, ws, ok, err := balanceRead(ctx, p) + if err == nil && amount == "" && len(ws) > 0 { + left := ws[0].Limit - ws[0].Amount + for _, w := range ws[1:] { + left = min(left, w.Limit-w.Amount) + } + amount = money("$", max(0, left)) + } return amount, ok, err } -// balanceParts is Balance, with the amounts of a balance field the user -// wrote each apart (nil for a vendor magpie reads itself). -func balanceParts(ctx context.Context, p Provider) (amount string, parts []BalancePart, ok bool, err error) { +// balanceRead is what Balance reads, with the amounts of a balance field +// the user wrote each apart (nil for a vendor magpie reads itself), and the +// key's own usage windows when its reply tells them (a sub2api key's +// limits, readSub2APIKeyLimits). +func balanceRead(ctx context.Context, p Provider) (amount string, parts []BalancePart, ws []QuotaWindow, ok bool, err error) { ctx = p.Via(ctx) src, ok := balanceSourceOf(p) if !ok || p.Account != nil || p.Key == "" { - return "", nil, false, nil + return "", nil, nil, false, nil } if src.token != "" && p.BalanceURL != "" && balanceURLPath(p.BalanceURL) == newAPIKeyUsage { - return "", nil, true, errKeyUsageWithToken(p.BalanceURL) + return "", nil, nil, true, errKeyUsageWithToken(p.BalanceURL) } req, err := http.NewRequestWithContext(ctx, http.MethodGet, withBalanceKey(src.url, p.Key, true), nil) if err != nil { - return "", nil, true, err + return "", nil, nil, true, err } if src.token != "" { // a named endpoint still gets the provider's headers, which is @@ -708,9 +718,9 @@ func balanceParts(ctx context.Context, p Provider) (amount string, parts []Balan if err != nil { if k := url.QueryEscape(p.Key); k != "" && strings.Contains(err.Error(), k) { // the URL in the error has the key in it: not shown on a card - return "", nil, true, errors.New(strings.ReplaceAll(err.Error(), k, Mask(p.Key))) + return "", nil, nil, true, errors.New(strings.ReplaceAll(err.Error(), k, Mask(p.Key))) } - return "", nil, true, err + return "", nil, nil, true, err } defer res.Body.Close() b, _ := io.ReadAll(io.LimitReader(res.Body, 1<<20)) @@ -721,28 +731,43 @@ func balanceParts(ctx context.Context, p Provider) (amount string, parts []Balan if res.StatusCode >= 300 { msg = res.Status + ": " + msg } - return "", nil, true, errors.New(msg) + return "", nil, nil, true, errors.New(msg) } if res.StatusCode >= 300 { // what a JSON reply says, on one line; a page of HTML says nothing msg := strings.Join(strings.Fields(string(b)), " ") if !strings.HasPrefix(msg, "{") { - return "", nil, true, errors.New(res.Status) + return "", nil, nil, true, errors.New(res.Status) } if r := []rune(msg); len(r) > 200 { msg = string(r[:200]) + "…" } - return "", nil, true, fmt.Errorf("%s: %s", res.Status, msg) + return "", nil, nil, true, fmt.Errorf("%s: %s", res.Status, msg) } if p.BalanceURL != "" { + if sub2APIKeyLimits(p) { + ws, _ = readSub2APIKeyLimits(b) + } // the field the user wrote: its amounts each apart, as well if parts, err = readBalanceParts(b, balancePathOf(p)); err != nil { - return "", nil, true, err + if len(ws) > 0 && strings.TrimSpace(p.BalancePath) == "" { + // a key given limits and no quota of its own has no + // "remaining": its windows are what it has left + return "", nil, ws, true, nil + } + return "", nil, nil, true, err } - return joinBalanceParts(parts), parts, true, nil + return joinBalanceParts(parts), parts, ws, true, nil } amount, err = src.read(b) - return amount, nil, true, err + return amount, nil, nil, true, err +} + +// sub2APIKeyLimits says the provider's Balance URL is a sub2api panel's +// query for a key, whose reply tells the key's own limits when it has +// some. +func sub2APIKeyLimits(p Provider) bool { + return p.BalanceURL != "" && balanceURLPath(p.BalanceURL) == sub2APIKeyUsage } // balanceKeyNames are what a Balance URL or a header's value names the @@ -775,6 +800,8 @@ func ForgetBalances() { keyBalanceCache.Lock() keyBalanceCache.data = nil keyBalanceCache.Unlock() + // nor is a key's allowance as last read, which may be another key's + forgetKeyAllowances() } // KeyBalances is the balance of every provider magpie can ask one of, the @@ -859,14 +886,17 @@ func KeyBalances(ctx context.Context) []SubscriptionQuota { out[i] = keepReading(ctx, q, keyTag("balance", j.p.Key)) return } - amount, parts, _, err := balanceParts(ctx, j.p) + amount, parts, ws, _, err := balanceRead(ctx, j.p) if err != nil { q.Error = err.Error() } else { q.Balance = amount q.BalanceParts = cardParts(parts) + q.Windows = append(q.Windows, ws...) now := time.Now() q.ReadAt = &now + // what routing goes by, read just now rather than in a minute + noteKeyAllowance(j.p, ws, now) } // the vendor failing a while shows the balance last read out[i] = keepReading(ctx, q, keyTag("balance", j.p.Key)) diff --git a/internal/provider/key_allowance.go b/internal/provider/key_allowance.go new file mode 100644 index 00000000..0ce099e2 --- /dev/null +++ b/internal/provider/key_allowance.go @@ -0,0 +1,202 @@ +package provider + +// A key's own usage windows, for routing, as Allowances is a subscription +// account's: a sub2api key its owner gave a 5-hour, day or 7-day limit +// tells what it has spent of each, and routing weighs and rests it by +// them. They are read behind the request, never waited for, and read +// again once a minute has passed; until the first reading is back, the +// one last shown on the key's card stands in for it. A reading that finds +// a window the key was out of no longer full — its limit raised in the +// panel, or its usage reset — tells OnRenewed's hooks, as a Codex reset +// spent does an account's, so what sat out for it comes back. + +import ( + "context" + "strings" + "sync" + "time" +) + +var keyAllowances struct { + sync.Mutex + m map[string]*keyAllowance // provider id#key id + // gen is bumped as they are all forgotten: a reading begun before + // says nothing of the keys as they are now + gen int +} + +type keyAllowance struct { + a Allowance + at time.Time // when read; zero to read again at the next ask + loading bool +} + +// keyAllowanceAge is how long a key's windows as read go unasked again. +const keyAllowanceAge = time.Minute + +// readsKeyWindows says p is a key whose own usage windows magpie reads: +// one whose Balance URL is a sub2api panel's query for a key. Said +// without asking anything, so any other key costs nothing here. +func readsKeyWindows(p Provider) bool { + return p.Account == nil && p.Key != "" && sub2APIKeyLimits(p) +} + +func keyAllowanceID(p Provider) string { return p.ID + "#" + keyID(p.Key) } + +// KeyAllowanceID is what OnRenewed's hooks are told, as the user of agent +// "", of a key whose own windows magpie reads; empty for any other. +func KeyAllowanceID(p Provider) string { + if !readsKeyWindows(p) { + return "" + } + return keyAllowanceID(p) +} + +// KeyAllowance is the provider's key in use's own usage windows as last +// read, asked for again in the background when that was over a minute +// ago; ok is false for a key without any, or one not read yet. It never +// waits: a key not read yet goes by the last reading on its card, else +// counts as not known. +func KeyAllowance(p Provider) (Allowance, bool) { + if !readsKeyWindows(p) { + return nil, false + } + id := keyAllowanceID(p) + c := &keyAllowances + c.Lock() + _, had := c.m[id] + c.Unlock() + var seed Allowance + if !had { + seed = lastKeyAllowance(p) // read before the lock: it takes the cards' + } + c.Lock() + defer c.Unlock() + if c.m == nil { + c.m = map[string]*keyAllowance{} + } + e := c.m[id] + if e == nil { + e = &keyAllowance{a: seed} + c.m[id] = e + } + if !e.loading && time.Since(e.at) > keyAllowanceAge { + e.loading = true + go readKeyAllowance(p, id, c.gen) + } + return e.a, len(e.a) > 0 +} + +// readKeyAllowance reads p's key's windows into the cache. A reading that +// fails keeps what was known, and is tried again a minute later: a failed +// read is not known, never empty. +func readKeyAllowance(p Provider, id string, gen int) { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + ws, err := keyWindows(ctx, p) + now := time.Now() + c := &keyAllowances + c.Lock() + e := c.m[id] + if e == nil || c.gen != gen { + c.Unlock() + return // forgotten while it was read: read again + } + e.loading, e.at = false, now + renewed := err == nil && e.keep(allowanceOf(ws, now), SpentShareOf(p.Routing), now) + c.Unlock() + if renewed { + tellRenewed("", id) + } +} + +// keep has the key go by a, read at now, and says whether a window it +// was full in (at share, as routing counts it full) is full no more, or +// till sooner: its limit raised, its usage reset, or the window started +// again. +func (e *keyAllowance) keep(a Allowance, share float64, now time.Time) (renewed bool) { + was := e.a.Full("", share, now) + e.a, e.at = a, now + return was.After(a.Full("", share, now)) +} + +// keyWindows asks the vendor for p's key's windows, as the key's card +// is asked (KeyBalances): the provider as saved, with the key, rather +// than as a request has it, which leaves out the base URLs of the +// protocols a key isn't made for. +func keyWindows(ctx context.Context, p Provider) ([]QuotaWindow, error) { + q := p + if saved, err := Find(p.ID); err == nil { + q = *saved + q.Key = p.Key + } + _, _, ws, _, err := balanceRead(ctx, q) + return ws, err +} + +// lastKeyAllowance is p's key's windows as its card last read them, kept +// on disk (keepLast), each started again whose reset has passed. +func lastKeyAllowance(p Provider) Allowance { + c := &lastQuotas + c.Lock() + defer c.Unlock() + c.load() + q, ok := c.reading(p.ID + "/" + strings.ToLower(keyTag("balance", p.Key))) + if !ok || len(q.Windows) == 0 { + return nil + } + return allowanceOf(q.Windows, time.Now()) +} + +// noteKeyAllowance keeps the windows the key's card just read, for +// routing to go by at once. +func noteKeyAllowance(p Provider, ws []QuotaWindow, now time.Time) { + if !readsKeyWindows(p) { + return + } + c := &keyAllowances + c.Lock() + if c.m == nil { + c.m = map[string]*keyAllowance{} + } + id := keyAllowanceID(p) + e := c.m[id] + if e == nil { + e = &keyAllowance{} + c.m[id] = e + } + renewed := e.keep(allowanceOf(ws, now), SpentShareOf(p.Routing), now) + c.Unlock() + if renewed { + tellRenewed("", id) + } +} + +// StaleKeyAllowance has the next KeyAllowance ask the vendor again for +// p's key's windows rather than trust what it last said: the key just +// answered that it is out of them. +func StaleKeyAllowance(p Provider) { + if !readsKeyWindows(p) { + return + } + c := &keyAllowances + c.Lock() + if e := c.m[keyAllowanceID(p)]; e != nil { + e.at = time.Time{} + } + c.Unlock() +} + +// forgetKeyAllowances has every key's windows read again at its next ask: +// a key or its Balance URL changed (ForgetBalances). What was read stands +// until then — fresher than its card's, which routing's reads never write +// — and a key changed is another key, with windows of its own. +func forgetKeyAllowances() { + c := &keyAllowances + c.Lock() + for _, e := range c.m { + e.at, e.loading = time.Time{}, false + } + c.gen++ + c.Unlock() +} diff --git a/internal/provider/routing.go b/internal/provider/routing.go index c8500dbd..9804088b 100644 --- a/internal/provider/routing.go +++ b/internal/provider/routing.go @@ -447,7 +447,10 @@ func Allowances(agent string) map[string]Allowance { } // OnRenewed has f told when an account's usage windows were started again -// (a Codex reset spent), so what sat out waiting for them can come back. +// (a Codex reset spent), so what sat out waiting for them can come back; +// and, as agent "" and the key's KeyAllowanceID, when a reading of a key's +// own windows finds one it was full in full no more: its limit raised in +// its panel, or its usage reset. func OnRenewed(f func(agent, user string)) { renewedHooks.Lock() renewedHooks.fs = append(renewedHooks.fs, f) @@ -465,6 +468,11 @@ var renewedHooks struct { // holds it no more. func renewedNow(agent, user string) { forgetAllowance(agent, user) + tellRenewed(agent, user) +} + +// tellRenewed tells those OnRenewed asked. +func tellRenewed(agent, user string) { renewedHooks.Lock() fs := renewedHooks.fs renewedHooks.Unlock() diff --git a/internal/provider/sub2api_usage.go b/internal/provider/sub2api_usage.go new file mode 100644 index 00000000..c431cd56 --- /dev/null +++ b/internal/provider/sub2api_usage.go @@ -0,0 +1,61 @@ +package provider + +// A sub2api key's own limits: its panel lets the owner give a key a +// spending limit per 5 hours, day and 7 days (rate_limit_5h/1d/7d), and +// once it has one, the key's /v1/usage answers in its quota_limited mode +// with each window's spend, limit and reset — and, unless the key has a +// total quota too, no "remaining" at all. + +import ( + "encoding/json" + "strings" + "time" +) + +// sub2APIWindows are the windows a sub2api key's rate_limits name, with +// the names magpie's cards already give a 5-hour, a day's and a week's. +var sub2APIWindows = []struct { + window, name string + span time.Duration +}{ + {"5h", "5 hours", 5 * time.Hour}, + {"1d", "1 day", 24 * time.Hour}, + {"7d", "7 days", 7 * 24 * time.Hour}, +} + +// readSub2APIKeyLimits is the windows of a sub2api /v1/usage reply's +// rate_limits, in dollars: {"window":"7d","limit":800,"used":12.5, +// "remaining":787.5,"window_start":…,"reset_at":…}. A window with no +// reset_at hasn't started, or has run out since (sub2api says used 0 for +// it then, and gives no window a limit of 0): it starts on the key's next +// request, so it renews within its span. ok is false for a reply without them — a key on the wallet, or not +// sub2api's. +func readSub2APIKeyLimits(body []byte) (ws []QuotaWindow, ok bool) { + var r struct { + RateLimits []struct { + Window string `json:"window"` + Limit float64 `json:"limit"` + Used float64 `json:"used"` + ResetAt *time.Time `json:"reset_at"` + } `json:"rate_limits"` + } + if json.Unmarshal(body, &r) != nil { + return nil, false + } + for _, s := range sub2APIWindows { + for _, l := range r.RateLimits { + if strings.TrimSpace(l.Window) != s.window || l.Limit <= 0 { + continue + } + w := QuotaWindow{Name: s.name, Span: s.span, Amount: max(0, l.Used), Limit: l.Limit, Unit: "USD"} + w.Used = min(100, max(0, l.Used/l.Limit*100)) + if l.ResetAt != nil && !l.ResetAt.IsZero() { + t := *l.ResetAt + w.ResetsAt = &t + } + ws = append(ws, w) + break + } + } + return ws, len(ws) > 0 +} diff --git a/internal/provider/sub2api_usage_test.go b/internal/provider/sub2api_usage_test.go new file mode 100644 index 00000000..8b9ac54d --- /dev/null +++ b/internal/provider/sub2api_usage_test.go @@ -0,0 +1,533 @@ +package provider + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/yetone/magpie/internal/agentenv" +) + +// sub2apiUsage is a sub2api key's /v1/usage as sub2api v0.1.149 answered +// it, the key given a $300 day and an $800 week in its panel: before its +// next request (not started) and after one. +func sub2apiUsage(t *testing.T, started bool) []byte { + t.Helper() + name := "testdata/sub2api_usage_key_limits_not_started.json" + if started { + name = "testdata/sub2api_usage_key_limits.json" + } + b, err := os.ReadFile(name) + if err != nil { + t.Fatal(err) + } + return b +} + +// weekUsed is the started reply with $used of the week spent, its day and +// week resetting a day and three days from now, as a test run on any day +// sees them. +func weekUsed(t *testing.T, used string, week time.Time) []byte { + t.Helper() + b := sub2apiUsage(t, true) + b = bytes.Replace(b, []byte(`"used":0.0000958,"window":"7d"`), []byte(`"used":`+used+`,"window":"7d"`), 1) + b = bytes.Replace(b, []byte("2026-10-07T00:00:00+08:00"), []byte(time.Now().Add(24*time.Hour).Format(time.RFC3339)), 1) + return bytes.Replace(b, []byte("2026-10-13T00:00:00+08:00"), []byte(week.Format(time.RFC3339)), 1) +} + +// A key given limits answers with its windows' spend in dollars; a +// window not started (or run out since) says used 0 and no reset_at. +func TestReadSub2APIKeyLimits(t *testing.T) { + ws, ok := readSub2APIKeyLimits(sub2apiUsage(t, true)) + if !ok || len(ws) != 2 { + t.Fatalf("windows = %+v, %v", ws, ok) + } + day, week := ws[0], ws[1] + reset := func(s string) time.Time { + v, err := time.Parse(time.RFC3339, s) + if err != nil { + t.Fatal(err) + } + return v + } + if day.Name != "1 day" || day.Span != 24*time.Hour || day.Limit != 300 || day.Unit != "USD" || + day.ResetsAt == nil || !day.ResetsAt.Equal(reset("2026-10-07T00:00:00+08:00")) || + day.Amount != 0.0000958 || day.Used <= 0 || day.Used > 0.0001 { + t.Fatalf("day = %+v", day) + } + if week.Name != "7 days" || week.Span != 7*24*time.Hour || week.Limit != 800 || + week.ResetsAt == nil || !week.ResetsAt.Equal(reset("2026-10-13T00:00:00+08:00")) { + t.Fatalf("week = %+v", week) + } + + ws, ok = readSub2APIKeyLimits(sub2apiUsage(t, false)) + if !ok || len(ws) != 2 || ws[0].ResetsAt != nil || ws[1].ResetsAt != nil || ws[1].Used != 0 || ws[1].Span != 7*24*time.Hour { + t.Fatalf("not started = %+v, %v", ws, ok) + } + + // $720 of the week's $800 is 90%, over it no more than 100 + b := bytes.Replace(sub2apiUsage(t, true), []byte(`"used":0.0000958,"window":"7d"`), []byte(`"used":720,"window":"7d"`), 1) + if ws, _ = readSub2APIKeyLimits(b); len(ws) != 2 || ws[1].Used != 90 || ws[1].Amount != 720 { + t.Fatalf("$720 of $800 = %+v", ws) + } + b = bytes.Replace(sub2apiUsage(t, true), []byte(`"used":0.0000958,"window":"7d"`), []byte(`"used":812.5,"window":"7d"`), 1) + if ws, _ = readSub2APIKeyLimits(b); len(ws) != 2 || ws[1].Used != 100 { + t.Fatalf("$812.50 of $800 = %+v", ws) + } + + // a window without a limit is none + b = bytes.Replace(sub2apiUsage(t, true), []byte(`"limit":300`), []byte(`"limit":0`), 1) + if ws, _ = readSub2APIKeyLimits(b); len(ws) != 1 || ws[0].Name != "7 days" { + t.Fatalf("a day's limit of 0 = %+v", ws) + } + + // a key on the wallet has none + if ws, ok := readSub2APIKeyLimits([]byte(`{"mode":"unrestricted","isValid":true,"planName":"钱包余额","remaining":7.25,"unit":"USD","balance":7.25}`)); ok || ws != nil { + t.Fatalf("wallet = %+v, %v", ws, ok) + } +} + +// keyLimitsHome gives a test a magpie of its own, without any key or +// window another test read. +func keyLimitsHome(t *testing.T) { + t.Helper() + isolate(t) + h := t.TempDir() + t.Setenv("HOME", h) + t.Setenv("USERPROFILE", h) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(h, ".config")) + t.Setenv("XDG_CACHE_HOME", filepath.Join(h, ".cache")) + t.Setenv("PATH", h) + for _, v := range agentenv.Vars { + t.Setenv(v, "") + } + restart := func() { + lastQuotas.Lock() + lastQuotas.m, lastQuotas.loaded = nil, false + lastQuotas.Unlock() + ForgetBalances() + keyAllowances.Lock() + keyAllowances.m = nil + keyAllowances.Unlock() + } + restart() + t.Cleanup(func() { + // the reads behind the test, and the sign-ins they look up, end + // before the next test's isolate changes what they call + deadline := time.Now().Add(10 * time.Second) + for busy := true; busy && time.Now().Before(deadline); time.Sleep(5 * time.Millisecond) { + busy = false + keyAllowances.Lock() + for _, e := range keyAllowances.m { + busy = busy || e.loading + } + keyAllowances.Unlock() + devinStatus.Lock() + busy = busy || devinStatus.refreshing + devinStatus.Unlock() + } + restart() + }) +} + +// sub2apiServer answers /v1/usage for sk-limited with body(), counting +// the asks; any other key is refused as sub2api refuses it. +func sub2apiServer(t *testing.T, body func() ([]byte, int)) (*httptest.Server, *atomic.Int32) { + t.Helper() + var asked atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/usage" { + t.Errorf("asked %s", r.URL.Path) + } + if r.Header.Get("Authorization") != "Bearer sk-limited" { + w.WriteHeader(http.StatusUnauthorized) + w.Write([]byte(`{"error":{"type":"authentication_error","message":"Invalid API key"}}`)) + return + } + asked.Add(1) + b, status := body() + w.WriteHeader(status) + w.Write(b) + })) + t.Cleanup(srv.Close) + return srv, &asked +} + +func sub2apiKey(srv *httptest.Server) Provider { + return Provider{ID: "sub2api", Name: "Sub2API", Chat: srv.URL + "/v1", Responses: srv.URL + "/v1", + Key: "sk-limited", BalanceURL: srv.URL + "/v1/usage"} +} + +// A key its owner gave limits has no "remaining" of its own: its card +// shows the windows it has spent of, not "nothing at remaining"; a field +// the user wrote is still read, beside them. +func TestSub2APIKeyLimitsOnItsCard(t *testing.T) { + keyLimitsHome(t) + var body atomic.Pointer[[]byte] + b := sub2apiUsage(t, true) + body.Store(&b) + srv, _ := sub2apiServer(t, func() ([]byte, int) { return *body.Load(), 200 }) + p := sub2apiKey(srv) + if err := Save(p); err != nil { + t.Fatal(err) + } + card := func() SubscriptionQuota { + ForgetBalances() + for _, q := range KeyBalances(context.Background()) { + if q.Provider == "sub2api" { + return q + } + } + t.Fatal("no card for the key") + return SubscriptionQuota{} + } + q := card() + if q.Error != "" || q.Balance != "" || len(q.Windows) != 2 || q.Windows[0].Name != "1 day" || q.Windows[1].Name != "7 days" || q.ReadAt == nil { + t.Fatalf("card = %+v", q) + } + // checked from its editor, or magpie provider show: what it can + // spend now, the least its windows have left + if amount, ok, err := Balance(context.Background(), p); err != nil || !ok || amount != "$300.00" { + t.Fatalf("Balance = %q %v %v", amount, ok, err) + } + week := weekUsed(t, "720", time.Now().Add(72*time.Hour)) + body.Store(&week) + if amount, _, err := Balance(context.Background(), p); err != nil || amount != "$80.00" { + t.Fatalf("Balance with $720 of the week's $800 spent = %q %v", amount, err) + } + body.Store(&b) + + p.BalancePath = "usage.total.actual_cost" + if err := Save(p); err != nil { + t.Fatal(err) + } + if q = card(); q.Error != "" || q.Balance != "192.33" || len(q.Windows) != 2 { + t.Fatalf("with a field: %+v", q) + } + // a field the user wrote that isn't there still says so + p.BalancePath = "remaining" + if err := Save(p); err != nil { + t.Fatal(err) + } + if q = card(); q.Error == "" || len(q.Windows) != 0 { + t.Fatalf("a field not there: %+v", q) + } +} + +// waitAllowance waits for p's key's windows to be read. +func waitAllowance(t *testing.T, p Provider, until func(Allowance) bool) Allowance { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for { + if a, ok := KeyAllowance(p); ok && until(a) { + return a + } + if time.Now().After(deadline) { + t.Fatal("the key's windows were never read") + } + time.Sleep(5 * time.Millisecond) + } +} + +// Routing asks for a key's windows on every request: it must never wait +// for sub2api, nor ask it more than once at a time. +func TestKeyAllowanceNeverWaits(t *testing.T) { + keyLimitsHome(t) + week := time.Now().Add(72 * time.Hour).Truncate(time.Second) + release := make(chan struct{}) + srv, asked := sub2apiServer(t, func() ([]byte, int) { + <-release + return weekUsed(t, "720", week), 200 + }) + p := sub2apiKey(srv) + if err := Save(p); err != nil { + t.Fatal(err) + } + var wg sync.WaitGroup + for range 10 { + wg.Add(1) + go func() { + defer wg.Done() + start := time.Now() + if a, ok := KeyAllowance(p); ok || a != nil { + t.Errorf("known before it was read: %+v", a) + } + if d := time.Since(start); d > time.Second { + t.Errorf("waited %v", d) + } + }() + } + wg.Wait() + close(release) + a := waitAllowance(t, p, func(Allowance) bool { return true }) + if used, _ := a.For("gpt-6-astra", time.Now()); used != 90 { + t.Fatalf("used = %v of %+v", used, a) + } + if n := asked.Load(); n != 1 { + t.Fatalf("asked %d times", n) + } + // read within the minute: not asked again + KeyAllowance(p) + time.Sleep(20 * time.Millisecond) + if n := asked.Load(); n != 1 { + t.Fatalf("asked %d times within the minute", n) + } +} + +// A read that fails says nothing of the key: what was known stands, and +// is asked for again a minute on rather than on every request. +func TestKeyAllowanceFailedReadKeepsLast(t *testing.T) { + keyLimitsHome(t) + week := time.Now().Add(72 * time.Hour).Truncate(time.Second) + var down atomic.Bool + srv, asked := sub2apiServer(t, func() ([]byte, int) { + if down.Load() { + return []byte(`{"error":{"type":"api_error","message":"internal error"}}`), 502 + } + return weekUsed(t, "720", week), 200 + }) + p := sub2apiKey(srv) + if err := Save(p); err != nil { + t.Fatal(err) + } + waitAllowance(t, p, func(Allowance) bool { return true }) + down.Store(true) + StaleKeyAllowance(p) + KeyAllowance(p) // asks again, and is told nothing + deadline := time.Now().Add(5 * time.Second) + for asked.Load() < 2 { + if time.Now().After(deadline) { + t.Fatal("not asked again once stale") + } + time.Sleep(5 * time.Millisecond) + } + for range 20 { + a, ok := KeyAllowance(p) + if used, _ := a.For("gpt-6-astra", time.Now()); !ok || used != 90 { + t.Fatalf("after a failed read: %+v, %v", a, ok) + } + time.Sleep(5 * time.Millisecond) + } + if n := asked.Load(); n != 2 { + t.Fatalf("asked %d times after failing", n) + } +} + +// A read under way as the keys are forgotten is dropped, and the key is +// read again at its next ask — not left waiting on the read dropped. +func TestKeyAllowanceForgottenMidRead(t *testing.T) { + keyLimitsHome(t) + release := make(chan struct{}) + var first atomic.Bool + srv, asked := sub2apiServer(t, func() ([]byte, int) { + if first.CompareAndSwap(false, true) { + <-release + } + return sub2apiUsage(t, true), 200 + }) + var once sync.Once + t.Cleanup(func() { once.Do(func() { close(release) }) }) + p := sub2apiKey(srv) + KeyAllowance(p) + for asked.Load() < 1 { + time.Sleep(5 * time.Millisecond) + } + ForgetBalances() + once.Do(func() { close(release) }) + waitAllowance(t, p, func(a Allowance) bool { return len(a) == 2 }) + if n := asked.Load(); n != 2 { + t.Fatalf("asked %d times", n) + } +} + +// Until the key is first read (magpie just started), its windows are the +// ones its card last showed, kept on disk; a window whose reset passed +// since is empty again. +func TestKeyAllowanceSeededFromItsCard(t *testing.T) { + keyLimitsHome(t) + week := time.Now().Add(72 * time.Hour).Truncate(time.Second) + var down atomic.Bool + srv, _ := sub2apiServer(t, func() ([]byte, int) { + if down.Load() { + return nil, 503 + } + return weekUsed(t, "720", week), 200 + }) + p := sub2apiKey(srv) + if err := Save(p); err != nil { + t.Fatal(err) + } + KeyBalances(context.Background()) + down.Store(true) + // as after a restart: nothing in memory + lastQuotas.Lock() + lastQuotas.m, lastQuotas.loaded = nil, false + lastQuotas.Unlock() + keyAllowances.Lock() + keyAllowances.m = nil + keyAllowances.Unlock() + a, ok := KeyAllowance(p) + now := time.Now() + if used, _ := a.For("gpt-6-astra", now); !ok || used != 90 { + t.Fatalf("seeded = %+v, %v", a, ok) + } + // the week renews on the reset sub2api told, not a week from now + if r := a.Renewal("gpt-6-astra", now); len(r) != 2 || !r[0].Equal(week) { + t.Fatalf("renewal = %v", r) + } +} + +// Only a key whose Balance URL is sub2api's /v1/usage is asked: any other +// key, or an account, costs nothing. +func TestKeyAllowanceOnlyForSub2APIKeys(t *testing.T) { + keyLimitsHome(t) + srv, asked := sub2apiServer(t, func() ([]byte, int) { return sub2apiUsage(t, true), 200 }) + for _, p := range []Provider{ + {ID: "relay", Chat: srv.URL + "/v1", Key: "sk-limited", BalanceURL: srv.URL + "/api/usage/token"}, + {ID: "plain", Chat: srv.URL + "/v1", Key: "sk-limited"}, + {ID: "sub2api", Chat: srv.URL + "/v1", BalanceURL: srv.URL + "/v1/usage"}, + {ID: "codex", Chat: srv.URL + "/v1", Key: "sk-limited", BalanceURL: srv.URL + "/v1/usage", Account: &Account{Agent: "codex", User: "a@example.com"}}, + } { + if a, ok := KeyAllowance(p); ok || a != nil { + t.Errorf("%s: %+v", p.ID, a) + } + } + time.Sleep(20 * time.Millisecond) + if n := asked.Load(); n != 0 { + t.Fatalf("asked %d times", n) + } +} + +// Each key of a provider has windows of its own; a key or Balance URL +// changed (ForgetBalances) has them all read again. +func TestKeyAllowancePerKeyAndForgotten(t *testing.T) { + keyLimitsHome(t) + srv, asked := sub2apiServer(t, func() ([]byte, int) { return sub2apiUsage(t, true), 200 }) + p := sub2apiKey(srv) + now := time.Now() + week := now.Add(72 * time.Hour) + noteKeyAllowance(p, []QuotaWindow{{Name: "7 days", Span: 7 * 24 * time.Hour, Used: 95, ResetsAt: &week}}, now) + other := p + other.Key = "sk-other" + if a, ok := KeyAllowance(other); ok || a != nil { + t.Fatalf("another key's: %+v", a) + } + if _, ok := KeyAllowance(p); !ok { + t.Fatal("the key's own not known") + } + // a read begun before ForgetBalances doesn't land after it + keyAllowances.Lock() + gen := keyAllowances.gen + keyAllowances.Unlock() + ForgetBalances() + readKeyAllowance(p, keyAllowanceID(p), gen) + keyAllowances.Lock() + e := keyAllowances.m[keyAllowanceID(p)] + var kept Allowance + if e != nil { + kept = e.a + } + keyAllowances.Unlock() + if len(kept) != 1 || kept[0].Used != 95 { + t.Fatalf("after ForgetBalances, and a read begun before it: %+v", kept) + } + // what was read stands, not the card's older reading, until the key + // is read again at its next ask + if a, ok := KeyAllowance(p); !ok || len(a) != 1 || a[0].Used != 95 { + t.Fatalf("after ForgetBalances: %+v, %v", a, ok) + } + waitAllowance(t, p, func(a Allowance) bool { return len(a) == 2 }) + if n := asked.Load(); n != 2 { // the read dropped, and the one after + t.Fatalf("asked %d times after ForgetBalances", n) + } + if strings.Contains(keyAllowanceID(p), "sk-limited") { + t.Fatal("the key itself names its windows") + } +} + +// A key out of a window is told renewed (OnRenewed, as agent "" and its +// KeyAllowanceID) once a reading finds that window full no more — its +// limit raised in the panel, or its usage reset — whether routing read it +// or its card did; not on its first reading, while it stays full, as it +// fills, or when a read fails. +func TestKeyAllowanceRenewed(t *testing.T) { + keyLimitsHome(t) + week := time.Now().Add(72 * time.Hour).Truncate(time.Second) + var body atomic.Pointer[[]byte] + srv, _ := sub2apiServer(t, func() ([]byte, int) { + if b := *body.Load(); b != nil { + return b, 200 + } + return []byte(`{"error":{"type":"api_error","message":"internal error"}}`), 502 + }) + p := sub2apiKey(srv) + if err := Save(p); err != nil { + t.Fatal(err) + } + renewedHooks.Lock() + hooks := renewedHooks.fs + renewedHooks.Unlock() + t.Cleanup(func() { + renewedHooks.Lock() + renewedHooks.fs = hooks + renewedHooks.Unlock() + }) + var mu sync.Mutex + var told []string + OnRenewed(func(agent, user string) { + mu.Lock() + told = append(told, agent+"/"+user) + mu.Unlock() + }) + id := KeyAllowanceID(p) + if id != keyAllowanceID(p) { + t.Fatalf("KeyAllowanceID = %q", id) + } + reply := func(b []byte) { body.Store(&b) } + // read is a reading as routing's, done here rather than behind it + read := func(b []byte) { + reply(b) + keyAllowances.Lock() + gen := keyAllowances.gen + keyAllowances.Unlock() + readKeyAllowance(p, id, gen) + } + tellsSo := func(what string, n int) { + t.Helper() + mu.Lock() + defer mu.Unlock() + if len(told) != n { + t.Fatalf("%s: told %v", what, told) + } + for _, s := range told { + if s != "/"+id { + t.Fatalf("%s: told %v", what, told) + } + } + } + full := weekUsed(t, "800", week) + reply(full) + waitAllowance(t, p, func(a Allowance) bool { u, _ := a.For("gpt-6-astra", time.Now()); return u == 100 }) + tellsSo("its first reading, out of its week", 0) + read(full) + tellsSo("its week still used up", 0) + read(bytes.Replace(full, []byte(`"limit":800`), []byte(`"limit":1000`), 1)) + tellsSo("its week's limit raised to $1,000", 1) + read(full) + tellsSo("its week used up again", 1) + read(nil) + tellsSo("a read that failed", 1) + // its usage reset in the panel, as its card reads it + reply(sub2apiUsage(t, false)) + ForgetBalances() + KeyBalances(context.Background()) + tellsSo("its usage reset, read on its card", 2) +} diff --git a/internal/provider/testdata/sub2api_usage_key_limits.json b/internal/provider/testdata/sub2api_usage_key_limits.json new file mode 100644 index 00000000..12789673 --- /dev/null +++ b/internal/provider/testdata/sub2api_usage_key_limits.json @@ -0,0 +1 @@ +{"daily_usage":[{"date":"2026-10-05","requests":1,"input_tokens":547,"output_tokens":14,"cache_read_tokens":3840,"cache_write_tokens":0,"total_tokens":4401,"cost":0.005075,"actual_cost":0.005075},{"date":"2026-10-06","requests":83,"input_tokens":606765,"output_tokens":22715,"cache_read_tokens":7153280,"cache_write_tokens":0,"total_tokens":7782760,"cost":12.8069815,"actual_cost":12.8069815}],"isValid":true,"mode":"quota_limited","model_stats":[{"model":"gpt-6-astra","requests":79,"input_tokens":451420,"output_tokens":22622,"cache_creation_tokens":0,"cache_read_tokens":7145216,"total_tokens":7619258,"cost":12.790516,"actual_cost":12.790516,"account_cost":12.790516},{"model":"gpt-6-luna","requests":3,"input_tokens":155182,"output_tokens":81,"cache_creation_tokens":0,"cache_read_tokens":3840,"total_tokens":159103,"cost":0.0155971,"actual_cost":0.0155971,"account_cost":0.0155971},{"model":"gpt-5.5","requests":1,"input_tokens":547,"output_tokens":14,"cache_creation_tokens":0,"cache_read_tokens":3840,"total_tokens":4401,"cost":0.005075,"actual_cost":0.005075,"account_cost":0.005075},{"model":"gpt-6.1-sol","requests":1,"input_tokens":163,"output_tokens":12,"cache_creation_tokens":0,"cache_read_tokens":4224,"total_tokens":4399,"cost":0.0008684,"actual_cost":0.0008684,"account_cost":0.0008684}],"rate_limits":[{"limit":300,"remaining":299.9999042,"reset_at":"2026-10-07T00:00:00+08:00","used":0.0000958,"window":"1d","window_start":"2026-10-06T00:00:00+08:00"},{"limit":800,"remaining":799.9999042,"reset_at":"2026-10-13T00:00:00+08:00","used":0.0000958,"window":"7d","window_start":"2026-10-06T00:00:00+08:00"}],"status":"active","usage":{"average_duration_ms":19312.833627278073,"rpm":0,"today":{"actual_cost":12.8069815,"cache_creation_tokens":0,"cache_read_tokens":7153280,"cost":12.8069815,"input_tokens":606765,"output_tokens":22715,"requests":83,"total_tokens":7782760},"total":{"actual_cost":192.32549926,"cache_creation_tokens":0,"cache_read_tokens":196550784,"cost":192.32549926,"input_tokens":12386780,"output_tokens":1085918,"requests":1701,"total_tokens":210023482},"tpm":878}} \ No newline at end of file diff --git a/internal/provider/testdata/sub2api_usage_key_limits_not_started.json b/internal/provider/testdata/sub2api_usage_key_limits_not_started.json new file mode 100644 index 00000000..8ae2d04e --- /dev/null +++ b/internal/provider/testdata/sub2api_usage_key_limits_not_started.json @@ -0,0 +1 @@ +{"daily_usage":[{"date":"2026-10-05","requests":1,"input_tokens":547,"output_tokens":14,"cache_read_tokens":3840,"cache_write_tokens":0,"total_tokens":4401,"cost":0.005075,"actual_cost":0.005075},{"date":"2026-10-06","requests":82,"input_tokens":606216,"output_tokens":22710,"cache_read_tokens":7149440,"cache_write_tokens":0,"total_tokens":7778366,"cost":12.8068857,"actual_cost":12.8068857}],"isValid":true,"mode":"quota_limited","model_stats":[{"model":"gpt-6-astra","requests":79,"input_tokens":451420,"output_tokens":22622,"cache_creation_tokens":0,"cache_read_tokens":7145216,"total_tokens":7619258,"cost":12.790516,"actual_cost":12.790516,"account_cost":12.790516},{"model":"gpt-6-luna","requests":2,"input_tokens":154633,"output_tokens":76,"cache_creation_tokens":0,"cache_read_tokens":0,"total_tokens":154709,"cost":0.0155013,"actual_cost":0.0155013,"account_cost":0.0155013},{"model":"gpt-5.5","requests":1,"input_tokens":547,"output_tokens":14,"cache_creation_tokens":0,"cache_read_tokens":3840,"total_tokens":4401,"cost":0.005075,"actual_cost":0.005075,"account_cost":0.005075},{"model":"gpt-6.1-sol","requests":1,"input_tokens":163,"output_tokens":12,"cache_creation_tokens":0,"cache_read_tokens":4224,"total_tokens":4399,"cost":0.0008684,"actual_cost":0.0008684,"account_cost":0.0008684}],"rate_limits":[{"limit":300,"remaining":300,"used":0,"window":"1d","window_start":null},{"limit":800,"remaining":800,"used":0,"window":"7d","window_start":null}],"status":"active","usage":{"average_duration_ms":19323.06,"rpm":0,"today":{"actual_cost":12.8068857,"cache_creation_tokens":0,"cache_read_tokens":7149440,"cost":12.8068857,"input_tokens":606216,"output_tokens":22710,"requests":82,"total_tokens":7778366},"total":{"actual_cost":192.32540346,"cache_creation_tokens":0,"cache_read_tokens":196546944,"cost":192.32540346,"input_tokens":12386231,"output_tokens":1085913,"requests":1700,"total_tokens":210019088},"tpm":0}} \ No newline at end of file