Skip to content

Commit 9483ce5

Browse files
lawrencecchenlawrencecchen
authored
Merge pull request #35 from manaflow-ai/feat-claude-529-overload-retry
Retry Claude requests on Anthropic 529/5xx overload with bounded backoff
2 parents f2fc755 + 31fc08a commit 9483ce5

2 files changed

Lines changed: 299 additions & 0 deletions

File tree

‎internal/proxy/claude_ratelimit_routing_test.go‎

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -722,3 +722,203 @@ func TestClientRemoteIP(t *testing.T) {
722722
}
723723
}
724724
}
725+
726+
// recordSleep returns a sleep fn that records waits without actually sleeping.
727+
func recordSleep(waits *[]time.Duration) func(context.Context, time.Duration) error {
728+
return func(ctx context.Context, d time.Duration) error {
729+
*waits = append(*waits, d)
730+
return ctx.Err()
731+
}
732+
}
733+
734+
// TestClaudeOverloadRetrySucceeds: a brief 529 blip is absorbed on the SAME
735+
// account with backoff; the client sees the eventual 200 and the account is
736+
// never marked exhausted.
737+
func TestClaudeOverloadRetrySucceeds(t *testing.T) {
738+
server, store := claudeFailoverServer(t)
739+
if _, err := store.Put("claude", "session-529", "cooked@example.com", ""); err != nil {
740+
t.Fatal(err)
741+
}
742+
var calls int
743+
stub := &stubRoundTripper{responses: func(req *http.Request) *http.Response {
744+
calls++
745+
if calls <= 2 {
746+
return &http.Response{StatusCode: 529, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(`{"type":"error","error":{"type":"overloaded_error"}}`))}
747+
}
748+
return &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(`{"id":"msg_ok"}`))}
749+
}}
750+
var waits []time.Duration
751+
transport := usageLimitRetryTransport{
752+
base: stub, server: &server, provider: accounts.ProviderClaude,
753+
agent: "claude", session: "session-529", account: "cooked@example.com",
754+
method: http.MethodPost, path: "/v1/messages", maxAttempts: 6,
755+
sleep: recordSleep(&waits),
756+
}
757+
req, _ := http.NewRequest(http.MethodPost, "https://api.anthropic.com/v1/messages", bytes.NewReader([]byte(`{"model":"claude-opus-4-8"}`)))
758+
req.Header.Set("Authorization", "Bearer tok-cooked")
759+
response, err := transport.RoundTrip(req)
760+
if err != nil {
761+
t.Fatal(err)
762+
}
763+
defer response.Body.Close()
764+
if response.StatusCode != http.StatusOK {
765+
t.Fatalf("status = %d, want 200 after overload retries", response.StatusCode)
766+
}
767+
if calls != 3 {
768+
t.Fatalf("upstream calls = %d, want 3 (529, 529, 200)", calls)
769+
}
770+
if len(waits) != 2 || waits[0] != time.Second || waits[1] != 2*time.Second {
771+
t.Fatalf("backoff waits = %v, want [1s 2s]", waits)
772+
}
773+
if server.SchedulerRef.Get().Exhausted(accounts.ProviderClaude, "cooked@example.com") {
774+
t.Fatal("an overload must NOT mark the account exhausted")
775+
}
776+
}
777+
778+
// TestClaudeOverloadRetryGivesUpAfterBudget: a sustained outage passes the 5xx
779+
// through after the small retry budget so the client's own backoff takes over.
780+
func TestClaudeOverloadRetryGivesUpAfterBudget(t *testing.T) {
781+
server, _ := claudeFailoverServer(t)
782+
var calls int
783+
stub := &stubRoundTripper{responses: func(req *http.Request) *http.Response {
784+
calls++
785+
return &http.Response{StatusCode: 529, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(`{"type":"error","error":{"type":"overloaded_error"}}`))}
786+
}}
787+
var waits []time.Duration
788+
transport := usageLimitRetryTransport{
789+
base: stub, server: &server, provider: accounts.ProviderClaude,
790+
agent: "claude", session: "s", account: "cooked@example.com",
791+
method: http.MethodPost, path: "/v1/messages", maxAttempts: 6,
792+
sleep: recordSleep(&waits),
793+
}
794+
req, _ := http.NewRequest(http.MethodPost, "https://api.anthropic.com/v1/messages", bytes.NewReader([]byte(`{}`)))
795+
req.Header.Set("Authorization", "Bearer tok-cooked")
796+
response, err := transport.RoundTrip(req)
797+
if err != nil {
798+
t.Fatal(err)
799+
}
800+
defer response.Body.Close()
801+
if response.StatusCode != 529 {
802+
t.Fatalf("status = %d, want 529 passed through after budget", response.StatusCode)
803+
}
804+
if calls != 1+claudeOverloadMaxRetries {
805+
t.Fatalf("upstream calls = %d, want %d", calls, 1+claudeOverloadMaxRetries)
806+
}
807+
}
808+
809+
// TestClaudeOverloadRetryPreservesFailoverAccount is the header-revert
810+
// regression: after failover moved the request to a second account, an overload
811+
// retry must keep that account's auth, not silently revert to the first.
812+
func TestClaudeOverloadRetryPreservesFailoverAccount(t *testing.T) {
813+
server, store := claudeFailoverServer(t)
814+
if _, err := store.Put("claude", "session-mix", "cooked@example.com", ""); err != nil {
815+
t.Fatal(err)
816+
}
817+
var freshCalls int
818+
var authsSeen []string
819+
stub := &stubRoundTripper{responses: func(req *http.Request) *http.Response {
820+
auth := req.Header.Get("Authorization")
821+
authsSeen = append(authsSeen, auth)
822+
if strings.Contains(auth, "tok-cooked") {
823+
h := http.Header{}
824+
h.Set("Anthropic-Ratelimit-Unified-Status", "rejected")
825+
return &http.Response{StatusCode: http.StatusTooManyRequests, Header: h, Body: io.NopCloser(strings.NewReader(`{"type":"error","error":{"type":"rate_limit_error"}}`))}
826+
}
827+
freshCalls++
828+
if freshCalls == 1 {
829+
return &http.Response{StatusCode: 529, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(`{"type":"error","error":{"type":"overloaded_error"}}`))}
830+
}
831+
return &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(`{"id":"msg_fresh"}`))}
832+
}}
833+
var waits []time.Duration
834+
transport := usageLimitRetryTransport{
835+
base: stub, server: &server, provider: accounts.ProviderClaude,
836+
agent: "claude", session: "session-mix", account: "cooked@example.com",
837+
method: http.MethodPost, path: "/v1/messages", maxAttempts: 6,
838+
sleep: recordSleep(&waits),
839+
}
840+
req, _ := http.NewRequest(http.MethodPost, "https://api.anthropic.com/v1/messages", bytes.NewReader([]byte(`{}`)))
841+
req.Header.Set("Authorization", "Bearer tok-cooked")
842+
response, err := transport.RoundTrip(req)
843+
if err != nil {
844+
t.Fatal(err)
845+
}
846+
defer response.Body.Close()
847+
body, _ := io.ReadAll(response.Body)
848+
if response.StatusCode != http.StatusOK || !strings.Contains(string(body), "msg_fresh") {
849+
t.Fatalf("status=%d body=%s, want fresh 200 (429 -> failover -> 529 -> retry same fresh account)", response.StatusCode, string(body))
850+
}
851+
// Sequence must be cooked, fresh (529), fresh (200): the overload retry
852+
// stays on the failover account.
853+
if len(authsSeen) != 3 || !strings.Contains(authsSeen[1], "tok-fresh") || !strings.Contains(authsSeen[2], "tok-fresh") {
854+
t.Fatalf("auth sequence = %v, want cooked then fresh twice", authsSeen)
855+
}
856+
}
857+
858+
func TestClaudeOverloadBackoff(t *testing.T) {
859+
h := http.Header{}
860+
if d := claudeOverloadBackoff(h, 0); d != time.Second {
861+
t.Fatalf("retry0 = %v, want 1s", d)
862+
}
863+
if d := claudeOverloadBackoff(h, 1); d != 2*time.Second {
864+
t.Fatalf("retry1 = %v, want 2s", d)
865+
}
866+
h.Set("Retry-After", "3")
867+
if d := claudeOverloadBackoff(h, 0); d != 3*time.Second {
868+
t.Fatalf("retry-after 3 = %v, want 3s", d)
869+
}
870+
h.Set("Retry-After", "9999")
871+
if d := claudeOverloadBackoff(h, 0); d != claudeOverloadMaxWait {
872+
t.Fatalf("retry-after 9999 = %v, want cap %v", d, claudeOverloadMaxWait)
873+
}
874+
if !claudeOverloadStatus(529) || !claudeOverloadStatus(500) || claudeOverloadStatus(429) || claudeOverloadStatus(200) {
875+
t.Fatal("claudeOverloadStatus classification wrong")
876+
}
877+
}
878+
879+
// TestClaudeRejected5xxFailsOverNotOverloadRetried: a 5xx that carries the
880+
// rejected unified-status header is a depleted account, not API overload, so it
881+
// must fail over to a healthy account instead of burning same-account retries.
882+
func TestClaudeRejected5xxFailsOverNotOverloadRetried(t *testing.T) {
883+
server, store := claudeFailoverServer(t)
884+
if _, err := store.Put("claude", "session-r5", "cooked@example.com", ""); err != nil {
885+
t.Fatal(err)
886+
}
887+
var cookedCalls, freshCalls int
888+
stub := &stubRoundTripper{responses: func(req *http.Request) *http.Response {
889+
if strings.Contains(req.Header.Get("Authorization"), "tok-cooked") {
890+
cookedCalls++
891+
h := http.Header{}
892+
h.Set("Anthropic-Ratelimit-Unified-Status", "rejected")
893+
return &http.Response{StatusCode: http.StatusInternalServerError, Header: h, Body: io.NopCloser(strings.NewReader(`{"type":"error"}`))}
894+
}
895+
freshCalls++
896+
return &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(`{"id":"msg_fresh"}`))}
897+
}}
898+
var waits []time.Duration
899+
transport := usageLimitRetryTransport{
900+
base: stub, server: &server, provider: accounts.ProviderClaude,
901+
agent: "claude", session: "session-r5", account: "cooked@example.com",
902+
method: http.MethodPost, path: "/v1/messages", maxAttempts: 6,
903+
sleep: recordSleep(&waits),
904+
}
905+
req, _ := http.NewRequest(http.MethodPost, "https://api.anthropic.com/v1/messages", bytes.NewReader([]byte(`{}`)))
906+
req.Header.Set("Authorization", "Bearer tok-cooked")
907+
response, err := transport.RoundTrip(req)
908+
if err != nil {
909+
t.Fatal(err)
910+
}
911+
defer response.Body.Close()
912+
if response.StatusCode != http.StatusOK {
913+
t.Fatalf("status = %d, want 200 via failover", response.StatusCode)
914+
}
915+
if cookedCalls != 1 || freshCalls != 1 {
916+
t.Fatalf("calls cooked=%d fresh=%d, want 1/1 (failover, not same-account overload retry)", cookedCalls, freshCalls)
917+
}
918+
if len(waits) != 0 {
919+
t.Fatalf("overload backoff waits = %v, want none for a rejected 5xx", waits)
920+
}
921+
if !server.SchedulerRef.Get().Exhausted(accounts.ProviderClaude, "cooked@example.com") {
922+
t.Fatal("a rejected 5xx must mark the depleted account exhausted")
923+
}
924+
}

