Skip to content

Expire exhaustion marks at the upstream's own reset time #38

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
lawrencecchen merged 5 commits into from
Jul 2, 2026
Merged

Expire exhaustion marks at the upstream's own reset time #38

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
8df07ce
Expire exhaustion marks at the upstream's own reset time
lawrencecchen Jul 2, 2026
430bfa1
Keep mark expiries through partial refreshes
lawrencecchen Jul 2, 2026
1618082
Document retry-once-on-lapse semantics; test the re-mark loop
lawrencecchen Jul 2, 2026
628ce6a
TTL per failure class; body-inspect re-mark keeps header expiry
lawrencecchen Jul 2, 2026
7b619e0
Distinguish fresh zeros from carried seeds; re-anchor expiries to new…
lawrencecchen Jul 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 127 additions & 0 deletions internal/proxy/claude_ratelimit_routing_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"fmt"
"io"
"log/slog"
"strconv"
"net/http"
"net/http/httptest"
"net/url"
Expand Down Expand Up @@ -922,3 +923,129 @@ func TestClaudeRejected5xxFailsOverNotOverloadRetried(t *testing.T) {
t.Fatal("a rejected 5xx must mark the depleted account exhausted")
}
}

func TestClaudeExhaustionExpiry(t *testing.T) {
now := time.Unix(1_750_000_000, 0)
h := http.Header{}
// No headers: scheduler default TTL.
if got := claudeExhaustionExpiry(h, now); !got.Equal(now.Add(selectacct.DefaultExhaustedTTL)) {
t.Fatalf("default = %v, want now+%v", got, selectacct.DefaultExhaustedTTL)
}
// Authoritative unified-reset epoch wins.
h.Set("Anthropic-Ratelimit-Unified-Reset", "1750003600")
if got := claudeExhaustionExpiry(h, now); !got.Equal(time.Unix(1_750_003_600, 0)) {
t.Fatalf("unified-reset = %v, want epoch 1750003600", got)
}
// A reset already in the past floors to now+1m (clock skew guard).
h.Set("Anthropic-Ratelimit-Unified-Reset", "1749990000")
if got := claudeExhaustionExpiry(h, now); !got.Equal(now.Add(time.Minute)) {
t.Fatalf("past reset = %v, want floor now+1m", got)
}
// Far-future reset is capped at 8d.
h.Set("Anthropic-Ratelimit-Unified-Reset", "1999999999")
if got := claudeExhaustionExpiry(h, now); !got.Equal(now.Add(8*24*time.Hour)) {
t.Fatalf("far reset = %v, want cap now+8d", got)
}
// Retry-After honored when unified-reset absent.
h = http.Header{}
h.Set("Retry-After", "300")
if got := claudeExhaustionExpiry(h, now); !got.Equal(now.Add(5*time.Minute)) {
t.Fatalf("retry-after = %v, want now+5m", got)
}
}

// TestClaudeFailoverTriesRecoveredAccount is the live regression: an account
// whose exhaustion mark has lapsed (window reset) must be reachable by failover
// again. Before the fix its zero score persisted until a successful usage
// refresh, so failover burned its attempts on genuinely-cooked accounts and
// never reached the recovered one.
func TestClaudeFailoverTriesRecoveredAccount(t *testing.T) {
store, err := session.NewStore(filepath.Join(t.TempDir(), "sessions.json"))
if err != nil {
t.Fatal(err)
}
if _, err := store.Put("claude", "session-rec", "cooked@example.com", ""); err != nil {
t.Fatal(err)
}
ref := selectacct.NewSchedulerRef(selectacct.NewScheduler(nil))
// recovered@ was marked exhausted while cooked, but its window has reset.
ref.MarkExhaustedUntil(accounts.ProviderClaude, "recovered@example.com", time.Now().Add(-time.Second))
server := Server{
Accounts: []accounts.Account{
{ID: "cooked@example.com", Provider: accounts.ProviderClaude, AuthMode: accounts.AuthModeOAuth, Token: "tok-cooked"},
{ID: "recovered@example.com", Provider: accounts.ProviderClaude, AuthMode: accounts.AuthModeOAuth, Token: "tok-recovered"},
},
Sessions: store,
SchedulerRef: ref,
}
stub := &stubRoundTripper{responses: func(req *http.Request) *http.Response {
if strings.Contains(req.Header.Get("Authorization"), "tok-recovered") {
return &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(`{"id":"msg_recovered"}`))}
}
h := http.Header{}
h.Set("Anthropic-Ratelimit-Unified-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-Reset", "1999999999")
return &http.Response{StatusCode: http.StatusTooManyRequests, Header: h, Body: io.NopCloser(strings.NewReader(`{"type":"error","error":{"type":"rate_limit_error","message":"Error"}}`))}
}}
transport := usageLimitRetryTransport{
base: stub, server: &server, provider: accounts.ProviderClaude,
agent: "claude", session: "session-rec", account: "cooked@example.com",
method: http.MethodPost, path: "/v1/messages", maxAttempts: 6,
}
req, _ := http.NewRequest(http.MethodPost, "https://api.anthropic.com/v1/messages", bytes.NewReader([]byte(`{}`)))
req.Header.Set("Authorization", "Bearer tok-cooked")
response, err := transport.RoundTrip(req)
if err != nil {
t.Fatal(err)
}
defer response.Body.Close()
body, _ := io.ReadAll(response.Body)
if response.StatusCode != http.StatusOK || !strings.Contains(string(body), "msg_recovered") {
t.Fatalf("status=%d body=%s, want the recovered account to serve after its mark lapsed", response.StatusCode, string(body))
}
}

