Adjust http client to support rate limits

cmd/api/api.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"net/http"
 	"os"
 	"strings"
 	"time"
@@ -114,17 +115,19 @@ func (c *ApiCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
 
 	fullEndpoint := buildFullEndpoint(c.Endpoint, f.Config.OrganizationSlug(), c.Analytics)
 
-	// Create an HTTP client with appropriate configuration
+	// Create an HTTP client with rate-limit retry via the shared transport.
+	rl := httpClient.NewRateLimitTransport(nil)
+	rl.MaxRetryDelay = 60 * time.Second
+	rl.OnRateLimit = func(attempt int, delay time.Duration) {
+		if c.Verbose {
+			fmt.Fprintf(os.Stderr, "WARNING: Rate limit exceeded, retrying in %v @ %q (attempt %d)\n", delay, time.Now().Add(delay).Format(time.RFC3339), attempt)
+		}
+	}
+
 	client := httpClient.NewClient(
 		f.Config.APIToken(),
 		httpClient.WithBaseURL(f.RestAPIClient.BaseURL.String()),
-		httpClient.WithMaxRetries(3),
-		httpClient.WithMaxRetryDelay(60*time.Second),
-		httpClient.WithOnRetry(func(attempt int, delay time.Duration) {
-			if c.Verbose {
-				fmt.Fprintf(os.Stderr, "WARNING: Rate limit exceeded, retrying in %v @ %q (attempt %d)\n", delay, time.Now().Add(delay).Format(time.RFC3339), attempt)
-			}
-		}),
+		httpClient.WithHTTPClient(&http.Client{Transport: rl}),
 	)
 
 	// Process custom headers

cmd/preflight/preflight.go

@@ -18,6 +18,7 @@ import (
 	"github.com/buildkite/cli/v3/internal/build/watch"
 	"github.com/buildkite/cli/v3/internal/cli"
 	bkErrors "github.com/buildkite/cli/v3/internal/errors"
+	bkhttp "github.com/buildkite/cli/v3/internal/http"
 	"github.com/buildkite/cli/v3/internal/pipeline/resolver"
 	"github.com/buildkite/cli/v3/internal/preflight"
 	"github.com/buildkite/cli/v3/pkg/cmd/factory"
@@ -43,7 +44,8 @@ func (c *PreflightCmd) Help() string {
 }
 
 func (c *PreflightCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
-	f, err := newFactory(factory.WithDebug(globals.EnableDebug()))
+	rlTransport := bkhttp.NewRateLimitTransport(http.DefaultTransport)
+	f, err := newFactory(factory.WithDebug(globals.EnableDebug()), factory.WithTransport(rlTransport))
 	if err != nil {
 		return bkErrors.NewInternalError(err, "failed to initialize CLI", "This is likely a bug", "Report to Buildkite")
 	}
@@ -90,6 +92,17 @@ func (c *PreflightCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error
 
 	renderer := newRenderer(os.Stdout, c.JSON, c.Text, stop)
 
+	rlTransport.OnRateLimit = func(attempt int, delay time.Duration) {
+		if globals.EnableDebug() {
+			_ = renderer.Render(Event{
+				Type:        EventOperation,
+				Time:        time.Now(),
+				PreflightID: preflightID.String(),
+				Title:       fmt.Sprintf("Rate limited by API, waiting %s before retrying (attempt %d/%d)...", delay.Truncate(time.Second), attempt+1, rlTransport.MaxRetries),
+			})
+		}
+	}
+
 	_ = renderer.Render(Event{Type: EventOperation, Time: time.Now(), PreflightID: preflightID.String(), Title: "Pushing snapshot of working tree..."})
 
 	var opts []preflight.SnapshotOption

go.mod

@@ -6,7 +6,6 @@ require (
 	github.com/alecthomas/kong v1.15.0
 	github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be
 	github.com/buildkite/go-buildkite/v4 v4.19.0
-	github.com/buildkite/roko v1.4.0
 	github.com/buildkite/termoji v0.0.0-20260330080310-c0aa4ebee0d1
 	github.com/charmbracelet/bubbles v1.0.0
 	github.com/charmbracelet/bubbletea v1.3.10

go.sum

@@ -35,8 +35,6 @@ github.com/bradleyjkemp/cupaloy/v2 v2.6.0 h1:knToPYa2xtfg42U3I6punFEjaGFKWQRXJwj
 github.com/bradleyjkemp/cupaloy/v2 v2.6.0/go.mod h1:bm7JXdkRd4BHJk9HpwqAI8BoAY1lps46Enkdqw6aRX0=
 github.com/buildkite/go-buildkite/v4 v4.19.0 h1:HPc6+V/Ky6v8eDWHxyvzHtxuhruCLUyuRrChOKvJE1I=
 github.com/buildkite/go-buildkite/v4 v4.19.0/go.mod h1:8+7GiWBKwEPAWoZnRU/kpNCt46j1iVH8kFMMbD4YDfc=
-github.com/buildkite/roko v1.4.0 h1:DxixoCdpNqxu4/1lXrXbfsKbJSd7r1qoxtef/TT2J80=
-github.com/buildkite/roko v1.4.0/go.mod h1:0vbODqUFEcVf4v2xVXRfZZRsqJVsCCHTG/TBRByGK4E=
 github.com/buildkite/termoji v0.0.0-20260330080310-c0aa4ebee0d1 h1:aaEl0QZURcwC+KOfFTzSp66xknw5eTmFZ1NgB87s2xk=
 github.com/buildkite/termoji v0.0.0-20260330080310-c0aa4ebee0d1/go.mod h1:ZTEvQlMN3+qzjROvjRb1p0X+xDQxxKpkMFhMSnaTrpw=
 github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
@@ -225,5 +223,3 @@ gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
 gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-gotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q=
-gotest.tools/v3 v3.5.2/go.mod h1:LtdLGcnqToBH83WByAAi/wiwSFCArdFIUV/xxN4pcjA=

internal/http/client.go

@@ -8,11 +8,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
-	"strconv"
 	"strings"
-	"time"
-
-	"github.com/buildkite/roko"
 )
 
 // ErrorResponse represents an error response from the API
@@ -44,10 +40,6 @@ type Client struct {
 	token     string
 	userAgent string
 	client    *http.Client
-
-	maxRetries    int
-	maxRetryDelay time.Duration
-	onRetry       OnRetryFunc
 }
 
 // ClientOption is a function that modifies a Client
@@ -74,30 +66,6 @@ func WithHTTPClient(client *http.Client) ClientOption {
 	}
 }
 
-// WithMaxRetries sets the maximum number of retries for rate-limited requests.
-func WithMaxRetries(n int) ClientOption {
-	return func(c *Client) {
-		c.maxRetries = n
-	}
-}
-
-// WithMaxRetryDelay sets the maximum delay between retries
-func WithMaxRetryDelay(d time.Duration) ClientOption {
-	return func(c *Client) {
-		c.maxRetryDelay = d
-	}
-}
-
-// WithOnRetry sets a callback that is invoked before each retry sleep.
-func WithOnRetry(f OnRetryFunc) ClientOption {
-	return func(c *Client) {
-		c.onRetry = f
-	}
-}
-
-// OnRetryFunc is called before each retry sleep with the attempt number and delay duration.
-type OnRetryFunc func(attempt int, delay time.Duration)
-
 // NewClient creates a new HTTP client with the given token and options
 func NewClient(token string, opts ...ClientOption) *Client {
 	c := &Client{
@@ -164,23 +132,6 @@ func (e *ErrorResponse) IsTooManyRequests() bool {
 	return e.StatusCode == http.StatusTooManyRequests
 }
 
-// RetryAfter returns the duration to wait before retrying, based on the RateLimit-Reset header.
-// Returns 0 if the header is missing or invalid.
-func (e *ErrorResponse) RetryAfter() time.Duration {
-	if e.Headers == nil {
-		return 0
-	}
-	resetStr := e.Headers.Get("RateLimit-Reset")
-	if resetStr == "" {
-		return 0
-	}
-	seconds, err := strconv.Atoi(resetStr)
-	if err != nil || seconds < 0 {
-		return 0
-	}
-	return time.Duration(seconds) * time.Second
-}
-
 // Do performs an HTTP request with the given method, endpoint, and body.
 func (c *Client) Do(ctx context.Context, method, endpoint string, body interface{}, v interface{}) error {
 	// Ensure endpoint starts with "/"
@@ -215,21 +166,7 @@ func (c *Client) Do(ctx context.Context, method, endpoint string, body interface
 		}
 	}
 
-	r := roko.NewRetrier(
-		roko.WithMaxAttempts(c.maxRetries+1),
-		roko.WithStrategy(roko.Constant(0)),
-	)
-
-	respBody, err := roko.DoFunc(ctx, r, func(r *roko.Retrier) ([]byte, error) {
-		resp, err := c.send(ctx, method, reqURL, bodyBytes)
-		if err != nil {
-			if !c.handleRetry(r, err) {
-				r.Break()
-			}
-			return nil, err
-		}
-		return resp, nil
-	})
+	respBody, err := c.send(ctx, method, reqURL, bodyBytes)
 	if err != nil {
 		return err
 	}
@@ -284,30 +221,3 @@ func (c *Client) send(ctx context.Context, method, reqURL string, body []byte) (
 
 	return respBody, nil
 }
-
-// handleRetry checks if an error is retryable and configures the retrier accordingly.
-// Returns true if the request should be retried, false otherwise.
-func (c *Client) handleRetry(r *roko.Retrier, err error) bool {
-	errResp, ok := err.(*ErrorResponse)
-	if !ok || !errResp.IsTooManyRequests() {
-		return false
-	}
-
-	attempt := r.AttemptCount()
-	delay := errResp.RetryAfter()
-	if attempt > 0 {
-		// Got rate-limited again means contention - back off exponentially
-		delay *= time.Duration(1 << attempt)
-	}
-
-	if c.maxRetryDelay > 0 {
-		delay = min(delay, c.maxRetryDelay)
-	}
-
-	if c.onRetry != nil {
-		c.onRetry(attempt, delay)
-	}
-
-	r.SetNextInterval(delay)
-	return true
-}