‎internal/proxy/proxy.go‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"net/http/httputil"
2020
"net/url"
2121
"sort"
22+
"strconv"
2223
"strings"
2324
"sync"
2425
"sync/atomic"
@@ -2666,6 +2667,64 @@ type usageLimitRetryTransport struct {
26662667
path string
26672668
upstream string
26682669
maxAttempts int
2670+
// sleep waits for the backoff duration or until the context is cancelled.
2671+
// Injectable for tests; nil means a real timer wait.
2672+
sleep func(context.Context, time.Duration) error
2673+
}
2674+
2675+
// claudeOverloadMaxRetries bounds the same-account retries subrouter itself
2676+
// performs on an Anthropic 5xx/529 overload before passing the error through to
2677+
// the client (which retries further with its own backoff). Small on purpose:
2678+
// subrouter absorbs brief blips without stacking long waits on top of the
2679+
// client's retry budget or amplifying load during a sustained outage.
2680+
const claudeOverloadMaxRetries = 2
2681+
2682+
// claudeOverloadMaxWait caps a single overload backoff wait, including one
2683+
// requested via Retry-After, so a pathological header cannot hold a proxied
2684+
// request hostage.
2685+
const claudeOverloadMaxWait = 10 * time.Second
2686+
2687+
// claudeOverloadStatus reports whether the upstream response is an
2688+
// Anthropic-side server failure (529 overloaded_error or another 5xx) worth
2689+
// retrying on the SAME account: it is not account-specific, so rotating
2690+
// accounts cannot help and would only burn the failover budget.
2691+
func claudeOverloadStatus(code int) bool {
2692+
return code >= 500 && code <= 599
2693+
}
2694+
2695+
// claudeOverloadBackoff picks the wait before an overload retry: Retry-After
2696+
// when the upstream sent one (capped), else 1s << retry (1s, 2s, ...).
2697+
func claudeOverloadBackoff(header http.Header, retry int) time.Duration {
2698+
if ra := strings.TrimSpace(claudeHeaderGet(header, "Retry-After")); ra != "" {
2699+
if secs, err := strconv.Atoi(ra); err == nil && secs > 0 {
2700+
wait := time.Duration(secs) * time.Second
2701+
if wait > claudeOverloadMaxWait {
2702+
return claudeOverloadMaxWait
2703+
}
2704+
return wait
2705+
}
2706+
}
2707+
wait := time.Second << retry
2708+
if wait > claudeOverloadMaxWait {
2709+
return claudeOverloadMaxWait
2710+
}
2711+
return wait
2712+
}
2713+
2714+
// sleepCtx waits for d or until ctx is cancelled, using the injected sleep when
2715+
// present (tests) and a real timer otherwise.
2716+
func (t usageLimitRetryTransport) sleepCtx(ctx context.Context, d time.Duration) error {
2717+
if t.sleep != nil {
2718+
return t.sleep(ctx, d)
2719+
}
2720+
timer := time.NewTimer(d)
2721+
defer timer.Stop()
2722+
select {
2723+
case <-ctx.Done():
2724+
return ctx.Err()
2725+
case <-timer.C:
2726+
return nil
2727+
}
26692728
}
26702729

