Skip to content

Commit a930c16

Browse files
tim-thacker-nullifyclaude
tim-thacker-nullify
and
claude
authored
fix(client): harden HTTP retry and add 401-triggered token refresh (#151)
* fix(client): harden HTTP retry and add 401-triggered token refresh The HTTP plumbing the CLI (and long-running MCP server) depends on had two reliability gaps and no test coverage. retry.go: - Make retries method-aware. 5xx is now retried only for idempotent methods (GET/HEAD/PUT/DELETE/OPTIONS/TRACE); a POST/PATCH that may have committed server-side before erroring is no longer blindly replayed (was a duplicate-scan / duplicate-autofix hazard). 429 is still retried for any method since the server rejected it without processing. - Honor the server's Retry-After header (seconds or HTTP-date), capped at maxDelay, instead of only client-side exponential backoff. - Drain the response body before retrying so the connection can be reused. - Switch to math/rand/v2. refreshing_transport.go: - React to 401. The cached token can become invalid before its 5-minute TTL (revocation, server session kill, clock skew); previously every request kept sending the dead token for up to 5 minutes. Now a 401 forces a token refresh and retries the request once. The forced refresh is storm-safe: concurrent 401s collapse to a single refresh, and the request body is buffered so it replays correctly. Adds the first tests for internal/client: retry status/method gating, Retry-After parsing + cap, body replay, context cancellation, and the 401 refresh-and-retry / storm-safety / TTL paths. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * fix(client): pass context.Background() in TTL-refresh test staticcheck SA1012: do not pass a nil Context. The TTL-expiry test calls `tr.getToken(...)` and discards the context internally (the refresh path doesn't need it), but passing nil trips the lint check. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent e67ac98 commit a930c16

4 files changed

Lines changed: 479 additions & 11 deletions

File tree

‎internal/client/refreshing_transport.go‎

Lines changed: 60 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
package client
22

33
import (
4+
"bytes"
45
"context"
6+
"io"
57
"net/http"
68
"strings"
79
"sync"
@@ -100,11 +102,65 @@ func (t *refreshingAuthTransport) getToken(ctx context.Context) string {
100102
return t.cachedToken
101103
}
102104

105+
// forceRefresh fetches a new token regardless of the cache TTL. It is storm-safe:
106+
// if another goroutine already replaced staleToken (e.g. several concurrent
107+
// requests hit a 401 at once), the already-refreshed token is reused instead of
108+
// fetching again.
109+
func (t *refreshingAuthTransport) forceRefresh(ctx context.Context, staleToken string) string {
110+
t.mu.Lock()
111+
defer t.mu.Unlock()
112+
if t.cachedToken != staleToken {
113+
return t.cachedToken
114+
}
115+
newToken, err := t.tokenProvider()
116+
if err != nil {
117+
logger.L(ctx).Warn("forced token refresh failed after 401", logger.Err(err))
118+
return ""
119+
}
120+
t.cachedToken = newToken
121+
t.cachedAt = time.Now()
122+
return newToken
123+
}
124+
103125
func (t *refreshingAuthTransport) RoundTrip(req *http.Request) (*http.Response, error) {
126+
// Buffer the body so the request can be replayed if the first attempt 401s.
127+
var bodyBytes []byte
128+
if req.Body != nil {
129+
b, err := io.ReadAll(req.Body)
130+
if err != nil {
131+
return nil, err
132+
}
133+
req.Body.Close()
134+
bodyBytes = b
135+
}
136+
137+
attempt := func(token string) (*http.Response, error) {
138+
r := req.Clone(req.Context())
139+
if bodyBytes != nil {
140+
r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
141+
r.ContentLength = int64(len(bodyBytes))
142+
}
143+
r.Header.Set("Authorization", "Bearer "+token)
144+
r.Header.Set("User-Agent", "Nullify-CLI/mcp")
145+
return t.transport.RoundTrip(r)
146+
}
147+
104148
token := t.getToken(req.Context())
149+
resp, err := attempt(token)
150+
if err != nil {
151+
return nil, err
152+
}
153+
154+
// The cached token can be invalid before its TTL elapses (revocation, server
155+
// session kill, clock skew). On a 401, force a refresh and retry once before
156+
// surfacing the failure.
157+
if resp.StatusCode == http.StatusUnauthorized {
158+
if newToken := t.forceRefresh(req.Context(), token); newToken != "" && newToken != token {
159+
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<10))
160+
resp.Body.Close()
161+
return attempt(newToken)
162+
}
163+
}
105164

