From 3c0655477f9519c03d28f8910b8bdb98b3dab7ee Mon Sep 17 00:00:00 2001 From: bdchatham Date: Thu, 10 Sep 2026 16:00:50 -0700 Subject: [PATCH] fix(omni): retry a mint the server rate-limited instead of refused Measured today: rolling out several PRs in quick succession, each minting under the same client_id, tripped the server's own rate limiter -- POST /oauth/token answered 429 (slow_down). mintOnce classified every non-2xx status as a final refusal (the deliberate policy for a 5xx, which this endpoint has been measured returning for a genuinely bad credential), so a rate limit that clears in seconds read identically to a permanently wrong secret. 429 is carved out from that rule: the server names the grant valid and asks only that this wait, which has no "the credential is wrong" reading the way a 5xx might. mintToken now retries on it, honouring the server's own Retry-After when it names one (capped at 5s, so a confused or hostile server cannot hold one attempt hostage), and falling back to the existing fixed backoff when it names none. Not routed through retryUnreached: that helper's backoff is fixed, and this needs to prefer a server-named wait over it. mintToken goes back to its own loop for this reason. The lookup path's retry (using retryUnreached, unaffected) and its own APIError handling are unrelated: this is scoped to the mint path only. Co-Authored-By: Claude Sonnet 5 --- sei-agent-driver/internal/omni/mint.go | 113 ++++++++++++++++---- sei-agent-driver/internal/omni/mint_test.go | 95 ++++++++++++++++ 2 files changed, 189 insertions(+), 19 deletions(-) diff --git a/sei-agent-driver/internal/omni/mint.go b/sei-agent-driver/internal/omni/mint.go index 8fb88a88..f7517b67 100644 --- a/sei-agent-driver/internal/omni/mint.go +++ b/sei-agent-driver/internal/omni/mint.go @@ -9,6 +9,7 @@ import ( "net" "net/http" "net/url" + "strconv" "strings" "time" @@ -70,6 +71,50 @@ func (u unreached) Error() string { func (u unreached) Unwrap() error { return u.err } func (u unreached) Is(target error) bool { return target == errUnreached } +// errRateLimited marks a mint the server explicitly deferred rather than +// refused: reached, and answered on purpose "not yet, slow down" -- a third +// class distinct from [errUnreached] (never reached at all) and every other +// non-2xx status (refused outright, and not retried; see the comment on that +// branch in [mintOnce]). Measured in production: a burst of concurrent +// reviews under the same client_id got exactly this back, and until this the +// driver read it as a bad credential and gave up on a call the server asked +// it to simply try again. +var errRateLimited = errors.New("rate limited") + +// rateLimited carries a 429 and the server's own Retry-After, when it named +// one. Matched with errors.As so mintToken can read retryAfter off it, and +// with errors.Is against errRateLimited by anything that only needs to know +// it happened. +type rateLimited struct { + err error + retryAfter time.Duration // zero when the server named none +} + +func (r rateLimited) Error() string { return r.err.Error() } +func (r rateLimited) Unwrap() error { return r.err } +func (r rateLimited) Is(target error) bool { return target == errRateLimited } + +// maxRateLimitWait caps how long one attempt waits on the server's own +// Retry-After. Honouring it uncapped would let a single header value hold a +// run hostage for however long the server named; the cap keeps this attempt's +// wait inside the same order of magnitude as the fixed backoff below it uses +// when the server names nothing. +const maxRateLimitWait = 5 * time.Second + +// parseRetryAfter reads the Retry-After header as a delay in seconds (RFC +// 9110 §10.2.3), which is the only form this endpoint sends. The HTTP-date +// form is not parsed: a clock skewed against either side would misread a date +// as sooner or later than the server meant, and a delay in seconds carries no +// clock to skew. Empty, unparseable, zero, or negative all return zero, which +// callers read as "the server named none." +func parseRetryAfter(v string) time.Duration { + secs, err := strconv.Atoi(strings.TrimSpace(v)) + if err != nil || secs <= 0 { + return 0 + } + return time.Duration(secs) * time.Second +} + // mintToken exchanges the machine client's credentials for a short-lived access // token at POST /oauth/token. // @@ -122,19 +167,34 @@ func mintToken( } client = &noRedirect - var token string - var ttl time.Duration - err := retryUnreached(ctx, - func(err error) bool { return errors.Is(err, errUnreached) }, - func() error { - var err error - token, ttl, err = mintOnce(ctx, client, baseURL, clientID, clientSecret) - return err - }) - if err != nil { - return "", 0, err + // Not retryUnreached: that helper's backoff is fixed, and a 429 names its + // own wait, which is a better source than a constant chosen before this + // specific slow_down existed. The two retryable reasons still share + // transportBackoff's table and attempt count for whichever one keeps + // firing without a server-named wait -- only the source of the wait + // itself diverges, not the budget. + for attempt := 1; ; attempt++ { + token, ttl, err := mintOnce(ctx, client, baseURL, clientID, clientSecret) + if err == nil { + return token, ttl, nil + } + + var limited rateLimited + retryable := errors.Is(err, errUnreached) || errors.As(err, &limited) + if attempt == transportAttempts || !retryable { + return "", 0, err + } + + wait := transportBackoff[attempt-1] + if limited.retryAfter > 0 { + wait = min(limited.retryAfter, maxRateLimitWait) + } + select { + case <-ctx.Done(): + return "", 0, err + case <-time.After(wait): + } } - return token, ttl, nil } // mintOnce is one exchange. Its failures are classified rather than merged: @@ -179,6 +239,19 @@ func mintOnce( return "", 0, fmt.Errorf("%w: reading response: %w", driver.ErrMint, err) } + if resp.StatusCode == http.StatusTooManyRequests { + // The one status that is a refusal's opposite: the server names the + // grant valid and asks only that this wait. Every other status is a + // refusal (see below) and this one status carries an explicit "not yet" + // that has no such reading -- there is no credential a 429 could be + // naming as wrong. + return "", 0, rateLimited{ + err: fmt.Errorf("%w: the token endpoint returned 429 (%s), asking to slow down", + driver.ErrMint, oauthErrorCode(body)), + retryAfter: parseRetryAfter(resp.Header.Get("Retry-After")), + } + } + if resp.StatusCode != http.StatusOK { // The OAuth error code is safe to surface and is the one thing that says // what to fix — invalid_client means the id or secret is wrong, @@ -186,13 +259,15 @@ func mintOnce( // The rest of the body is withheld: a non-2xx here need not have come // from this API at all. // - // Every status is a refusal here, including a 5xx, and none of them is - // retried. That reads backwards -- a 503 is usually a server declining to - // answer right now -- but this endpoint was measured answering 503 to a - // malformed credential, a token with a trailing newline among them. Retrying - // on status would retry the one case that cannot succeed, and would report a - // bad secret as a deployment that is down. Reachability is what decides a - // retry here, and it is decided in the transport error above. + // Every OTHER status is a refusal here, including a 5xx, and none of + // them is retried. That reads backwards -- a 503 is usually a server + // declining to answer right now -- but this endpoint was measured + // answering 503 to a malformed credential, a token with a trailing + // newline among them. Retrying on status would retry the one case that + // cannot succeed, and would report a bad secret as a deployment that is + // down. Reachability is what decides a retry here, and it is decided in + // the transport error above -- 429 is carved out above this branch + // because it is the one status that is not a refusal at all. return "", 0, fmt.Errorf("%w: the token endpoint returned %d (%s)", driver.ErrMint, resp.StatusCode, oauthErrorCode(body)) } diff --git a/sei-agent-driver/internal/omni/mint_test.go b/sei-agent-driver/internal/omni/mint_test.go index 8ea32b7a..46f9abe0 100644 --- a/sei-agent-driver/internal/omni/mint_test.go +++ b/sei-agent-driver/internal/omni/mint_test.go @@ -389,6 +389,101 @@ func TestMintDoesNotRetryARefusal(t *testing.T) { } } +// TestMintRetriesOnRateLimit pins the one carve-out from +// TestMintDoesNotRetryARefusal's rule: 429 is not a refusal, so it is retried +// where every other status is not. +func TestMintRetriesOnRateLimit(t *testing.T) { + t.Parallel() + + calls := 0 + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + if calls < 3 { + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":"slow_down"}`)) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"tok","expires_in":3600}`)) + })) + defer srv.Close() + + token, _, err := mintToken(t.Context(), srv.Client(), srv.URL, "id", "secret") + if err != nil { + t.Fatalf("a mint that succeeds on the third attempt still failed: %v", err) + } + if token != "tok" { + t.Errorf("token = %q, want tok", token) + } + if calls != 3 { + t.Errorf("made %d calls, want 3", calls) + } +} + +// TestMintHonoursRetryAfter checks the server's own wait is used in place of +// the fixed backoff, capped rather than trusted outright. +func TestMintHonoursRetryAfter(t *testing.T) { + t.Parallel() + + t.Run("under the cap", func(t *testing.T) { + t.Parallel() + calls := 0 + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + if calls == 1 { + w.Header().Set("Retry-After", "0") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":"slow_down"}`)) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"tok","expires_in":3600}`)) + })) + defer srv.Close() + + start := time.Now() + if _, _, err := mintToken(t.Context(), srv.Client(), srv.URL, "id", "secret"); err != nil { + t.Fatalf("a mint that succeeds on the second attempt still failed: %v", err) + } + // Retry-After: 0 names an immediate retry, which parseRetryAfter reads + // as "named none" (zero and negative are indistinguishable from + // absent) -- so this in fact falls to transportBackoff's first step, + // asserted here as "well under the 5s cap" rather than as instant, so + // a change to that fallback does not make this test flake. + if elapsed := time.Since(start); elapsed > maxRateLimitWait { + t.Errorf("waited %s, want under the %s cap", elapsed, maxRateLimitWait) + } + }) + + t.Run("over the cap", func(t *testing.T) { + t.Parallel() + calls := 0 + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls++ + if calls == 1 { + // Far past maxRateLimitWait: a server that is confused, or + // hostile, must not be able to hold one attempt hostage for + // however long it names. + w.Header().Set("Retry-After", "3600") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":"slow_down"}`)) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"tok","expires_in":3600}`)) + })) + defer srv.Close() + + start := time.Now() + if _, _, err := mintToken(t.Context(), srv.Client(), srv.URL, "id", "secret"); err != nil { + t.Fatalf("a mint that succeeds on the second attempt still failed: %v", err) + } + if elapsed := time.Since(start); elapsed > maxRateLimitWait+time.Second { + t.Errorf("waited %s, want capped near %s", elapsed, maxRateLimitWait) + } + }) +} + // TestMintStopsWhenTheCallerDoes keeps the backoff from outliving the run that // is waiting on it. func TestMintStopsWhenTheCallerDoes(t *testing.T) {