various: add auth-error-retries flag to up

Add an optional auth-error-retries flag to up that allows for
controlling the amount of retries for transient errors (e.g, 5XX or
client-side deadline exceeded errors) during OAuth or workload identity
federation auth exchange operations.

This defaults to 1 (e.g., no retries). Specifying a larger number of
retries enables retrying the auth exchange with exponential backoff
instead of the default of failing after a single try.

Updates https://github.com/tailscale/tailscale/issues/20339

Signed-off-by: Mario Minardi <mario@tailscale.com>

cmd/k8s-operator/depaware.txt

@@ -716,7 +716,7 @@ tailscale.com/cmd/k8s-operator dependencies: (generated by github.com/tailscale/
         tailscale.com/appc                                           from tailscale.com/ipn/ipnlocal
      💣 tailscale.com/atomicfile                                     from tailscale.com/ipn+
         tailscale.com/client/local                                   from tailscale.com/client/tailscale+
-        tailscale.com/client/tailscale                               from tailscale.com/internal/client/tailscale
+        tailscale.com/client/tailscale                               from tailscale.com/internal/client/tailscale+
         tailscale.com/client/tailscale/apitype                       from tailscale.com/client/tailscale+
         tailscale.com/client/tailscale/v2                            from tailscale.com/cmd/k8s-operator+
         tailscale.com/client/web                                     from tailscale.com/ipn/ipnlocal
@@ -883,6 +883,7 @@ tailscale.com/cmd/k8s-operator dependencies: (generated by github.com/tailscale/
         tailscale.com/types/structs                                  from tailscale.com/control/controlclient+
         tailscale.com/types/tkatype                                  from tailscale.com/client/local+
         tailscale.com/types/views                                    from tailscale.com/appc+
+        tailscale.com/util/authretry                                 from tailscale.com/feature/oauthkey
         tailscale.com/util/backoff                                   from tailscale.com/control/controlclient+
         tailscale.com/util/bufiox                                    from tailscale.com/types/key
         tailscale.com/util/checkchange                               from tailscale.com/ipn/ipnlocal+

cmd/tailscale/cli/up.go

@@ -104,6 +104,7 @@ func newUpFlagSet(goos string, upArgs *upArgsT, cmd string) *flag.FlagSet {
 	upf.StringVar(&upArgs.clientID, "client-id", "", "Client ID used to generate authkeys via workload identity federation")
 	upf.StringVar(&upArgs.clientSecretOrFile, "client-secret", "", `Client Secret used to generate authkeys via OAuth; if it begins with "file:", then it's a path to a file containing the secret`)
 	upf.StringVar(&upArgs.idTokenOrFile, "id-token", "", `ID token from the identity provider to exchange with the control server for workload identity federation; if it begins with "file:", then it's a path to a file containing the token`)
+	upf.IntVar(&upArgs.authErrorRetries, "auth-error-retries", 1, "max number of attempts to acquire an OAuth or workload identity federation token on transient errors, using exponential backoff; values <= 1 make a single attempt")
 
 	upf.StringVar(&upArgs.server, "login-server", ipn.DefaultControlURL, "base URL of control server")
 	upf.BoolVar(&upArgs.acceptRoutes, "accept-routes", acceptRouteDefault(goos), "accept routes advertised by other Tailscale nodes")
@@ -208,6 +209,7 @@ type upArgsT struct {
 	acceptedRisks          string
 	profileName            string
 	postureChecking        bool
+	authErrorRetries       int
 }
 
 // resolveValueFromFile returns the value as-is, or if it starts with "file:",
@@ -674,6 +676,7 @@ func runUp(ctx context.Context, cmd string, args []string, upArgs upArgsT) (retE
 			authKey, err = f(ctx, tailscale.ResolveAuthKeyArgs{
 				AuthKey:          clientSecret,
 				Tags:             prefs.AdvertiseTags,
+				AuthErrorRetries: upArgs.authErrorRetries,
 			})
 			if err != nil {
 				return err
@@ -693,6 +696,7 @@ func runUp(ctx context.Context, cmd string, args []string, upArgs upArgsT) (retE
 				IDToken:          idToken,
 				Audience:         upArgs.audience,
 				Tags:             prefs.AdvertiseTags,
+				AuthErrorRetries: upArgs.authErrorRetries,
 			})
 			if err != nil {
 				return err
@@ -992,7 +996,7 @@ func addPrefFlagMapping(flagName string, prefNames ...string) {
 // correspond to an ipn.Pref.
 func preflessFlag(flagName string) bool {
 	switch flagName {
-	case "auth-key", "force-reauth", "reset", "qr", "qr-format", "json", "timeout", "accept-risk", "host-routes", "client-id", "audience", "client-secret", "id-token":
+	case "auth-key", "force-reauth", "reset", "qr", "qr-format", "json", "timeout", "accept-risk", "host-routes", "client-id", "audience", "client-secret", "id-token", "auth-error-retries":
 		return true
 	}
 	return false

cmd/tailscale/cli/up_test.go

@@ -47,8 +47,27 @@ var validUpFlags = set.Of(
 	"client-secret",
 	"id-token",
 	"audience",
+	"auth-error-retries",
 )
 
+func TestAuthErrorRetriesFlag(t *testing.T) {
+	for _, cmd := range []string{"up", "login"} {
+		t.Run(cmd, func(t *testing.T) {
+			var args upArgsT
+			fs := newUpFlagSet("linux", &args, cmd)
+			if args.authErrorRetries != 1 {
+				t.Fatalf("default auth-error-retries = %d, want 1", args.authErrorRetries)
+			}
+			if err := fs.Parse([]string{"--auth-error-retries=4"}); err != nil {
+				t.Fatal(err)
+			}
+			if args.authErrorRetries != 4 {
+				t.Errorf("auth-error-retries = %d, want 4", args.authErrorRetries)
+			}
+		})
+	}
+}
+
 // TestUpFlagSetIsFrozen complains when new flags are added to tailscale up.
 func TestUpFlagSetIsFrozen(t *testing.T) {
 	upFlagSet.VisitAll(func(f *flag.Flag) {

cmd/tailscale/depaware.txt

@@ -184,7 +184,7 @@ tailscale.com/cmd/tailscale dependencies: (generated by github.com/tailscale/dep
    L    tailscale.com/client/freedesktop                             from tailscale.com/client/systray
         tailscale.com/client/local                                   from tailscale.com/client/tailscale+
    L    tailscale.com/client/systray                                 from tailscale.com/cmd/tailscale/cli
-        tailscale.com/client/tailscale                               from tailscale.com/internal/client/tailscale
+        tailscale.com/client/tailscale                               from tailscale.com/internal/client/tailscale+
         tailscale.com/client/tailscale/apitype                       from tailscale.com/client/tailscale+
         tailscale.com/client/web                                     from tailscale.com/cmd/tailscale/cli
         tailscale.com/clientupdate                                   from tailscale.com/cmd/tailscale/cli
@@ -296,7 +296,8 @@ tailscale.com/cmd/tailscale dependencies: (generated by github.com/tailscale/dep
         tailscale.com/types/structs                                  from tailscale.com/ipn+
         tailscale.com/types/tkatype                                  from tailscale.com/types/key+
         tailscale.com/types/views                                    from tailscale.com/tailcfg+
-        tailscale.com/util/backoff                                   from tailscale.com/cmd/tailscale/cli
+        tailscale.com/util/authretry                                 from tailscale.com/feature/identityfederation+
+        tailscale.com/util/backoff                                   from tailscale.com/cmd/tailscale/cli+
         tailscale.com/util/bufiox                                    from tailscale.com/types/key
         tailscale.com/util/cibuild                                   from tailscale.com/health+
         tailscale.com/util/clientmetric                              from tailscale.com/net/netcheck+

cmd/tsidp/depaware.txt

@@ -127,7 +127,7 @@ tailscale.com/cmd/tsidp dependencies: (generated by github.com/tailscale/depawar
         tailscale.com/appc                                           from tailscale.com/ipn/ipnlocal
      💣 tailscale.com/atomicfile                                     from tailscale.com/ipn+
         tailscale.com/client/local                                   from tailscale.com/client/web+
-        tailscale.com/client/tailscale                               from tailscale.com/internal/client/tailscale
+        tailscale.com/client/tailscale                               from tailscale.com/internal/client/tailscale+
         tailscale.com/client/tailscale/apitype                       from tailscale.com/client/local+
         tailscale.com/client/web                                     from tailscale.com/ipn/ipnlocal
         tailscale.com/control/controlbase                            from tailscale.com/control/controlhttp+
@@ -268,6 +268,7 @@ tailscale.com/cmd/tsidp dependencies: (generated by github.com/tailscale/depawar
         tailscale.com/types/structs                                  from tailscale.com/control/controlclient+
         tailscale.com/types/tkatype                                  from tailscale.com/client/local+
         tailscale.com/types/views                                    from tailscale.com/appc+
+        tailscale.com/util/authretry                                 from tailscale.com/feature/oauthkey
         tailscale.com/util/backoff                                   from tailscale.com/control/controlclient+
         tailscale.com/util/bufiox                                    from tailscale.com/types/key
         tailscale.com/util/checkchange                               from tailscale.com/ipn/ipnlocal+
@@ -364,7 +365,7 @@ tailscale.com/cmd/tsidp dependencies: (generated by github.com/tailscale/depawar
         golang.org/x/net/ipv6                                        from github.com/tailscale/wireguard-go/conn+
         golang.org/x/net/proxy                                       from tailscale.com/net/netns
    D    golang.org/x/net/route                                       from tailscale.com/net/netmon+
-        golang.org/x/oauth2                                          from golang.org/x/oauth2/clientcredentials
+        golang.org/x/oauth2                                          from golang.org/x/oauth2/clientcredentials+
         golang.org/x/oauth2/clientcredentials                        from tailscale.com/feature/oauthkey
         golang.org/x/oauth2/internal                                 from golang.org/x/oauth2+
         golang.org/x/sync/errgroup                                   from github.com/mdlayher/socket+

feature/identityfederation/identityfederation.go

@@ -19,6 +19,7 @@ import (
 	"tailscale.com/feature"
 	"tailscale.com/internal/client/tailscale"
 	"tailscale.com/ipn"
+	"tailscale.com/util/authretry"
 	"tailscale.com/wif"
 )
 
@@ -58,11 +59,13 @@ func resolveAuthKey(ctx context.Context, args tailscale.ResolveAuthKeyWIFArgs) (
 		return "", fmt.Errorf("failed to parse optional config attributes: %w", err)
 	}
 
-	accessToken, err := exchangeJWTForToken(ctx, tailscale.ExchangeJWTForTokenWIFArgs{
+	accessToken, err := authretry.RetryOnTransientFailure(ctx, "exchange-jwt", args.AuthErrorRetries, func() (string, error) {
+		return exchangeJWTForToken(ctx, tailscale.ExchangeJWTForTokenWIFArgs{
 			BaseURL:  args.BaseURL,
 			ClientID: strippedID,
 			IDToken:  args.IDToken,
 		})
+	})
 	if err != nil {
 		return "", fmt.Errorf("failed to exchange JWT for access token: %w", err)
 	}

feature/oauthkey/oauthkey.go

@@ -13,9 +13,11 @@ import (
 	"strconv"
 	"strings"
 
+	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/clientcredentials"
 	"tailscale.com/feature"
 	"tailscale.com/internal/client/tailscale"
+	"tailscale.com/util/authretry"
 )
 
 func init() {
@@ -54,9 +56,15 @@ func resolveAuthKey(ctx context.Context, args tailscale.ResolveAuthKeyArgs) (str
 		TokenURL:     baseURL + "/api/v2/oauth/token",
 	}
 
+	tokenSource := retryTokenSource{
+		ctx:      ctx,
+		source:   credentials.TokenSource(ctx),
+		attempts: args.AuthErrorRetries,
+	}
+
 	tsClient := tailscale.NewClient("-", nil)
 	tsClient.UserAgent = "tailscale-cli"
-	tsClient.HTTPClient = credentials.Client(ctx)
+	tsClient.HTTPClient = oauth2.NewClient(ctx, tokenSource)
 	tsClient.BaseURL = baseURL
 
 	caps := tailscale.KeyCapabilities{
@@ -77,6 +85,21 @@ func resolveAuthKey(ctx context.Context, args tailscale.ResolveAuthKeyArgs) (str
 	return authkey, nil
 }
 
+type retryTokenSource struct {
+	ctx      context.Context
+	source   oauth2.TokenSource
+	attempts int
+}
+
+func (s retryTokenSource) Token() (token *oauth2.Token, err error) {
+	_, err = authretry.RetryOnTransientFailure(s.ctx, "oauth-token", s.attempts, func() (string, error) {
+		var tokenErr error
+		token, tokenErr = s.source.Token()
+		return "", tokenErr
+	})
+	return token, err
+}
+
 func parseOptionalAttributes(clientSecret string) (strippedSecret string, ephemeral bool, preauth bool, baseURL string, err error) {
 	strippedSecret, named, _ := strings.Cut(clientSecret, "?")
 	attrs, err := url.ParseQuery(named)

internal/client/tailscale/identityfederation.go

@@ -20,6 +20,9 @@ type ResolveAuthKeyWIFArgs struct {
 	Audience string
 	// Tags is the list of tags to be associated with the auth key.
 	Tags []string
+	// AuthErrorRetries is the maximum number of token exchange attempts on
+	// transient errors, using exponential backoff. Values <= 1 make one attempt.
+	AuthErrorRetries int
 }
 
 type ExchangeJWTForTokenWIFArgs struct {

internal/client/tailscale/oauthkeys.go

@@ -16,6 +16,9 @@ type ResolveAuthKeyArgs struct {
 	// OAuth secret case, and required to be the same as the list of tags for which the OAuth
 	// secret is allowed to issue auth keys).
 	Tags []string
+	// AuthErrorRetries is the maximum number of token acquisition attempts on
+	// transient errors, using exponential backoff. Values <= 1 make one attempt.
+	AuthErrorRetries int
 }
 
 // HookResolveAuthKey resolves to [oauthkey.ResolveAuthKey] when the

tsnet/depaware.txt

@@ -123,7 +123,7 @@ tailscale.com/tsnet dependencies: (generated by github.com/tailscale/depaware)
         tailscale.com/appc                                           from tailscale.com/ipn/ipnlocal
      💣 tailscale.com/atomicfile                                     from tailscale.com/ipn+
         tailscale.com/client/local                                   from tailscale.com/client/web+
-        tailscale.com/client/tailscale                               from tailscale.com/internal/client/tailscale
+        tailscale.com/client/tailscale                               from tailscale.com/internal/client/tailscale+
         tailscale.com/client/tailscale/apitype                       from tailscale.com/client/local+
  LDWI    tailscale.com/client/web                                     from tailscale.com/ipn/ipnlocal
         tailscale.com/control/controlbase                            from tailscale.com/control/controlhttp+
@@ -263,6 +263,7 @@ tailscale.com/tsnet dependencies: (generated by github.com/tailscale/depaware)
         tailscale.com/types/structs                                  from tailscale.com/control/controlclient+
         tailscale.com/types/tkatype                                  from tailscale.com/client/local+
         tailscale.com/types/views                                    from tailscale.com/appc+
+        tailscale.com/util/authretry                                 from tailscale.com/feature/oauthkey
         tailscale.com/util/backoff                                   from tailscale.com/control/controlclient+
         tailscale.com/util/bufiox                                    from tailscale.com/types/key
         tailscale.com/util/checkchange                               from tailscale.com/ipn/ipnlocal+
@@ -359,7 +360,7 @@ tailscale.com/tsnet dependencies: (generated by github.com/tailscale/depaware)
         golang.org/x/net/ipv6                                        from github.com/tailscale/wireguard-go/conn+
  LDW    golang.org/x/net/proxy                                       from tailscale.com/net/netns
   DI    golang.org/x/net/route                                       from tailscale.com/net/netmon+
-        golang.org/x/oauth2                                          from golang.org/x/oauth2/clientcredentials
+        golang.org/x/oauth2                                          from golang.org/x/oauth2/clientcredentials+
         golang.org/x/oauth2/clientcredentials                        from tailscale.com/feature/oauthkey
         golang.org/x/oauth2/internal                                 from golang.org/x/oauth2+
         golang.org/x/sync/errgroup                                   from github.com/mdlayher/socket+

util/authretry/authretry.go

@@ -0,0 +1,67 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+// Package authretry provides retry logic for transient errors encountered
+// during auth key generation (OAuth token fetch, WIF token exchange).
+package authretry
+
+import (
+	"context"
+	"errors"
+	"log"
+	"net/http"
+	"time"
+
+	"golang.org/x/oauth2"
+	tsclient "tailscale.com/client/tailscale"
+	"tailscale.com/util/backoff"
+)
+
+// RetryOnTransientFailure calls fn up to attempts times, retrying transient
+// errors with exponential backoff. Values <= 1 call fn exactly once.
+func RetryOnTransientFailure(ctx context.Context, name string, attempts int, fn func() (string, error)) (string, error) {
+	attempts = max(1, attempts)
+	var bo *backoff.Backoff
+	if attempts > 1 {
+		bo = backoff.NewBackoff(name, log.Printf, 30*time.Second)
+	}
+	var lastErr error
+	for i := range attempts {
+		result, err := fn()
+		if err == nil {
+			return result, nil
+		}
+		lastErr = err
+		if i == attempts-1 || !isRetriableError(err) {
+			return "", err
+		}
+		log.Printf("%s: transient error (attempt %d/%d), retrying: %v", name, i+1, attempts, err)
+		bo.BackOff(ctx, err)
+		if ctx.Err() != nil {
+			return "", ctx.Err()
+		}
+	}
+	return "", lastErr
+}
+
+// isRetriableError reports whether err is a transient error that may succeed on
+// retry: e.g., any 5xx status (from either tsclient.ErrResponse or *oauth2.RetrieveError),
+// or context.DeadlineExceeded.
+//
+// TODO(mpminardi): add-in handling for 429 if / when these are returned from the upstream
+// endpoints that are hit in the auth exchange flows.
+func isRetriableError(err error) bool {
+	if err == nil {
+		return false
+	}
+	if errors.Is(err, context.DeadlineExceeded) {
+		return true
+	}
+	if errResp, ok := errors.AsType[tsclient.ErrResponse](err); ok {
+		return errResp.Status >= http.StatusInternalServerError
+	}
+	if retrieveErr, ok := errors.AsType[*oauth2.RetrieveError](err); ok && retrieveErr.Response != nil {
+		return retrieveErr.Response.StatusCode >= http.StatusInternalServerError
+	}
+	return false
+}

util/authretry/authretry_test.go

@@ -0,0 +1,243 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+package authretry
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"golang.org/x/oauth2"
+	tsclient "tailscale.com/client/tailscale"
+)
+
+func TestIsRetriable(t *testing.T) {
+	tests := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{"nil-error", nil, false},
+		{"non-retriable-error", errors.New("something went wrong"), false},
+		{"deadline-exceeded", context.DeadlineExceeded, true},
+		{"wrapped-deadline-exceeded", fmt.Errorf("request failed: %w", context.DeadlineExceeded), true},
+		{"err-response-500", tsclient.ErrResponse{Status: 500, Message: "internal server error"}, true},
+		{"err-response-502", tsclient.ErrResponse{Status: 502, Message: "bad gateway"}, true},
+		{"err-response-503", tsclient.ErrResponse{Status: 503, Message: "service unavailable"}, true},
+		{"err-response-504", tsclient.ErrResponse{Status: 504, Message: "gateway timeout"}, true},
+		{"err-response-400", tsclient.ErrResponse{Status: 400, Message: "bad request"}, false},
+		{"err-response-401", tsclient.ErrResponse{Status: 401, Message: "unauthorized"}, false},
+		{"err-response-403", tsclient.ErrResponse{Status: 403, Message: "forbidden"}, false},
+		{"err-response-404", tsclient.ErrResponse{Status: 404, Message: "not found"}, false},
+		{"wrapped-err-response-400", fmt.Errorf("create key: %w", tsclient.ErrResponse{Status: 400}), false},
+		{
+			"oauth2-retrieve-error-503",
+			&oauth2.RetrieveError{Response: &http.Response{StatusCode: 503}, Body: []byte("service unavailable")},
+			true,
+		},
+		{
+			"oauth2-retrieve-error-400",
+			&oauth2.RetrieveError{Response: &http.Response{StatusCode: 400}, Body: []byte("bad request")},
+			false,
+		},
+		{
+			"wrapped-oauth2-retrieve-error-504",
+			fmt.Errorf("token exchange: %w", &oauth2.RetrieveError{Response: &http.Response{StatusCode: 504}, Body: []byte("gateway timeout")}),
+			true,
+		},
+		{
+			"oauth2-retrieve-error-nil-response",
+			&oauth2.RetrieveError{Response: nil, Body: []byte("no response")},
+			false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := isRetriableError(tt.err); got != tt.want {
+				t.Errorf("isRetriable() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestRetryAttemptCount(t *testing.T) {
+	for _, attempts := range []int{-1, 0, 1, 2, 4} {
+		t.Run(fmt.Sprint(attempts), func(t *testing.T) {
+			calls := 0
+			wantErr := tsclient.ErrResponse{Status: 503, Message: "unavailable"}
+			_, err := RetryOnTransientFailure(context.Background(), "test", attempts, func() (string, error) {
+				calls++
+				return "", wantErr
+			})
+			if !errors.Is(err, wantErr) {
+				t.Errorf("got error %v, want %v", err, wantErr)
+			}
+			if want := max(1, attempts); calls != want {
+				t.Errorf("got %d calls, want %d", calls, want)
+			}
+		})
+	}
+}
+
+func TestRetryOnTransient(t *testing.T) {
+	t.Run("calls-once-on-retriable-error-with-one-attempt", func(t *testing.T) {
+		var calls atomic.Int32
+		_, err := RetryOnTransientFailure(context.Background(), "test", 1, func() (string, error) {
+			calls.Add(1)
+			return "", tsclient.ErrResponse{Status: 503, Message: "unavailable"}
+		})
+		if err == nil {
+			t.Fatal("expected error")
+		}
+		if got := calls.Load(); got != 1 {
+			t.Errorf("got %d calls, want 1", got)
+		}
+	})
+
+	t.Run("success-case-one-attempt", func(t *testing.T) {
+		result, err := RetryOnTransientFailure(context.Background(), "test", 1, func() (string, error) {
+			return "key-123", nil
+		})
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if result != "key-123" {
+			t.Errorf("got %q, want %q", result, "key-123")
+		}
+	})
+
+	t.Run("retries-on-retriable-error", func(t *testing.T) {
+		var calls atomic.Int32
+		result, err := RetryOnTransientFailure(context.Background(), "test", 3, func() (string, error) {
+			n := calls.Add(1)
+			if n < 3 {
+				return "", tsclient.ErrResponse{Status: 503, Message: "unavailable"}
+			}
+			return "key-123", nil
+		})
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if result != "key-123" {
+			t.Errorf("got %q, want %q", result, "key-123")
+		}
+		if got := calls.Load(); got != 3 {
+			t.Errorf("got %d calls, want 3", got)
+		}
+	})
+
+	t.Run("calls-once-on-non-retriable-error", func(t *testing.T) {
+		var calls atomic.Int32
+		_, err := RetryOnTransientFailure(context.Background(), "test", 3, func() (string, error) {
+			calls.Add(1)
+			return "", tsclient.ErrResponse{Status: 400, Message: "bad request"}
+		})
+		if err == nil {
+			t.Fatal("expected error")
+		}
+		if got := calls.Load(); got != 1 {
+			t.Errorf("got %d calls, want 1", got)
+		}
+	})
+
+	t.Run("exhausts-retries", func(t *testing.T) {
+		var calls atomic.Int32
+		_, err := RetryOnTransientFailure(context.Background(), "test", 3, func() (string, error) {
+			calls.Add(1)
+			return "", tsclient.ErrResponse{Status: 503, Message: "unavailable"}
+		})
+		if err == nil {
+			t.Fatal("expected error")
+		}
+		if got := calls.Load(); got != 3 {
+			t.Errorf("got %d calls, want %d", got, 3)
+		}
+	})
+
+	t.Run("respects-context-cancellation", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		var calls atomic.Int32
+		_, err := RetryOnTransientFailure(ctx, "test", 3, func() (string, error) {
+			n := calls.Add(1)
+			if n == 1 {
+				cancel()
+			}
+			return "", tsclient.ErrResponse{Status: 503, Message: "unavailable"}
+		})
+		if err == nil {
+			t.Fatal("expected error")
+		}
+		if got := calls.Load(); got > 2 {
+			t.Errorf("got %d calls, want at most 2", got)
+		}
+	})
+
+	t.Run("success-on-first-try", func(t *testing.T) {
+		var calls atomic.Int32
+		result, err := RetryOnTransientFailure(context.Background(), "test", 3, func() (string, error) {
+			calls.Add(1)
+			return "key-abc", nil
+		})
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if result != "key-abc" {
+			t.Errorf("got %q, want %q", result, "key-abc")
+		}
+		if got := calls.Load(); got != 1 {
+			t.Errorf("got %d calls, want 1", got)
+		}
+	})
+
+	t.Run("retries-deadline-exceeded", func(t *testing.T) {
+		var calls atomic.Int32
+		result, err := RetryOnTransientFailure(context.Background(), "test", 3, func() (string, error) {
+			n := calls.Add(1)
+			if n < 2 {
+				return "", context.DeadlineExceeded
+			}
+			return "key-456", nil
+		})
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if result != "key-456" {
+			t.Errorf("got %q, want %q", result, "key-456")
+		}
+		if got := calls.Load(); got != 2 {
+			t.Errorf("got %d calls, want 2", got)
+		}
+	})
+
+	t.Run("retry-after-respects-cancellation", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		var calls atomic.Int32
+		start := time.Now()
+		go func() {
+			time.Sleep(200 * time.Millisecond)
+			cancel()
+		}()
+		_, err := RetryOnTransientFailure(ctx, "test", 3, func() (string, error) {
+			calls.Add(1)
+			return "", &oauth2.RetrieveError{
+				Response: &http.Response{
+					StatusCode: http.StatusTooManyRequests,
+					Header:     http.Header{"Retry-After": []string{"10"}},
+				},
+				Body: []byte("rate limited"),
+			}
+		})
+		if err == nil {
+			t.Fatal("expected error")
+		}
+		elapsed := time.Since(start)
+		if elapsed > 2*time.Second {
+			t.Errorf("expected early exit on cancellation, took %v", elapsed)
+		}
+	})
+}