// TestMarkTTLSelection pins the TTL class per failure kind: rate-limit marks
// expire at the upstream reset, 401 dead-credential marks get the long
// credential TTL, and re-marking through the passive body-inspect path must not
// shorten a header-derived expiry.
func TestMarkTTLSelection(t *testing.T) {
server, _ := claudeFailoverServer(t)
resetEpoch := time.Now().Add(48 * time.Hour).Unix()
h := http.Header{}
h.Set("Anthropic-Ratelimit-Unified-Status", "rejected")
h.Set("Anthropic-Ratelimit-Unified-Reset", strconv.FormatInt(resetEpoch, 10))

// Rate-limit mark: expiry = upstream reset.
server.markAccountExhaustedFromResponse(accounts.ProviderClaude, "rl@example.com", http.StatusTooManyRequests, h)
until, ok := server.SchedulerRef.ExhaustedUntilFor(accounts.ProviderClaude, "rl@example.com")
if !ok || !until.Equal(time.Unix(resetEpoch, 0)) {
t.Fatalf("rate-limit mark until=%v ok=%v, want upstream reset %v", until, ok, time.Unix(resetEpoch, 0))
}
// Re-marking with the same headers (the passive body-inspect path) must not
// shorten the expiry to the default TTL.
server.markAccountExhaustedFromResponse(accounts.ProviderClaude, "rl@example.com", http.StatusTooManyRequests, h)
until2, _ := server.SchedulerRef.ExhaustedUntilFor(accounts.ProviderClaude, "rl@example.com")
if !until2.Equal(time.Unix(resetEpoch, 0)) {
t.Fatalf("re-mark shortened expiry to %v, want %v", until2, time.Unix(resetEpoch, 0))
}

// 401 dead credential: long credential TTL, not the 10m default.
server.markAccountExhaustedFromResponse(accounts.ProviderClaude, "dead@example.com", http.StatusUnauthorized, http.Header{})
deadUntil, ok := server.SchedulerRef.ExhaustedUntilFor(accounts.ProviderClaude, "dead@example.com")
if !ok || time.Until(deadUntil) < 50*time.Minute {
t.Fatalf("401 mark until=%v (in %v), want ~1h credential TTL", deadUntil, time.Until(deadUntil))
}

// Terminal refresh failure: credential TTL; transient: short default.
server.markAccountExhaustedRefreshFailure(accounts.ProviderClaude, "invalid@example.com", fmt.Errorf("400 Bad Request: invalid_grant"))
invUntil, _ := server.SchedulerRef.ExhaustedUntilFor(accounts.ProviderClaude, "invalid@example.com")
if time.Until(invUntil) < 50*time.Minute {
t.Fatalf("terminal refresh mark in %v, want ~1h", time.Until(invUntil))
}
server.markAccountExhaustedRefreshFailure(accounts.ProviderClaude, "blip@example.com", fmt.Errorf("dial tcp: connection refused"))
blipUntil, _ := server.SchedulerRef.ExhaustedUntilFor(accounts.ProviderClaude, "blip@example.com")
if time.Until(blipUntil) > 15*time.Minute {

@cubic-dev-ai cubic-dev-ai Bot Jul 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1: The inverted comparison for the transient-refresh assertion (blip@refresh-failure) is reversed: it checks time.Until(blipUntil) > 15*time.Minute but should check < 15*time.Minute. Because transient marks default to ~10m remaining, the current condition causes a false fatal and breaks the test. Swap the comparison so the test validates the expected TTL bound.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At internal/proxy/claude_ratelimit_routing_test.go, line 1048:

<comment>The inverted comparison for the transient-refresh assertion (`blip@refresh-failure`) is reversed: it checks `time.Until(blipUntil) > 15*time.Minute` but should check `< 15*time.Minute`. Because transient marks default to ~10m remaining, the current condition causes a false fatal and breaks the test. Swap the comparison so the test validates the expected TTL bound.</comment>

<file context>
@@ -1003,3 +1004,48 @@ func TestClaudeFailoverTriesRecoveredAccount(t *testing.T) {
+	}
+	server.markAccountExhaustedRefreshFailure(accounts.ProviderClaude, "blip@example.com", fmt.Errorf("dial tcp: connection refused"))
+	blipUntil, _ := server.SchedulerRef.ExhaustedUntilFor(accounts.ProviderClaude, "blip@example.com")
+	if time.Until(blipUntil) > 15*time.Minute {
+		t.Fatalf("transient refresh mark in %v, want ~10m default", time.Until(blipUntil))
+	}
</file context>

t.Fatalf("transient refresh mark in %v, want ~10m default", time.Until(blipUntil))
}
}
98 changes: 91 additions & 7 deletions internal/proxy/proxy.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1231,6 +1231,7 @@ func (s Server) scoreAccounts(ctx context.Context, available []accounts.Account)
seed := current.ScoreFor(account.Provider, account.ID)
seed.AccountID = account.ID
seed.Provider = account.Provider
seed.Fresh = false // carried forward until a fresh fetch overwrites it
if account.AuthMode == accounts.AuthModeAPIKey {
seed.Headroom = 0.01
seed.ShortHeadroom = 0.01
Expand Down Expand Up @@ -1282,7 +1283,9 @@ func (s Server) scoreAccounts(ctx context.Context, available []accounts.Account)
continue
}
if idx, ok := scoreByID[selectacct.ScoreKey(account.Provider, account.ID)]; ok {
scores[idx] = scoreFromUsageWindows(account.Provider, account.ID, windows)
next := scoreFromUsageWindows(account.Provider, account.ID, windows)
next.Fresh = true
scores[idx] = next
scored++
}
}
Expand Down Expand Up @@ -1678,6 +1681,79 @@ func (s Server) markAccountExhausted(provider accounts.Provider, accountID strin
}
}

