fix(codex): scope usage limit errors to credentials and parse flexible quota resets

- Mark Codex usage limit errors as credential-scoped across HTTP and WebSocket executors.
- Support both top-level and nested error structures with case-insensitive matching when parsing retry-after resets.
- Propagate prevalidated candidate context to session affinity and built-in selectors during auth selection.

Closes: #5529
This commit is contained in:
Luis Pater
2026-09-06 06:20:20 +08:00
parent 7c2f6ce0d1
commit 5ab0bca040
11 changed files with 682 additions and 25 deletions

View File

@@ -2,12 +2,106 @@ package executor
import (
"encoding/json"
"errors"
"net/http"
"strconv"
"strings"
"testing"
"time"
)
func TestCodexQuotaErrorCredentialScope(t *testing.T) {
quota := `{"type":"usage_limit_reached","message":"You've hit your usage limit.","resets_in_seconds":3600}`
for _, test := range []struct {
name string
path string
body string
want bool
}{
{name: "http usage limit", path: "http", body: `{"error":` + quota + `}`, want: true},
{name: "http top-level usage limit", path: "http", body: quota, want: true},
{name: "http usage limit without reset", path: "http", body: `{"error":{"type":"usage_limit_reached"}}`, want: true},
{name: "sse error", path: "terminal", body: `{"type":"error","error":` + quota + `}`, want: true},
{name: "sse response failed", path: "terminal", body: `{"type":"response.failed","response":{"error":` + quota + `}}`, want: true},
{name: "websocket error", path: "websocket", body: `{"type":"error","status":429,"error":` + quota + `}`, want: true},
{name: "websocket body error", path: "websocket", body: `{"type":"error","status":429,"body":{"error":` + quota + `}}`, want: true},
{name: "model capacity", path: "http", body: `{"error":{"message":"Selected model is at capacity. Please try a different model."}}`},
{name: "transient rate limit", path: "http", body: `{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`},
{name: "websocket connection limit", path: "websocket", body: `{"type":"error","status":429,"error":{"code":"websocket_connection_limit_reached"}}`},
{name: "generic provider error", path: "generic", body: `{"error":` + quota + `}`},
} {
t.Run(test.name, func(t *testing.T) {
var err error
switch test.path {
case "http":
err = newCodexStatusErr(http.StatusTooManyRequests, []byte(test.body))
case "terminal":
terminalErr, _, ok := codexTerminalStreamErr([]byte(test.body))
if !ok {
t.Fatal("terminal error was not recognized")
}
err = terminalErr
case "websocket":
var ok bool
err, ok = parseCodexWebsocketError([]byte(test.body))
if !ok {
t.Fatal("websocket error was not recognized")
}
default:
err = statusErr{code: http.StatusTooManyRequests, msg: test.body}
}
var scoped interface{ IsCredentialScoped() bool }
got := errors.As(err, &scoped) && scoped.IsCredentialScoped()
if got != test.want {
t.Errorf("credential scope = %t, want %t; actual error type %T", got, test.want, err)
}
if test.want && strings.Contains(test.body, "resets_in_seconds") {
var retry interface{ RetryAfter() *time.Duration }
if !errors.As(err, &retry) || retry.RetryAfter() == nil || *retry.RetryAfter() != time.Hour {
t.Errorf("real Codex quota error must preserve its one-hour reset: %v", err)
}
}
})
}
}
func TestParseCodexRetryAfterQuotaLayouts(t *testing.T) {
now := time.Unix(1_700_000_000, 0)
for _, layout := range []string{"nested", "top-level"} {
for _, test := range []struct {
name string
body string
want time.Duration
}{
{name: "relative reset", body: `{"type":"usage_limit_reached","resets_in_seconds":3600}`, want: time.Hour},
{name: "absolute reset wins", body: `{"type":"usage_limit_reached","resets_at":1700000300,"resets_in_seconds":1}`, want: 5 * time.Minute},
{name: "expired absolute reset falls back", body: `{"type":"usage_limit_reached","resets_at":1699999940,"resets_in_seconds":77}`, want: 77 * time.Second},
{name: "type matching agrees with quota classification", body: `{"type":" USAGE_LIMIT_REACHED ","resets_in_seconds":30}`, want: 30 * time.Second},
{name: "missing reset", body: `{"type":"usage_limit_reached"}`},
{name: "expired reset", body: `{"type":"usage_limit_reached","resets_at":1699999940}`},
{name: "nonpositive reset", body: `{"type":"usage_limit_reached","resets_at":0,"resets_in_seconds":-1}`},
{name: "transient limit", body: `{"type":"rate_limit_error","resets_in_seconds":30}`},
} {
t.Run(layout+"/"+test.name, func(t *testing.T) {
body := []byte(test.body)
if layout == "nested" {
body = []byte(`{"error":` + test.body + `}`)
}
got := parseCodexRetryAfter(http.StatusTooManyRequests, body, now)
if test.want == 0 {
if got != nil {
t.Fatalf("retryAfter = %v, want nil", *got)
}
return
}
if got == nil || *got != test.want {
t.Fatalf("retryAfter = %v, want %v", got, test.want)
}
})
}
}
}
func TestParseCodexRetryAfter(t *testing.T) {
now := time.Unix(1_700_000_000, 0)

View File

@@ -285,11 +285,12 @@ func codexTerminalErrorIsContextLength(body []byte) bool {
func newCodexStatusErr(statusCode int, body []byte) statusErr {
errCode := statusCode
if isCodexModelCapacityError(body) || isCodexUsageLimitError(body) {
credentialScoped := isCodexUsageLimitError(body)
if isCodexModelCapacityError(body) || credentialScoped {
errCode = http.StatusTooManyRequests
}
body = classifyCodexStatusError(errCode, body)
err := statusErr{code: errCode, msg: string(body)}
err := statusErr{code: errCode, msg: string(body), credentialScoped: credentialScoped}
if retryAfter := parseCodexRetryAfter(errCode, body, time.Now()); retryAfter != nil {
err.retryAfter = retryAfter
}
@@ -390,20 +391,22 @@ func parseCodexRetryAfter(statusCode int, errorBody []byte, now time.Time) *time
if statusCode != http.StatusTooManyRequests || len(errorBody) == 0 {
return nil
}
if strings.TrimSpace(gjson.GetBytes(errorBody, "error.type").String()) != "usage_limit_reached" {
return nil
}
if resetsAt := gjson.GetBytes(errorBody, "error.resets_at").Int(); resetsAt > 0 {
resetAtTime := time.Unix(resetsAt, 0)
if resetAtTime.After(now) {
retryAfter := resetAtTime.Sub(now)
for _, quota := range []gjson.Result{gjson.GetBytes(errorBody, "error"), gjson.ParseBytes(errorBody)} {
if !strings.EqualFold(strings.TrimSpace(quota.Get("type").String()), "usage_limit_reached") {
continue
}
if resetsAt := quota.Get("resets_at").Int(); resetsAt > 0 {
resetAtTime := time.Unix(resetsAt, 0)
if resetAtTime.After(now) {
retryAfter := resetAtTime.Sub(now)
return &retryAfter
}
}
if resetsInSeconds := quota.Get("resets_in_seconds").Int(); resetsInSeconds > 0 {
retryAfter := time.Duration(resetsInSeconds) * time.Second
return &retryAfter
}
}
if resetsInSeconds := gjson.GetBytes(errorBody, "error.resets_in_seconds").Int(); resetsInSeconds > 0 {
retryAfter := time.Duration(resetsInSeconds) * time.Second
return &retryAfter
}
return nil
}

View File

@@ -43,7 +43,7 @@ func parseCodexWebsocketError(payload []byte) (error, bool) {
out := buildCodexWebsocketErrorPayload(payload, status)
headers := parseCodexWebsocketErrorHeaders(payload)
statusError := statusErr{code: status, msg: string(out)}
statusError := statusErr{code: status, msg: string(out), credentialScoped: isCodexUsageLimitError(out)}
if retryAfter := parseCodexRetryAfter(status, out, time.Now()); retryAfter != nil {
statusError.retryAfter = retryAfter
} else if isCodexWebsocketConnectionLimitError(payload) {

View File

@@ -1011,9 +1011,10 @@ func openAICompatStreamDataError(payload []byte, eventName string) (statusErr, b
}
type statusErr struct {
code int
msg string
retryAfter *time.Duration
code int
msg string
retryAfter *time.Duration
credentialScoped bool
}
func (e statusErr) Error() string {
@@ -1024,6 +1025,7 @@ func (e statusErr) Error() string {
}
func (e statusErr) StatusCode() int { return e.code }
func (e statusErr) RetryAfter() *time.Duration { return e.retryAfter }
func (e statusErr) IsCredentialScoped() bool { return e.credentialScoped }
const openAICompatTPMFallbackRetryAfter = time.Minute

View File

@@ -3,6 +3,7 @@ package auth
import (
"context"
"errors"
"net/http"
"testing"
"time"
@@ -11,6 +12,105 @@ import (
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
)
func TestManagerAliasQuotaFailoverWithUnobservedTargetModel(t *testing.T) {
withQuotaCooldownEnabled(t)
for name, newSelector := range map[string]func() Selector{
"round-robin": func() Selector { return &RoundRobinSelector{} },
"weighted-round-robin": func() Selector { return &WeightedRoundRobinSelector{} },
"fill-first": func() Selector { return &FillFirstSelector{} },
} {
for _, path := range []string{"select", "execute", "stream"} {
t.Run(name+"/"+path, func(t *testing.T) {
const routeModel, targetModel, otherModel = "quota-route", "quota-target", "quota-other"
manager := NewManager(nil, newSelector(), nil)
manager.SetRetryConfig(3, 30*time.Second, 0)
manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{
"codex": {{Name: targetModel, Alias: routeModel, Fork: true}},
})
highID, lowID := "alias-quota-high-"+t.Name(), "alias-quota-low-"+t.Name()
for _, candidate := range []*Auth{
{ID: highID, Provider: "codex", Status: StatusActive, Attributes: map[string]string{"priority": "4", AttributeWeight: "1"}},
{ID: lowID, Provider: "codex", Status: StatusActive, Attributes: map[string]string{"priority": "3", AttributeWeight: "1"}},
} {
registry.GetGlobalRegistry().RegisterClient(candidate.ID, "codex", []*registry.ModelInfo{{ID: routeModel}, {ID: targetModel}, {ID: otherModel}})
t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(candidate.ID) })
if _, errRegister := manager.Register(context.Background(), candidate); errRegister != nil {
t.Fatal(errRegister)
}
}
retryAfter := time.Hour
manager.MarkResult(context.Background(), Result{
AuthID: highID, Provider: "codex", Model: otherModel,
Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "other model rate limit"}, RetryAfter: &retryAfter,
})
high, _ := manager.GetByID(highID)
if !high.Unavailable || high.ModelStates[targetModel] != nil {
t.Fatal("expected aggregate cooldown with no state yet for the requested target")
}
var attempts []string
execute := func(_ context.Context, selected *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
attempts = append(attempts, selected.ID)
if req.Model != targetModel {
t.Errorf("upstream model = %s, want %s", req.Model, targetModel)
}
if selected.ID == highID {
return cliproxyexecutor.Response{}, streamQuotaError{
customStatusError: customStatusError{code: http.StatusTooManyRequests, msg: "account quota exhausted", retryAfter: &retryAfter},
credentialScoped: true,
}
}
return cliproxyexecutor.Response{Payload: []byte(lowID)}, nil
}
manager.RegisterExecutor(&customStreamMockExecutor{
identifier: "codex", mockCustomErrorExecutor: mockCustomErrorExecutor{executeFn: execute},
streamFn: func(ctx context.Context, selected *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
response, errExecute := execute(ctx, selected, req, opts)
if errExecute != nil {
return nil, errExecute
}
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
chunks <- cliproxyexecutor.StreamChunk{Payload: response.Payload}
close(chunks)
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
},
})
switch path {
case "select":
selected, errSelect := manager.SelectAuth(context.Background(), "codex", routeModel, cliproxyexecutor.Options{})
if errSelect != nil {
t.Fatal(errSelect)
}
if selected.ID != highID || !selected.Unavailable || !selected.Quota.Exceeded {
t.Fatalf("selection must preserve the selected auth's actual state: %+v", selected)
}
return
case "execute":
response, errExecute := manager.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: routeModel}, cliproxyexecutor.Options{})
if errExecute != nil {
t.Fatal(errExecute)
}
if string(response.Payload) != lowID {
t.Fatalf("response = %s, want healthy lower-priority account", response.Payload)
}
case "stream":
result, errStream := manager.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: routeModel}, cliproxyexecutor.Options{Stream: true})
if errStream != nil {
t.Fatal(errStream)
}
for chunk := range result.Chunks {
if chunk.Err != nil || string(chunk.Payload) != lowID {
t.Fatalf("chunk = %+v, want healthy lower-priority account", chunk)
}
}
}
if len(attempts) != 2 || attempts[0] != highID || attempts[1] != lowID {
t.Fatalf("attempts = %v, want exhausted account then healthy account", attempts)
}
})
}
}
}
func TestManagerExecute_ModelAliasRequestNotBlockedByOtherModelQuotaCooldown(t *testing.T) {
const (
provider = "antigravity"

View File

@@ -531,6 +531,19 @@ func selectionArgForSelector(selector Selector, routeModel string) string {
return routeModel
}
func selectorContextForAvailableAuths(ctx context.Context, selector Selector, routeModel string) context.Context {
ctx = withWeightedSelectorStateModel(ctx, selector, routeModel)
if !isBuiltInSelector(selector) {
if _, sessionAffinity := selector.(*SessionAffinitySelector); !sessionAffinity {
return ctx
}
}
if ctx == nil {
ctx = context.Background()
}
return context.WithValue(ctx, prevalidatedAuthCandidatesKey{}, true)
}
func restoreModelCooldownErrorModel(err error, requestedModel string) error {
if err == nil || requestedModel == "" {
return err
@@ -1531,7 +1544,7 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op
return nil, nil, errPick
}
if !handled {
selectorCtx := withWeightedSelectorStateModel(ctx, selector, model)
selectorCtx := selectorContextForAvailableAuths(ctx, selector, model)
selected, errPick = selector.Pick(selectorCtx, provider, selectionArgForSelector(selector, model), opts, selectorAuths)
if errPick != nil {
if isBuiltInSelector(selector) {
@@ -1864,7 +1877,7 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m
return nil, nil, "", errPick
}
if !handled {
selectorCtx := withWeightedSelectorStateModel(ctx, selector, model)
selectorCtx := selectorContextForAvailableAuths(ctx, selector, model)
selected, errPick = selector.Pick(selectorCtx, "mixed", selectionArgForSelector(selector, model), opts, selectorAuths)
if errPick != nil {
if isBuiltInSelector(selector) {

View File

@@ -0,0 +1,155 @@
package auth
import (
"context"
"fmt"
"net/http"
"testing"
"time"
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
)
func TestManagerSessionAffinityAliasCooldownPreservesSelection(t *testing.T) {
withQuotaCooldownEnabled(t)
for strategy, newFallback := range map[string]func() Selector{
"round-robin": func() Selector { return &RoundRobinSelector{} },
"weighted-round-robin": func() Selector { return &WeightedRoundRobinSelector{} },
"fill-first": func() Selector { return &FillFirstSelector{} },
} {
for _, mode := range []string{"no-session", "explicit-session", "lcp"} {
for _, path := range []string{"select", "execute", "stream"} {
t.Run(strategy+"/"+mode+"/"+path, func(t *testing.T) {
ctx := context.Background()
const routeModel, targetModel = "affinity-alias", "affinity-healthy-target"
selector := NewSessionAffinitySelector(newFallback())
t.Cleanup(selector.Stop)
manager := NewManager(nil, selector, nil)
manager.SetRetryConfig(0, 0, 0)
// Order the lower tier first to catch fallback paths that forget to
// narrow the manager's across-priority candidates before selecting.
highID, lowID := "z-high-"+t.Name(), "a-low-"+t.Name()
retryAfter := time.Hour
cool := func(id, model string) {
manager.MarkResult(ctx, Result{
AuthID: id, Provider: "codex", Model: model, RetryAfter: &retryAfter,
Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "model quota"},
})
}
for _, candidate := range []*Auth{
{ID: highID, Provider: "codex", Status: StatusActive, Attributes: map[string]string{"priority": "4", AttributeWeight: "1"}},
{ID: lowID, Provider: "codex", Status: StatusActive, Attributes: map[string]string{"priority": "3", AttributeWeight: "1"}},
} {
registry.GetGlobalRegistry().RegisterClient(candidate.ID, "codex", []*registry.ModelInfo{{ID: routeModel}, {ID: targetModel}})
t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(candidate.ID) })
if _, errRegister := manager.Register(ctx, candidate); errRegister != nil {
t.Fatal(errRegister)
}
cool(candidate.ID, routeModel)
}
// Introduce an alias while both credentials retain cooldowns under
// its old name. The newly resolved target is healthy on both.
manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{
"codex": {{Name: targetModel, Alias: routeModel, Fork: true}},
})
var attempts []string
execute := func(_ context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
attempts = append(attempts, auth.ID)
if mode == "lcp" && opts.Metadata[cliproxyexecutor.LCPAffinitySessionIDMetadataKey] == nil {
t.Error("expected executor to receive the LCP session ID")
}
if req.Model != targetModel {
t.Errorf("upstream model = %s, want %s", req.Model, targetModel)
}
return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil
}
manager.RegisterExecutor(&customStreamMockExecutor{
identifier: "codex", mockCustomErrorExecutor: mockCustomErrorExecutor{executeFn: execute},
streamFn: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
response, errExecute := execute(ctx, auth, req, opts)
if errExecute != nil {
return nil, errExecute
}
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
chunks <- cliproxyexecutor.StreamChunk{Payload: response.Payload}
close(chunks)
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
},
})
options := func(session string) cliproxyexecutor.Options {
opts := cliproxyexecutor.Options{Metadata: map[string]any{}}
switch mode {
case "explicit-session":
opts.Headers = http.Header{"X-Session-Id": []string{session}}
case "lcp":
opts.SourceFormat = sdktranslator.FormatOpenAI
opts.OriginalRequest = []byte(fmt.Sprintf(`{"messages":[{"role":"system","content":%q},{"role":"user","content":"hello"}]}`, session))
opts.Metadata[cliproxyexecutor.CallerScopeMetadataKey] = "alias-caller"
}
return opts
}
run := func(opts cliproxyexecutor.Options) (string, error) {
switch path {
case "select":
auth, errSelect := manager.SelectAuth(ctx, "codex", routeModel, opts)
if errSelect != nil {
return "", errSelect
}
return auth.ID, nil
case "execute":
response, errExecute := manager.Execute(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: routeModel}, opts)
return string(response.Payload), errExecute
default:
opts.Stream = true
result, errStream := manager.ExecuteStream(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: routeModel}, opts)
if errStream != nil {
return "", errStream
}
var payload []byte
for chunk := range result.Chunks {
if chunk.Err != nil {
t.Errorf("unexpected stream error: %v", chunk.Err)
}
payload = append(payload, chunk.Payload...)
}
return string(payload), nil
}
}
assertPick := func(label, session, wantID string) {
t.Helper()
opts := options(session)
gotID, errPick := run(opts)
if errPick != nil || gotID != wantID {
t.Fatalf("%s: auth=%q error=%v, want %s; upstream attempts=%v", label, gotID, errPick, wantID, attempts)
}
if mode == "lcp" && path == "select" && opts.Metadata[cliproxyexecutor.LCPAffinitySessionIDMetadataKey] == nil {
t.Fatal("expected LCP matching path to publish a session ID")
}
}
assertPick("cold selection ignores old alias cooldown", "stable", highID)
cool(highID, targetModel)
assertPick("actual target cooldown triggers failover", "stable", lowID)
expireSessionAffinityPriorityModelCooldown(t, manager, highID, targetModel)
wantAfterRecovery := lowID
if mode == "no-session" {
wantAfterRecovery = highID
}
assertPick("higher-priority recovery preserves binding", "stable", wantAfterRecovery)
assertPick("fresh session uses highest priority", "fresh", highID)
cool(highID, targetModel)
cool(lowID, targetModel)
before := len(attempts)
if _, errPick := run(options("stable")); statusCodeFromError(errPick) != http.StatusTooManyRequests {
t.Fatalf("all actual targets cooling: error=%v, want 429", errPick)
}
if len(attempts) != before {
t.Fatal("attempted an upstream request while all targets were cooling")
}
})
}
}
}
}

View File

@@ -139,6 +139,8 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re
rerr := resultErrorFromError(chunk.Err)
action, okAction := matchRequestScopedErrorAction(auth, chunk.Err, m.runtimeConfigSnapshot())
result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: opts}
result.RetryAfter = retryAfterFromError(chunk.Err)
result.CredentialScope = isCredentialScopedError(chunk.Err)
applyRequestScopedActionToResult(action, okAction, &result)
m.recordExecutionResult(ctx, result, auth, ephemeralResult)
}