26712730
// responseUsageLimited reports whether the upstream response means the current
@@ -2778,11 +2837,51 @@ func (t usageLimitRetryTransport) RoundTrip(req *http.Request) (*http.Response,
27782837
if accountID != "" {
27792838
tried[accountID] = struct{}{}
27802839
}
2840+
overloadRetries := 0
27812841
for attempt := 1; attempt <= maxAttempts; attempt++ {
27822842
response, err := base.RoundTrip(attemptReq)
27832843
if err != nil || req.GetBody == nil || req.Context().Err() != nil {
27842844
return response, err
27852845
}
2846+
// Anthropic overload (529/5xx): retry the SAME account after a bounded
2847+
// backoff. Overload is API-wide, not account-specific, so no failover, no
2848+
// exhaustion-marking, and no failover-budget consumption. Once the small
2849+
// overload budget is spent the 5xx passes through and the client's own
2850+
// backoff takes over. A 5xx that carries the rejected unified-status
2851+
// header is NOT overload-retried: rejected means this account is out of
2852+
// quota regardless of HTTP status, so it falls through to the usage-limit
2853+
// path below and fails over to a healthy account instead.
2854+
if t.provider == accounts.ProviderClaude && claudeOverloadStatus(response.StatusCode) && !claudeResponseRejected(response.Header) {
2855+
if overloadRetries >= claudeOverloadMaxRetries {
2856+
return response, nil
2857+
}
2858+
wait := claudeOverloadBackoff(response.Header, overloadRetries)
2859+
overloadRetries++
2860+
if t.logger != nil {
2861+
t.logger.Warn("retrying after anthropic overload", "agent", t.agent, "session", t.session, "account", accountID, "method", t.method, "path", t.path, "upstream", t.upstream, "status", response.StatusCode, "wait", wait.String(), "overload_retry", overloadRetries, "max_overload_retries", claudeOverloadMaxRetries)
2862+
}
2863+
if response.Body != nil {
2864+
_ = response.Body.Close()
2865+
}
2866+
if sleepErr := t.sleepCtx(req.Context(), wait); sleepErr != nil {
2867+
return nil, sleepErr
2868+
}
2869+
body, bodyErr := req.GetBody()
2870+
if bodyErr != nil {
2871+
return nil, bodyErr
2872+
}
2873+
// Preserve the CURRENT attempt's headers: after an earlier account
2874+
// failover attemptReq carries that account's auth, and cloning from the
2875+
// original req would silently revert to the first account.
2876+
currentHeader := attemptReq.Header.Clone()
2877+
attemptReq = req.Clone(req.Context())
2878+
attemptReq.Body = body
2879+
attemptReq.GetBody = req.GetBody
2880+
attemptReq.ContentLength = req.ContentLength
2881+
attemptReq.Header = currentHeader
2882+
attempt-- // overload retries do not consume the account-failover budget
2883+
continue
2884+
}
27862885
usageLimited, inspectErr := t.responseUsageLimited(response)
27872886
if inspectErr != nil {
27882887
if t.logger != nil {

0 commit comments

Comments
 (0)