// markAccountExhaustedFromResponse marks the account exhausted until the reset
// time the upstream itself reported, so the mark self-expires exactly when the
// account's window recovers. Observed live: an account whose weekly window had
// reset (real quota available) stayed zero-scored for hours because marks only
// cleared on a successful usage refresh, which the loaded usage endpoint kept
// failing; failover then burned its attempts on genuinely-cooked accounts and
// never reached the recovered one.
func (s Server) markAccountExhaustedFromResponse(provider accounts.Provider, accountID string, status int, header http.Header) {
if s.SchedulerRef == nil {
return
}
if status == http.StatusUnauthorized {
// Dead/expired credential, not a rate-limit window: no reset header
// exists and it will not self-heal on a schedule. A longer TTL avoids
// probing a dead account every few minutes while still picking it back
// up within the hour after a re-auth.
s.markAccountExhaustedCredential(provider, accountID)
return
}
s.SchedulerRef.MarkExhaustedUntil(provider, accountID, claudeExhaustionExpiry(header, time.Now()))

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Avoid expiring credential failures as rate-limit windows

When Claude returns a 401 authentication_error, claudeAccountExhaustedByResponse routes it through this helper with no reset headers, so claudeExhaustionExpiry stores only DefaultExhaustedTTL; after 10 minutes SchedulerRef.Get prunes the mark and the scheduler can pick the same dead OAuth credential again. A 401 is not a quota window that will recover at reset time, so these auth failures should keep using a refresh/reauth-driven exhaustion mark rather than the new rate-limit expiry.

Useful? React with 👍 / 👎.

}