106-
r := req.Clone(req.Context())
107-
r.Header.Set("Authorization", "Bearer "+token)
108-
r.Header.Set("User-Agent", "Nullify-CLI/mcp")
109-
return t.transport.RoundTrip(r)
165+
return resp, nil
110166
}

‎internal/client/refreshing_transport_test.go‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
package client
2+
3+
import (
4+
"context"
5+
"io"
6+
"net/http"
7+
"strings"
8+
"sync"
9+
"testing"
10+
"time"
11+
)
12+
13+
func newRefreshingTransport(initial string, tp TokenProvider, inner http.RoundTripper) *refreshingAuthTransport {
14+
return &refreshingAuthTransport{
15+
nullifyHost: "acme.nullify.ai",
16+
tokenProvider: tp,
17+
transport: inner,
18+
cachedToken: initial,
19+
cachedAt: time.Now(),
20+
cacheTTL: time.Hour, // keep getToken from refreshing on TTL during 401 tests
21+
}
22+
}
23+
24+
// bearerRouter returns the status mapped to the request's bearer token.
25+
func bearerRouter(byToken map[string]int) (http.RoundTripper, func() int) {
26+
var mu sync.Mutex
27+
calls := 0
28+
rt := rtFunc(func(r *http.Request) (*http.Response, error) {
29+
mu.Lock()
30+
calls++
31+
mu.Unlock()
32+
tok := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
33+
status, ok := byToken[tok]
34+
if !ok {
35+
status = http.StatusUnauthorized
36+
}
37+
return newResp(status, nil), nil
38+
})
39+
return rt, func() int { mu.Lock(); defer mu.Unlock(); return calls }
40+
}
41+
42+
func TestRefreshOn401RetriesWithNewToken(t *testing.T) {
43+
refreshCalls := 0
44+
tp := func() (string, error) { refreshCalls++; return "fresh", nil }
45+
inner, _ := bearerRouter(map[string]int{"stale": 401, "fresh": 200})
46+
tr := newRefreshingTransport("stale", tp, inner)
47+
48+
req, _ := http.NewRequest(http.MethodGet, "http://x", nil)
49+
resp, err := tr.RoundTrip(req)
50+
if err != nil {
51+
t.Fatalf("unexpected err: %v", err)
52+
}
53+
if resp.StatusCode != 200 {
54+
t.Errorf("status = %d, want 200 after refresh", resp.StatusCode)
55+
}
56+
if refreshCalls != 1 {
57+
t.Errorf("refreshCalls = %d, want 1", refreshCalls)
58+
}
59+
}
60+
61+
func TestNo401MeansNoRefresh(t *testing.T) {
62+
refreshCalls := 0
63+
tp := func() (string, error) { refreshCalls++; return "fresh", nil }
64+
inner, calls := bearerRouter(map[string]int{"good": 200})
65+
tr := newRefreshingTransport("good", tp, inner)
66+
67+
req, _ := http.NewRequest(http.MethodGet, "http://x", nil)
68+
resp, err := tr.RoundTrip(req)
69+
if err != nil {
70+
t.Fatalf("unexpected err: %v", err)
71+
}
72+
if resp.StatusCode != 200 {
73+
t.Errorf("status = %d, want 200", resp.StatusCode)
74+
}
75+
if refreshCalls != 0 {
76+
t.Errorf("refreshCalls = %d, want 0 (no 401)", refreshCalls)
77+
}
78+
if calls() != 1 {
79+
t.Errorf("inner calls = %d, want 1", calls())
80+
}
81+
}
82+
83+
func TestRefreshFailureSurfaces401(t *testing.T) {
84+
tp := func() (string, error) { return "", io.ErrUnexpectedEOF }
85+
inner, _ := bearerRouter(map[string]int{"stale": 401})
86+
tr := newRefreshingTransport("stale", tp, inner)
87+
88+
req, _ := http.NewRequest(http.MethodGet, "http://x", nil)
89+
resp, err := tr.RoundTrip(req)
90+
if err != nil {
91+
t.Fatalf("unexpected err: %v", err)
92+
}
93+
if resp.StatusCode != 401 {
94+
t.Errorf("status = %d, want 401 surfaced when refresh fails", resp.StatusCode)
95+
}
96+
}
97+
98+
func TestRefreshReplaysBodyOn401(t *testing.T) {
99+
tp := func() (string, error) { return "fresh", nil }
100+
var bodies []string
101+
var mu sync.Mutex
102+
inner := rtFunc(func(r *http.Request) (*http.Response, error) {
103+
b, _ := io.ReadAll(r.Body)
104+
mu.Lock()
105+
bodies = append(bodies, string(b))
106+
mu.Unlock()
107+
if strings.HasSuffix(r.Header.Get("Authorization"), "stale") {
108+
return newResp(401, nil), nil
109+
}
110+
return newResp(200, nil), nil
111+
})
112+
tr := newRefreshingTransport("stale", tp, inner)
113+
114+
req, _ := http.NewRequest(http.MethodPost, "http://x", strings.NewReader("payload"))
115+
resp, err := tr.RoundTrip(req)
116+
if err != nil {
117+
t.Fatalf("unexpected err: %v", err)
118+
}
119+
if resp.StatusCode != 200 {
120+
t.Errorf("status = %d, want 200", resp.StatusCode)
121+
}
122+
if len(bodies) != 2 {
123+
t.Fatalf("attempts = %d, want 2", len(bodies))
124+
}
125+
if bodies[0] != "payload" || bodies[1] != "payload" {
126+
t.Errorf("body not replayed across 401 retry: %q", bodies)
127+
}
128+
}
129+
130+
func TestForceRefreshIsStormSafe(t *testing.T) {
131+
var refreshCalls int
132+
var mu sync.Mutex
133+
tp := func() (string, error) {
134+
mu.Lock()
135+
refreshCalls++
136+
mu.Unlock()
137+
return "fresh", nil
138+
}
139+
inner, _ := bearerRouter(map[string]int{"stale": 401, "fresh": 200})
140+
tr := newRefreshingTransport("stale", tp, inner)
141+
142+
const n = 20
143+
var wg sync.WaitGroup
144+
wg.Add(n)
145+
for range n {
146+
go func() {
147+
defer wg.Done()
148+
req, _ := http.NewRequest(http.MethodGet, "http://x", nil)
149+
resp, err := tr.RoundTrip(req)
150+
if err == nil {
151+
resp.Body.Close()
152+
}
153+
}()
154+
}
155+
wg.Wait()
156+
157+
mu.Lock()
158+
defer mu.Unlock()
159+
if refreshCalls != 1 {
160+
t.Errorf("refreshCalls = %d, want 1 (storm should dedup)", refreshCalls)
161+
}
162+
}
163+
164+
func TestGetTokenRefreshesOnTTLExpiry(t *testing.T) {
165+
refreshCalls := 0
166+
tp := func() (string, error) { refreshCalls++; return "new", nil }
167+
tr := newRefreshingTransport("old", tp, rtFunc(func(*http.Request) (*http.Response, error) {
168+
return newResp(200, nil), nil
169+
}))
170+
tr.cacheTTL = time.Millisecond
171+
tr.cachedAt = time.Now().Add(-time.Hour) // expired
172+
173+
got := tr.getToken(context.Background())
174+
if got != "new" {
175+
t.Errorf("getToken = %q, want new (TTL expired)", got)
176+
}
177+
if refreshCalls != 1 {
178+
t.Errorf("refreshCalls = %d, want 1", refreshCalls)
179+
}
180+
}