View File

@@ -0,0 +1,117 @@
package auth
import (
"context"
"fmt"
"net/http"
"testing"
"time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
)
type streamQuotaError struct {
customStatusError
credentialScoped bool
}
func (e streamQuotaError) IsCredentialScoped() bool { return e.credentialScoped }
func TestExecuteStreamQuotaFailurePreservesCooldownAndScope(t *testing.T) {
withQuotaCooldownEnabled(t)
for _, credentialScoped := range []bool{true, false} {
for _, weighted := range []bool{false, true} {
name := fmt.Sprintf("credential_scope=%t/weighted=%t", credentialScoped, weighted)
t.Run(name, func(t *testing.T) {
var selector Selector = &RoundRobinSelector{}
if weighted {
selector = &WeightedRoundRobinSelector{}
}
manager := NewManager(nil, selector, nil)
manager.SetRetryConfig(3, 30*time.Second, 0)
model := "stream-quota-model"
siblingModel := "stream-quota-sibling"
highID := "stream-quota-high-" + name
lowID := "stream-quota-low-" + name
for _, candidate := range []*Auth{
{ID: highID, Provider: "codex", Status: StatusActive, Attributes: map[string]string{"priority": "4", AttributeWeight: "1"}},
{ID: lowID, Provider: "codex", Status: StatusActive, Attributes: map[string]string{"priority": "3", AttributeWeight: "1"}},
} {
registry.GetGlobalRegistry().RegisterClient(candidate.ID, "codex", []*registry.ModelInfo{{ID: model}, {ID: siblingModel}})
t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(candidate.ID) })
if _, errRegister := manager.Register(context.Background(), candidate); errRegister != nil {
t.Fatal(errRegister)
}
}
retryAfter := time.Hour
quotaErr := streamQuotaError{
customStatusError: customStatusError{
code: http.StatusTooManyRequests, msg: `{"error":{"type":"usage_limit_reached","resets_in_seconds":3600}}`, retryAfter: &retryAfter,
},
credentialScoped: credentialScoped,
}
if !credentialScoped {
quotaErr.msg = `{"error":{"type":"rate_limit_error","message":"Model rate limit exceeded"}}`
}
var attempts []string
manager.RegisterExecutor(&customStreamMockExecutor{
identifier: "codex",
streamFn: func(_ context.Context, selected *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
attempts = append(attempts, selected.ID)
chunks := make(chan cliproxyexecutor.StreamChunk, 2)
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.created\"}\n\n")}
chunks <- cliproxyexecutor.StreamChunk{Err: quotaErr}
close(chunks)
return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
},
})
before := time.Now()
result, errStream := manager.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true})
if errStream != nil {
t.Fatalf("ExecuteStream() error = %v", errStream)
}
var payloads, failures int
for chunk := range result.Chunks {
if len(chunk.Payload) != 0 {
payloads++
}
if chunk.Err != nil {
failures++
if chunk.Err != quotaErr {
t.Fatalf("stream error = %v, want original quota error", chunk.Err)
}
}
}
if payloads != 1 || failures != 1 || len(attempts) != 1 || attempts[0] != highID {
t.Fatalf("started stream must retain its payload and error without replay: payloads=%d failures=%d attempts=%v", payloads, failures, attempts)
}
high, _ := manager.GetByID(highID)
state := high.ModelStates[model]
if state == nil || state.NextRetryAfter.Before(before.Add(retryAfter)) {
t.Errorf("model cooldown lost upstream RetryAfter: state=%+v", state)
}
if credentialScoped && (high.Quota.Reason != "credential_quota" || high.Quota.NextRecoverAt.Before(before.Add(retryAfter))) {
t.Errorf("credential cooldown lost scope or RetryAfter: quota=%+v", high.Quota)
}
for _, requestedModel := range []string{model, siblingModel} {
wantID := lowID
if requestedModel == siblingModel && !credentialScoped {
wantID = highID
}
selected, _, _, errSelect := manager.pickNextMixed(context.Background(), []string{"codex"}, requestedModel, cliproxyexecutor.Options{}, nil)
if errSelect != nil {
t.Fatalf("pickNextMixed(%s) error = %v", requestedModel, errSelect)
}
if selected.ID != wantID {
t.Errorf("pickNextMixed(%s) = %s, want %s", requestedModel, selected.ID, wantID)
}
}
if blocked, _, _ := isAuthBlockedForModel(high, model, before.Add(time.Minute)); !blocked {
t.Error("exhausted model becomes selectable before the upstream reset")
}
})
}
}
}