// credentialExhaustionTTL is how long an account with a dead credential
// (401 / invalid_grant) stays out of routing before one re-probe. Credentials
// only recover via human re-auth, so probes are pure overhead; but the mark
// must still lapse so a re-authed account rejoins without waiting for a
// successful usage refresh.
const credentialExhaustionTTL = time.Hour

func (s Server) markAccountExhaustedCredential(provider accounts.Provider, accountID string) {
if s.SchedulerRef == nil {
return
}
s.SchedulerRef.MarkExhaustedUntil(provider, accountID, time.Now().Add(credentialExhaustionTTL))
}

// markAccountExhaustedRefreshFailure picks the mark TTL by failure class: a
// terminal credential error gets the long credential TTL, anything transient
// gets the short default so the account rejoins quickly.
func (s Server) markAccountExhaustedRefreshFailure(provider accounts.Provider, accountID string, err error) {
if isTerminalCredentialError(err) {
s.markAccountExhaustedCredential(provider, accountID)
return
}
s.markAccountExhausted(provider, accountID)
}

// claudeExhaustionExpiry picks when an exhaustion mark should lapse:
// anthropic-ratelimit-unified-reset (epoch seconds, the authoritative window
// reset) when present, else Retry-After seconds, else the scheduler default.
// Clamped to [now+1m, now+8d]: the floor guards clock skew / already-passed
// resets, the ceiling guards a nonsense far-future header pinning an account
// out forever.
func claudeExhaustionExpiry(header http.Header, now time.Time) time.Time {
until := now.Add(selectacct.DefaultExhaustedTTL)
if raw := strings.TrimSpace(claudeHeaderGet(header, "anthropic-ratelimit-unified-reset")); raw != "" {
if epoch, err := strconv.ParseInt(raw, 10, 64); err == nil && epoch > 0 {
until = time.Unix(epoch, 0)
}
} else if ra := strings.TrimSpace(claudeHeaderGet(header, "Retry-After")); ra != "" {
if secs, err := strconv.Atoi(ra); err == nil && secs > 0 {
until = now.Add(time.Duration(secs) * time.Second)
}
Comment on lines +1739 to +1746

@coderabbitai coderabbitai Bot Jul 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Locate the relevant function and nearby uses of the headers.
FILE="internal/proxy/proxy.go"

echo "== File size =="
wc -l "$FILE"

echo
echo "== Outline =="
ast-grep outline "$FILE" --view expanded || true

echo
echo "== Relevant region =="
sed -n '1670,1735p' "$FILE" | cat -n

echo
echo "== Search for header parsing and Retry-After handling =="
rg -n 'anthropic-ratelimit-unified-reset|Retry-After|ParseTime|strconv\.Atoi|time\.Unix|DefaultExhaustedTTL' "$FILE"

Repository: manaflow-ai/subrouter

Length of output: 22700

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== Search for Retry-After parsing across repo =="
rg -n 'Retry-After|ParseTime\(' .

echo
echo "== Relevant nearby code for overload backoff =="
sed -n '2728,2748p' internal/proxy/proxy.go | cat -n

echo
echo "== Search for tests around claudeExhaustionExpiry / overload backoff =="
rg -n 'claudeExhaustionExpiry|claudeOverloadBackoff|Retry-After' internal -g '*_test.go'

Repository: manaflow-ai/subrouter

Length of output: 4202

Fall back to Retry-After when reset parsing fails.
A malformed anthropic-ratelimit-unified-reset currently bypasses the Retry-After path; run the fallback when reset parsing doesn’t succeed, not only when the header is absent.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@internal/proxy/proxy.go` around lines 1703 - 1710, The fallback logic in the
header parsing block should use Retry-After whenever
anthropic-ratelimit-unified-reset is missing or invalid, not just when it is
absent. Update the parsing around claudeHeaderGet, strconv.ParseInt, and
strconv.Atoi so that a malformed reset value does not short-circuit the
Retry-After path; only keep until from the reset header when parsing succeeds,
otherwise continue to the Retry-After branch.

}
if min := now.Add(time.Minute); until.Before(min) {
return min
}
if max := now.Add(8 * 24 * time.Hour); until.After(max) {
return max
}
return until
}

func usageLimitJSON(body []byte) bool {
var event map[string]any
if err := json.Unmarshal(body, &event); err != nil {
Expand Down Expand Up @@ -1851,7 +1927,7 @@ func (s Server) captureResponseBody(response *http.Response, agentType, sessionI
// the rejected header). A transient "allowed"/"allowed_warning" 429 still
// fails over for this request but must not mark a healthy account exhausted.
if claudeAccountExhaustedByResponse(response.StatusCode, response.Header) {
s.markAccountExhausted(provider, accountID)
s.markAccountExhaustedFromResponse(provider, accountID, response.StatusCode, response.Header)
}
// Surface the genuine upstream rate-limit signal (headers now, body
// prefix below). Anthropic conveys subscription exhaustion only via the
Expand Down Expand Up @@ -1894,7 +1970,10 @@ func (s Server) captureResponseBody(response *http.Response, agentType, sessionI
loggedBody := false
inspect = func(body []byte) {
if inspectUsageLimit && usageLimitJSON(body) {
s.markAccountExhausted(provider, accountID)
// Use the response's headers so a header-derived reset expiry set
// above is recomputed identically, not overwritten with the short
// default TTL.
s.markAccountExhaustedFromResponse(provider, accountID, response.StatusCode, response.Header)
}
// Only log the body for the original hard rate-limit statuses
// (429/401), whose body is a known rate-limit/auth error envelope. A
Expand Down Expand Up @@ -2525,7 +2604,7 @@ func (s Server) refreshSelectedAccount(ctx context.Context, provider accounts.Pr
s.Logger.Warn("selected Claude account refresh failed, trying another account", "account", account.ID, "error", err)
}
tried := map[string]struct{}{account.ID: {}}
s.markAccountExhausted(provider, account.ID)
s.markAccountExhaustedRefreshFailure(provider, account.ID, err)
lastErr := err
for {
next, pickErr := s.retryAccount(ctx, provider, agentType, sessionID, userEmail, tried)
Expand All @@ -2538,7 +2617,7 @@ func (s Server) refreshSelectedAccount(ctx context.Context, provider accounts.Pr
return refreshed, nil
}
lastErr = err
s.markAccountExhausted(provider, next.ID)
s.markAccountExhaustedRefreshFailure(provider, next.ID, err)
if s.Logger != nil {
s.Logger.Warn("retry Claude account refresh failed", "account", next.ID, "error", err)
}
Expand Down Expand Up @@ -2904,7 +2983,10 @@ func (t usageLimitRetryTransport) RoundTrip(req *http.Request) (*http.Response,
exhausted = claudeAccountExhaustedByResponse(response.StatusCode, response.Header)
}
if t.server != nil && exhausted {
t.server.markAccountExhausted(t.provider, accountID)
// Use the response's own reset time so the mark self-expires when the
// window recovers (codex responses lack these headers and fall back
// to the default TTL inside claudeExhaustionExpiry).
t.server.markAccountExhaustedFromResponse(t.provider, accountID, response.StatusCode, response.Header)
}
if attempt == maxAttempts || t.server == nil {
reason := "max_attempts"
Expand Down Expand Up @@ -3108,7 +3190,9 @@ func (s Server) oauthRetryAccount(ctx context.Context, provider accounts.Provide
lastErr = err
terminal := isTerminalCredentialError(err)
if terminal {
s.markAccountExhausted(provider, account.ID)
// Credential TTL, not the short default: a dead token only heals
// via human re-auth, so frequent probes are pure overhead.
s.markAccountExhaustedCredential(provider, account.ID)
}
if s.Logger != nil {
s.Logger.Warn("usage-limit retry skipping OAuth account with failed refresh",
Expand Down
6 changes: 6 additions & 0 deletions internal/selectacct/scheduler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,12 @@ type Score struct {
ExpiryPressure float64
Sessions int
ModelScores map[string]Score
// Fresh marks a score computed from a successful, current usage fetch, as
// opposed to a seed carried forward from the previous scheduler (fetch
// failed/stale) or a request-time exhaustion mark. Expiry reconciliation
// uses it to tell "fresh evidence re-confirmed exhausted" apart from "old
// zero score dragged along".
Fresh bool
}

type Scheduler struct {
Expand Down
Loading