‎internal/client/retry.go‎

Lines changed: 64 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,16 @@ import (
44
"bytes"
55
"io"
66
"math"
7-
"math/rand"
7+
"math/rand/v2"
88
"net/http"
9+
"strconv"
10+
"strings"
911
"time"
1012
)
1113

1214
// retryTransport wraps an http.RoundTripper and retries on 429 and 5xx errors
13-
// with exponential backoff.
15+
// with exponential backoff. 5xx is only retried for idempotent methods so a
16+
// POST/PATCH that may have committed server-side is never replayed.
1417
type retryTransport struct {
1518
transport http.RoundTripper
1619
maxRetries int
@@ -54,14 +57,17 @@ func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
5457
return nil, err
5558
}
5659

57-
if !t.shouldRetry(resp.StatusCode) || attempt == t.maxRetries {
60+
if !shouldRetry(req.Method, resp.StatusCode) || attempt == t.maxRetries {
5861
return resp, nil
5962
}
6063

61-
// Drain and close the response body before retrying
64+
delay := t.retryDelay(resp, attempt)
65+
66+
// Drain and close the response body before retrying so the underlying
67+
// connection can be reused.
68+
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<10))
6269
resp.Body.Close()
6370

64-
delay := t.backoffDelay(attempt)
6571
select {
6672
case <-time.After(delay):
6773
case <-req.Context().Done():
@@ -72,8 +78,59 @@ func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
7278
return resp, err
7379
}
7480

