lywsvip/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-09-05 23:50:29 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
main
CLIProxyAPI/internal/runtime/executor/helps/codex_quota_test.go
sususu98 ca601db05d feat: observe upstream provider quota signals (#5211) 
Codex and Claude already emit credential-level quota watermarks on ordinary
responses. CPA used to drop them. Keep the latest watermark in memory and
return it from the management auth-file API.

Hard rule: this is observation only. It must not change scheduling, cooldown
selection, or auth-file persistence.

Snapshot, not accumulation
- QuotaState now has ObservedAt and a bounded Signals map. MarkResult fills
  them from the response headers already recorded on the request.
- Signals is the current response, not a union of earlier ones. Retry-After
  and "limit reached" only appear on the response that produced them; merging
  across responses would keep an expired value forever.
- A response with no quota header (transport failure, 5xx, unrelated endpoint)
  leaves the previous snapshot in place.
- ObservedAt is the time of the current snapshot. It advances even when the
  values did not change, so a consumer can tell a fresh reading from a stale
  one.
- When two model states merge, keep the newer snapshot. Do not union keys
  captured at different times.

What is observed, and what is not
- One predicate, ProviderSupportsQuotaObservation, decides the provider set.
- Keep Codex and Claude. Drop Kimi, xAI/Grok, Antigravity, and the Gemini
  family (gemini/vertex/aistudio): their ordinary headers are not a reliable
  credential-level remaining quota.
- Count-tokens reuses the credential but is not generation traffic.
  ExecuteCount sets SkipQuotaObservation so those headers cannot replace the
  last generation snapshot. Cooldown and success/failure accounting still run.

Cooldown must not overwrite the last snapshot
- Observation writes only ObservedAt and Signals.
- Cooldown writes only Exceeded, Reason, NextRecoverAt, and BackoffLevel,
  through applyCooldownFields. Never assign a fresh QuotaState{...} over a
  live value: that would zero the snapshot on 429, Cloudflare, credential-
  scope sibling updates, and cooldown clears.
- If a credential-quota cooldown is still active, MarkResult still observes
  an already-present model state. It does not create scheduler state just to
  record a watermark.
- .cds files persist cooldownFieldsOf(Quota) only. Restore keeps the newer
  ObservedAt, so reloading cooldown cannot clobber a newer in-memory snapshot.
- cooldownQuotaEqual still ignores observation fields, so a watermark change
  cannot by itself persist cooldown or move the scheduler.
- The management payload omits every cooldown field, so it cannot be mistaken
  for scheduler state or wired back into scheduling.
- Manual ResetQuota still clears the full QuotaState.

Codex websocket events
- Codex WS reports quota as codex.rate_limits frames, not HTTP headers.
  ParseCodexQuotaEventHeaders turns one event into the same bounded header
  shape, and MergeResponseHeaders folds it into the request-scoped holder.
  additional_rate_limits is accepted as an object (websocket) or an array
  (/wham/usage).
- Parse only through AppendCodexAPIWebsocketResponse. The shared
  AppendAPIWebsocketResponse is also used by xAI, and xAI error frames really
  do carry x-ratelimit-* headers. Parsing every frame as Codex quota would
  forge Codex headers into another provider's request log.
- Also capture code_review_rate_limits.
- A malformed active-limit name drops only that one header, not the window
  watermarks parsed from the same event.
- The type discriminator scans a bounded frame prefix, not every byte of
  every frame.
- HTTP namespaces an extra limit by short name (x-codex-bengalfox-*); WS
  namespaces it by limit name (GPT-5.3-Codex-Spark). The two paths cannot
  emit the same header names. The X-Codex-Additional- prefix marks the WS
  origin, and snapshot replacement keeps the two spellings from piling up.

Hardening
- Reject observed values with control characters. These strings reach the
  plain-text request log, and Limit-Name is upstream-controlled, so CR/LF
  could forge a header line.
- When the header cap is hit, keep plan/credits/primary ahead of
  additional-limit namespaces, then sort names so truncation is deterministic.
- QuotaState.Clone deep-copies Signals and is used by Auth.Clone and
  ModelState.Clone.
- Token stores still serialize credential metadata only, so observation adds
  no auth-file writes.
2026-08-24 17:15:37 +08:00

341 lines
14 KiB
Go
Raw Permalink Blame History

package helps
import (
"context"
"net/http"
"strings"
"testing"
internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
)
func TestParseCodexQuotaEventHeadersPreservesActiveLimit(t *testing.T) {
headers := ParseCodexQuotaEventHeaders([]byte(`{
"type":"codex.rate_limits",
"plan_type":"pro",
"metered_limit_name":"codex_bengalfox",
"rate_limits":{
"primary":{"used_percent":2,"window_minutes":10080,"reset_at":1782951970},
"secondary":null
}
}`))
if headers.Get("X-Codex-Active-Limit") != "codex_bengalfox" ||
headers.Get("X-Codex-Primary-Used-Percent") != "2" ||
headers.Get("X-Codex-Primary-Window-Minutes") != "10080" ||
headers.Get("X-Codex-Primary-Reset-At") != "1782951970" ||
headers.Get("X-Codex-Plan-Type") != "pro" {
t.Fatalf("unexpected quota headers: %#v", headers)
}
}
func TestParseCodexQuotaEventHeadersPreservesAdditionalLimitsAndCredits(t *testing.T) {
headers := ParseCodexQuotaEventHeaders([]byte(`{
"type":"codex.rate_limits",
"plan_type":"pro",
"rate_limits":{
"allowed":true,
"limit_reached":false,
"primary":{"used_percent":48,"window_minutes":10080,"reset_after_seconds":523210,"reset_at":1786677299},
"secondary":null
},
"additional_rate_limits":{
"GPT-5.3-Codex-Spark":{
"allowed":true,
"limit_reached":false,
"primary":{"used_percent":3,"window_minutes":300,"reset_after_seconds":10148,"reset_at":1787231961},
"secondary":{"used_percent":63,"window_minutes":10080,"reset_after_seconds":75420,"reset_at":1787290791}
}
},
"credits":{"has_credits":false,"unlimited":false,"balance":"0"}
}`))
for key, want := range map[string]string{
"X-Codex-Primary-Used-Percent": "48",
"X-Codex-Primary-Window-Minutes": "10080",
"X-Codex-Primary-Reset-After-Seconds": "523210",
"X-Codex-Additional-GPT-5.3-Codex-Spark-Limit-Name": "GPT-5.3-Codex-Spark",
"X-Codex-Additional-GPT-5.3-Codex-Spark-Primary-Used-Percent": "3",
"X-Codex-Additional-GPT-5.3-Codex-Spark-Primary-Window-Minutes": "300",
"X-Codex-Additional-GPT-5.3-Codex-Spark-Secondary-Used-Percent": "63",
"X-Codex-Credits-Has-Credits": "false",
"X-Codex-Credits-Unlimited": "false",
"X-Codex-Credits-Balance": "0",
} {
if got := headers.Get(key); got != want {
t.Fatalf("header %s = %q, want %q; all headers = %#v", key, got, want, headers)
}
}
}
func TestParseCodexQuotaEventHeadersReadsQuotaHeadersFromErrorFrames(t *testing.T) {
headers := ParseCodexQuotaEventHeaders([]byte(`{
"type":"error",
"status_code":429,
"headers":{
"X-Codex-Primary-Used-Percent":"100",
"X-Codex-Primary-Window-Minutes":"10080",
"X-Codex-Primary-Reset-After-Seconds":"437380",
"X-Codex-Credits-Balance":"0",
"X-Codex-Turn-State":"must-not-be-observed",
"Set-Cookie":"must-not-be-observed"
}
}`))
if headers.Get("X-Codex-Primary-Used-Percent") != "100" ||
headers.Get("X-Codex-Primary-Reset-After-Seconds") != "437380" ||
headers.Get("X-Codex-Credits-Balance") != "0" {
t.Fatalf("unexpected error quota headers: %#v", headers)
}
if headers.Get("Set-Cookie") != "" || headers.Get("X-Codex-Turn-State") != "" {
t.Fatalf("non-quota error header leaked into quota observation: %#v", headers)
}
}
func TestParseCodexQuotaEventHeadersDoesNotInventActiveLimitAndRejectsIncompleteWindows(t *testing.T) {
headers := ParseCodexQuotaEventHeaders([]byte(`{
"type":"codex.rate_limits",
"rate_limits":{"primary":{"used_percent":42,"window_minutes":300,"reset_after_seconds":60}}
}`))
if headers.Get("X-Codex-Active-Limit") != "" || headers.Get("X-Codex-Primary-Reset-After-Seconds") != "60" {
t.Fatalf("unexpected quota headers: %#v", headers)
}
if got := ParseCodexQuotaEventHeaders([]byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42}}}`)); got != nil {
t.Fatalf("incomplete quota event produced headers: %#v", got)
}
headers = ParseCodexQuotaEventHeaders([]byte(`{
"type":"codex.rate_limits",
"rate_limits":{
"primary":{"used_percent":42,"window_minutes":300},
"secondary":{"used_percent":17,"window_minutes":10080,"reset_after_seconds":60}
}
}`))
if headers.Get("X-Codex-Primary-Used-Percent") != "" ||
headers.Get("X-Codex-Secondary-Used-Percent") != "17" {
t.Fatalf("partially valid quota event produced incomplete headers: %#v", headers)
}
}
func TestParseCodexQuotaEventHeadersSupportsAdditionalArrayAndCamelCase(t *testing.T) {
headers := ParseCodexQuotaEventHeaders([]byte(`{
"type":"codex.rate_limits",
"planType":"pro",
"meteredLimitName":"premium",
"rateLimit":{
"allowed":true,
"limitReached":false,
"primary":{"usedPercent":71,"windowMinutes":10080,"resetAfterSeconds":60}
},
"additionalRateLimits":[{
"limitName":"GPT-5.3-Codex-Spark",
"rateLimit":{
"primary":{"usedPercent":13,"windowMinutes":300,"resetAt":1787399817},
"secondary":{"usedPercent":51,"windowMinutes":10080,"resetAfterSeconds":120}
}
}]
}`))
for key, want := range map[string]string{
"X-Codex-Active-Limit": "premium",
"X-Codex-Plan-Type": "pro",
"X-Codex-Primary-Used-Percent": "71",
"X-Codex-Additional-GPT-5.3-Codex-Spark-Limit-Name": "GPT-5.3-Codex-Spark",
"X-Codex-Additional-GPT-5.3-Codex-Spark-Primary-Used-Percent": "13",
"X-Codex-Additional-GPT-5.3-Codex-Spark-Secondary-Reset-After-Seconds": "120",
} {
if got := headers.Get(key); got != want {
t.Fatalf("header %s = %q, want %q; all headers = %#v", key, got, want, headers)
}
}
}
func TestParseCodexQuotaEventHeadersKeepsOnlySafeErrorHeaders(t *testing.T) {
headers := ParseCodexQuotaEventHeaders([]byte(`{
"type":"error",
"headers":{
"X-Ratelimit-Remaining-Requests":"7",
"Retry-After":"60",
"X-Codex-Primary-Over-Secondary-Limit-Percent":"0",
"X-Codex-Turn-State":"secret",
"Authorization":"Bearer secret"
}
}`))
for key, want := range map[string]string{
"X-Ratelimit-Remaining-Requests": "7",
"Retry-After": "60",
"X-Codex-Primary-Over-Secondary-Limit-Percent": "0",
} {
if got := headers.Get(key); got != want {
t.Fatalf("header %s = %q, want %q", key, got, want)
}
}
if headers.Get("X-Codex-Turn-State") != "" || headers.Get("Authorization") != "" {
t.Fatalf("unsafe error headers leaked: %#v", headers)
}
}
func TestParseCodexQuotaEventHeadersNonQuotaEventDoesNotAllocate(t *testing.T) {
payload := []byte(`{"type":"response.output_text.delta","delta":"hello"}`)
allocs := testing.AllocsPerRun(1000, func() {
if headers := ParseCodexQuotaEventHeaders(payload); headers != nil {
t.Fatalf("non-quota event produced headers: %#v", headers)
}
})
if allocs != 0 {
t.Fatalf("non-quota event allocations = %v, want 0", allocs)
}
}
// Large non-quota frames must stay allocation-free: only the bounded frame
// prefix is scanned for the type discriminator.
func TestParseCodexQuotaEventHeadersLargeNonQuotaFrameDoesNotAllocate(t *testing.T) {
big := strings.Repeat("x", 4096)
payload := []byte(`{"sequence_number":7,"type":"response.output_text.delta","delta":"` + big + `"}`)
allocs := testing.AllocsPerRun(100, func() {
if headers := ParseCodexQuotaEventHeaders(payload); headers != nil {
t.Fatalf("non-quota event produced headers: %#v", headers)
}
})
if allocs != 0 {
t.Fatalf("large non-quota frame allocations = %v, want 0", allocs)
}
}
// Quota markers live in the frame prefix, so large quota events must still be parsed.
func TestParseCodexQuotaEventHeadersLargeRateLimitsFrameIsParsed(t *testing.T) {
padding := strings.Repeat("p", 4096)
payload := []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42,"limit_percent":100,"window_minutes":10080,"reset_after_seconds":3600,"used_minutes":120},"secondary":{"used_percent":1,"limit_percent":100,"window_minutes":300,"reset_after_seconds":60,"used_minutes":2},"active_limit":"primary","credits":{"total_credits":"10","available_credits":"5","pct_remaining":"50"}},"padding":"` + padding + `"}`)
headers := ParseCodexQuotaEventHeaders(payload)
if headers == nil {
t.Fatal("large rate_limits frame produced no headers")
}
if got := headers.Get("X-Codex-Primary-Used-Percent"); got != "42" {
t.Fatalf("X-Codex-Primary-Used-Percent = %q, want 42", got)
}
}
func TestParseCodexQuotaEventHeadersRejectsInvalidAndEmptyEvents(t *testing.T) {
for _, payload := range []string{
``,
`{"type":"response.completed"}`,
`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":101,"window_minutes":300,"reset_after_seconds":60}}}`,
} {
if headers := ParseCodexQuotaEventHeaders([]byte(payload)); headers != nil {
t.Fatalf("payload %q produced unexpected headers: %#v", payload, headers)
}
}
}
// A malformed active-limit name must not discard the window watermarks that were
// parsed successfully from the same event.
func TestParseCodexQuotaEventHeadersKeepsWindowsWhenActiveLimitIsInvalid(t *testing.T) {
headers := ParseCodexQuotaEventHeaders([]byte(`{
"type":"codex.rate_limits",
"metered_limit_name":"bad limit",
"rate_limits":{"primary":{"used_percent":1,"window_minutes":300,"reset_after_seconds":60}}
}`))
if headers == nil {
t.Fatal("valid window watermarks were discarded")
}
if got := headers.Get("X-Codex-Primary-Used-Percent"); got != "1" {
t.Fatalf("X-Codex-Primary-Used-Percent = %q, want \"1\"", got)
}
if got := headers.Get("X-Codex-Active-Limit"); got != "" {
t.Fatalf("malformed active limit was retained: %q", got)
}
}
// Real websocket events namespace additional limits by limit name and carry a
// code_review_rate_limits sibling; both must survive parsing.
func TestParseCodexQuotaEventHeadersObservedProductionEvent(t *testing.T) {
headers := ParseCodexQuotaEventHeaders([]byte(`{
"type":"codex.rate_limits",
"plan_type":"pro",
"rate_limits":{"allowed":true,"limit_reached":false,"primary":{"used_percent":81,"window_minutes":10080,"reset_after_seconds":137160,"reset_at":1787588999},"secondary":null},
"code_review_rate_limits":{"primary":{"used_percent":4,"window_minutes":300,"reset_after_seconds":900}},
"additional_rate_limits":{"GPT-5.3-Codex-Spark":{"allowed":true,"limit_reached":false,"primary":{"used_percent":0,"window_minutes":300,"reset_after_seconds":18000}}},
"credits":{"has_credits":false,"unlimited":false,"balance":"0"},
"promo":null
}`))
if headers == nil {
t.Fatal("observed production event produced no headers")
}
for key, want := range map[string]string{
"X-Codex-Plan-Type": "pro",
"X-Codex-Primary-Used-Percent": "81",
"X-Codex-Primary-Reset-At": "1787588999",
"X-Codex-Code-Review-Primary-Used-Percent": "4",
"X-Codex-Additional-Gpt-5.3-Codex-Spark-Limit-Name": "GPT-5.3-Codex-Spark",
"X-Codex-Additional-Gpt-5.3-Codex-Spark-Primary-Used-Percent": "0",
"X-Codex-Credits-Balance": "0",
"X-Codex-Credits-Has-Credits": "false",
} {
if got := headers.Get(key); got != want {
t.Fatalf("header %s = %q, want %q", key, got, want)
}
}
// secondary is null upstream and must not be invented.
if got := headers.Get("X-Codex-Secondary-Used-Percent"); got != "" {
t.Fatalf("null secondary window was materialized: %q", got)
}
}
// An upstream-controlled limit name containing control characters must not reach
// the plain-text request log.
func TestParseCodexQuotaEventHeadersRejectsControlCharactersInLimitName(t *testing.T) {
headers := ParseCodexQuotaEventHeaders([]byte(`{
"type":"codex.rate_limits",
"additional_rate_limits":[{"limit_name":"evil\r\nX-Injected: 1","primary":{"used_percent":3,"window_minutes":300,"reset_after_seconds":60}}]
}`))
if headers == nil {
t.Fatal("expected window watermarks to survive")
}
for key, values := range headers {
for _, value := range values {
if strings.ContainsAny(value, "\r\n") {
t.Fatalf("header %s carried control characters: %q", key, value)
}
}
}
}
// Codex quota parsing must stay on the codex-specific websocket entry point.
// The generic entry point is shared with other providers (xAI), whose frames
// must never be reinterpreted as codex quota data: an xAI error frame really
// does carry x-ratelimit-* headers, and merging them here would forge codex
// quota headers into an unrelated provider's request log.
func TestGenericWebsocketEntryPointDoesNotObserveCodexQuota(t *testing.T) {
frame := []byte(`{"type":"error","headers":{"x-ratelimit-remaining-requests":"0","retry-after":"60"}}`)
genericCtx := internallogging.WithResponseHeadersHolder(context.Background())
internallogging.SetResponseHeaders(genericCtx, http.Header{"X-Request-Id": []string{"xai-1"}})
AppendAPIWebsocketResponse(genericCtx, nil, frame)
generic := internallogging.GetResponseHeaders(genericCtx)
if got := generic.Get("X-Ratelimit-Remaining-Requests"); got != "" {
t.Fatalf("generic websocket entry point observed codex quota: %q", got)
}
if got := generic.Get("Retry-After"); got != "" {
t.Fatalf("generic websocket entry point observed retry-after: %q", got)
}
if got := generic.Get("X-Request-Id"); got != "xai-1" {
t.Fatalf("generic websocket entry point disturbed response headers: %q", got)
}
codexCtx := internallogging.WithResponseHeadersHolder(context.Background())
AppendCodexAPIWebsocketResponse(codexCtx, nil, frame)
if got := internallogging.GetResponseHeaders(codexCtx).Get("X-Ratelimit-Remaining-Requests"); got != "0" {
t.Fatalf("codex websocket entry point lost quota observation: %q", got)
}
}
func TestMergeWebsocketQuotaHeaders(t *testing.T) {
ctx := internallogging.WithResponseHeadersHolder(context.Background())
internallogging.SetResponseHeaders(ctx, http.Header{"X-Request-Id": []string{"req-1"}})
internallogging.MergeResponseHeaders(ctx, ParseCodexQuotaEventHeaders([]byte(`{
"type":"codex.rate_limits",
"rate_limits":{"primary":{"used_percent":2,"window_minutes":10080,"reset_at":1782951970}},
"metered_limit_name":"codex_bengalfox"
}`)))
headers := internallogging.GetResponseHeaders(ctx)
if headers.Get("X-Request-Id") != "req-1" || headers.Get("X-Codex-Active-Limit") != "codex_bengalfox" {
t.Fatalf("merged response headers = %#v", headers)
}
}
Reference in New Issue View Git Blame Copy Permalink