View File

@@ -491,6 +491,32 @@ func getAvailableAuths(auths []*Auth, provider, model string, now time.Time) ([]
return getAvailableAuthsWithPriorityMode(auths, provider, model, now, false)
}
type prevalidatedAuthCandidatesKey struct{}
func getSelectorAvailableAuths(ctx context.Context, auths []*Auth, provider, model string, now time.Time) ([]*Auth, error) {
return getSelectorAvailableAuthsWithPriorityMode(ctx, auths, provider, model, now, false)
}
func getSelectorAvailableAuthsAcrossPriorities(ctx context.Context, auths []*Auth, provider, model string, now time.Time) ([]*Auth, error) {
return getSelectorAvailableAuthsWithPriorityMode(ctx, auths, provider, model, now, true)
}
func getSelectorAvailableAuthsWithPriorityMode(ctx context.Context, auths []*Auth, provider, model string, now time.Time, allPriorities bool) ([]*Auth, error) {
if ctx != nil {
if validated, _ := ctx.Value(prevalidatedAuthCandidatesKey{}).(bool); validated && len(auths) > 0 {
// The manager already resolved each credential's upstream model and supplied
// ID-sorted candidates. Rechecking the alias or an empty model would apply
// unrelated cooldowns. Affinity bindings may span all priority tiers, but
// fallback selection must still use the highest available tier.
if !allPriorities {
return highestPriorityAuths(auths), nil
}
return auths, nil
}
}
return getAvailableAuthsWithPriorityMode(auths, provider, model, now, allPriorities)
}
func getAvailableAuthsAcrossPriorities(auths []*Auth, provider, model string, now time.Time) ([]*Auth, error) {
return getAvailableAuthsWithPriorityMode(auths, provider, model, now, true)
}
@@ -589,7 +615,7 @@ func highestPriorityAuths(auths []*Auth) []*Auth {
func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
_ = opts
now := time.Now()
available, err := getAvailableAuths(auths, provider, model, now)
available, err := getSelectorAvailableAuths(ctx, auths, provider, model, now)
if err != nil {
return nil, err
}
@@ -647,7 +673,7 @@ func positiveWeightAuths(auths []*Auth) []*Auth {
// Pick selects the next available auth using smooth weighted round-robin.
func (s *WeightedRoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
_ = opts
available, errAvailable := getAvailableAuths(positiveWeightAuths(auths), provider, model, time.Now())
available, errAvailable := getSelectorAvailableAuths(ctx, positiveWeightAuths(auths), provider, model, time.Now())
if errAvailable != nil {
return nil, errAvailable
}
@@ -787,7 +813,7 @@ func saturatingAddInt64(value, delta int64) int64 {
func (s *FillFirstSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
_ = opts
now := time.Now()
available, err := getAvailableAuths(auths, provider, model, now)
available, err := getSelectorAvailableAuths(ctx, auths, provider, model, now)
if err != nil {
return nil, err
}
@@ -983,7 +1009,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri
availabilityCandidates = positiveWeightAuths(auths)
}
if primaryID == "" {
fallbackAuths, errAvailable := getAvailableAuths(availabilityCandidates, provider, model, now)
fallbackAuths, errAvailable := getSelectorAvailableAuths(ctx, availabilityCandidates, provider, model, now)
if errAvailable != nil {
return nil, errAvailable
}
@@ -993,7 +1019,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri
// A single availability pass serves both lookups: the bound credential is validated against
// every priority tier, while the fallback selector keeps seeing only the highest tier.
available, err := getAvailableAuthsAcrossPriorities(availabilityCandidates, provider, model, now)
available, err := getSelectorAvailableAuthsAcrossPriorities(ctx, availabilityCandidates, provider, model, now)
if err != nil {
return nil, err
}
@@ -1100,7 +1126,7 @@ func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model s
if _, weighted := s.fallback.(*WeightedRoundRobinSelector); weighted {
availabilityCandidates = positiveWeightAuths(auths)
}
available, errAvailable := getAvailableAuthsAcrossPriorities(availabilityCandidates, provider, model, time.Now())
available, errAvailable := getSelectorAvailableAuthsAcrossPriorities(ctx, availabilityCandidates, provider, model, time.Now())
if errAvailable != nil {
return nil, true, errAvailable
}

View File

@@ -0,0 +1,145 @@
package test
import (
"context"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/gorilla/websocket"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor"
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
)
func TestCodexTerminalQuotaCoolsAccountAcrossModels(t *testing.T) {
for _, transport := range []string{"sse", "websocket-error", "websocket-failed"} {
t.Run(transport, func(t *testing.T) {
const model, siblingModel = "gpt-5.4", "gpt-5.4-mini"
const created = `{"type":"response.created","response":{"id":"quota-test-response"}}`
const quota = `{"type":"usage_limit_reached","message":"You've hit your usage limit.","resets_in_seconds":3600}`
const completed = `{"type":"response.completed","response":{"id":"quota-test-success","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`
attempts := make(chan string, 8)
upgrader := websocket.Upgrader{}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
account := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
attempts <- account
terminal := completed
if account == "quota-high" {
terminal = `{"type":"error","status":429,"error":` + quota + `}`
if transport == "websocket-failed" {
terminal = `{"type":"response.failed","response":{"error":` + quota + `}}`
}
}
if transport == "sse" {
w.Header().Set("Content-Type", "text/event-stream")
if _, errWrite := fmt.Fprintf(w, "data: %s\n\ndata: %s\n\n", created, terminal); errWrite != nil {
t.Errorf("write SSE: %v", errWrite)
}
return
}
conn, errUpgrade := upgrader.Upgrade(w, r, nil)
if errUpgrade != nil {
t.Errorf("upgrade websocket: %v", errUpgrade)
return
}
defer func() {
if errClose := conn.Close(); errClose != nil {
t.Errorf("close websocket: %v", errClose)
}
}()
if _, _, errRead := conn.ReadMessage(); errRead != nil {
t.Errorf("read websocket request: %v", errRead)
return
}
for _, event := range []string{created, terminal} {
if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(event)); errWrite != nil {
t.Errorf("write websocket event: %v", errWrite)
return
}
}
}))
defer server.Close()
manager := cliproxyauth.NewManager(nil, &cliproxyauth.RoundRobinSelector{}, nil)
manager.SetRetryConfig(0, 0, 0)
cfg := &config.Config{}
if transport == "sse" {
manager.RegisterExecutor(runtimeexecutor.NewCodexExecutor(cfg))
} else {
manager.RegisterExecutor(runtimeexecutor.NewCodexWebsocketsExecutor(cfg))
}
highID, lowID := "quota-high-"+transport, "quota-low-"+transport
for i, id := range []string{highID, lowID} {
key := "quota-high"
if i == 1 {
key = "quota-low"
}
registry.GetGlobalRegistry().RegisterClient(id, "codex", []*registry.ModelInfo{{ID: model}, {ID: siblingModel}})
t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(id) })
if _, errRegister := manager.Register(context.Background(), &cliproxyauth.Auth{
ID: id, Provider: "codex", Status: cliproxyauth.StatusActive,
Attributes: map[string]string{"priority": fmt.Sprint(4 - i), "base_url": server.URL, "api_key": key},
Metadata: map[string]any{"disable_cooling": false},
}); errRegister != nil {
t.Fatal(errRegister)
}
}
run := func(requestModel string) ([]byte, error) {
t.Helper()
result, errStream := manager.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{
Model: requestModel, Payload: []byte(fmt.Sprintf(`{"model":%q,"input":"hello"}`, requestModel)),
}, cliproxyexecutor.Options{Stream: true, SourceFormat: sdktranslator.FromString("openai-response")})
if errStream != nil {
t.Fatalf("stream must start before the terminal quota error: %v", errStream)
}
var payload []byte
var terminalErr error
for chunk := range result.Chunks {
payload = append(payload, chunk.Payload...)
if chunk.Err != nil {
terminalErr = chunk.Err
}
}
return payload, terminalErr
}
before := time.Now()
payload, quotaErr := run(model)
if !strings.Contains(string(payload), "response.created") || quotaErr == nil {
t.Fatalf("expected payload then terminal quota error: payload=%s error=%v", payload, quotaErr)
}
var scoped interface{ IsCredentialScoped() bool }
if !errors.As(quotaErr, &scoped) || !scoped.IsCredentialScoped() {
t.Errorf("real Codex error is missing credential scope: %T %v", quotaErr, quotaErr)
}
high, _ := manager.GetByID(highID)
if high.Quota.Reason != "credential_quota" || high.Quota.NextRecoverAt.Before(before.Add(time.Hour)) {
t.Errorf("account cooldown missing scope or upstream reset: %+v", high.Quota)
}
payload, errSibling := run(siblingModel)
if errSibling != nil || !strings.Contains(string(payload), "response.completed") {
t.Errorf("sibling model did not fail over to the healthy account: payload=%s error=%v", payload, errSibling)
}
for _, want := range []string{"quota-high", "quota-low"} {
select {
case got := <-attempts:
if got != want {
t.Errorf("upstream account = %s, want %s", got, want)
}
default:
t.Fatalf("missing upstream attempt for %s", want)
}
}
if len(attempts) != 0 {
t.Errorf("unexpected extra upstream attempts: %d", len(attempts))
}
})
}
}