75-
func (t *retryTransport) shouldRetry(statusCode int) bool {
76-
return statusCode == http.StatusTooManyRequests || statusCode >= 500
81+
// shouldRetry decides whether a response is worth retrying. 429 is always safe to
82+
// retry (the server rejected the request without processing it). 5xx is retried
83+
// only for idempotent methods; replaying a POST/PATCH risks duplicating work that
84+
// may have already committed before the server errored.
85+
func shouldRetry(method string, statusCode int) bool {
86+
if statusCode == http.StatusTooManyRequests {
87+
return true
88+
}
89+
if statusCode >= 500 {
90+
return isIdempotent(method)
91+
}
92+
return false
93+
}
94+
95+
func isIdempotent(method string) bool {
96+
switch strings.ToUpper(method) {
97+
case http.MethodGet, http.MethodHead, http.MethodPut, http.MethodDelete, http.MethodOptions, http.MethodTrace:
98+
return true
99+
}
100+
return false
101+
}
102+
103+
// retryDelay honors a Retry-After header when the server provides one (capped at
104+
// maxDelay), otherwise falls back to jittered exponential backoff.
105+
func (t *retryTransport) retryDelay(resp *http.Response, attempt int) time.Duration {
106+
if d := parseRetryAfter(resp.Header.Get("Retry-After"), time.Now()); d > 0 {
107+
if d > t.maxDelay {
108+
return t.maxDelay
109+
}
110+
return d
111+
}
112+
return t.backoffDelay(attempt)
113+
}
114+
115+
// parseRetryAfter parses a Retry-After header value, which may be either an
116+
// integer number of seconds or an HTTP-date. Returns 0 if absent/invalid.
117+
func parseRetryAfter(v string, now time.Time) time.Duration {
118+
v = strings.TrimSpace(v)
119+
if v == "" {
120+
return 0
121+
}
122+
if secs, err := strconv.Atoi(v); err == nil {
123+
if secs <= 0 {
124+
return 0
125+
}
126+
return time.Duration(secs) * time.Second
127+
}
128+
if when, err := http.ParseTime(v); err == nil {
129+
if d := when.Sub(now); d > 0 {
130+
return d
131+
}
132+
}
133+
return 0
77134
}
78135

79136
func (t *retryTransport) backoffDelay(attempt int) time.Duration {

0 commit comments

Comments
 (0)