feat: implement and refactor network request retries

packages/agentproxy/ca.go

@@ -16,7 +16,7 @@ import (
 	"time"
 
 	"github.com/Infisical/infisical-merge/packages/api"
-	"github.com/go-resty/resty/v2"
+	"github.com/Infisical/infisical-merge/packages/util"
 	"github.com/rs/zerolog/log"
 )
 
@@ -166,7 +166,11 @@ func (c *caManager) resignIntermediateLocked() error {
 	}
 	pubPem := pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubDer})
 
-	client := resty.New().SetAuthToken(c.token())
+	client, err := util.GetRestyClientWithCustomHeaders()
+	if err != nil {
+		return err
+	}
+	client.SetAuthToken(c.token())
 	resp, err := api.CallSignAgentProxyIntermediateCa(client, api.SignAgentProxyIntermediateCaRequest{
 		PublicKey: string(pubPem),
 	})

packages/agentproxy/cache.go

@@ -11,7 +11,6 @@ import (
 
 	"github.com/Infisical/infisical-merge/packages/api"
 	"github.com/Infisical/infisical-merge/packages/util"
-	"github.com/go-resty/resty/v2"
 	"github.com/rs/zerolog/log"
 )
 
@@ -231,7 +230,11 @@ type resolveParams struct {
 // resolveServices lists the proxied services for a scope and attaches credential values. Shared by
 // both resolvers; the differences live in resolveParams.
 func resolveServices(scope agentScope, p resolveParams) ([]*resolvedService, error) {
-	client := resty.New().SetAuthToken(p.discoveryToken)
+	client, err := util.GetRestyClientWithCustomHeaders()
+	if err != nil {
+		return nil, err
+	}
+	client.SetAuthToken(p.discoveryToken)
 	listResp, err := api.CallListProxiedServices(client, api.ListProxiedServicesRequest{
 		ProjectID:   scope.projectID,
 		Environment: scope.environment,

packages/agentproxy/leases.go

@@ -11,7 +11,7 @@ import (
 	"time"
 
 	"github.com/Infisical/infisical-merge/packages/api"
-	"github.com/go-resty/resty/v2"
+	"github.com/Infisical/infisical-merge/packages/util"
 	"github.com/rs/zerolog/log"
 	"golang.org/x/sync/singleflight"
 )
@@ -99,7 +99,11 @@ func newLeaseStore(proxyToken func() string) *leaseStore {
 
 func defaultLeaseMinter(proxyToken func() string) leaseMinter {
 	return func(args leaseMintArgs) (leaseMintResult, error) {
-		client := resty.New().SetAuthToken(proxyToken())
+		client, err := util.GetRestyClientWithCustomHeaders()
+		if err != nil {
+			return leaseMintResult{}, err
+		}
+		client.SetAuthToken(proxyToken())
 		resp, err := api.CallCreateDynamicSecretLeaseV1(client, api.CreateDynamicSecretLeaseV1Request{
 			ProjectSlug:       args.projectSlug,
 			Environment:       args.environment,
@@ -116,8 +120,12 @@ func defaultLeaseMinter(proxyToken func() string) leaseMinter {
 
 func defaultLeaseRevoker(proxyToken func() string) leaseRevoker {
 	return func(leaseID, projectSlug, environment, path string) error {
-		client := resty.New().SetAuthToken(proxyToken())
-		_, err := api.CallRevokeDynamicSecretLeaseV1(client, api.RevokeDynamicSecretLeaseV1Request{
+		client, err := util.GetRestyClientWithCustomHeaders()
+		if err != nil {
+			return err
+		}
+		client.SetAuthToken(proxyToken())
+		_, err = api.CallRevokeDynamicSecretLeaseV1(client, api.RevokeDynamicSecretLeaseV1Request{
 			LeaseID:     leaseID,
 			ProjectSlug: projectSlug,
 			Environment: environment,

packages/agentproxy/proxy.go

@@ -19,7 +19,7 @@ import (
 	"time"
 
 	"github.com/Infisical/infisical-merge/packages/api"
-	"github.com/go-resty/resty/v2"
+	"github.com/Infisical/infisical-merge/packages/util"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 )
@@ -175,7 +175,12 @@ func (ps *proxyServer) flushUsage() {
 	ps.usage = make(map[string]struct{})
 	ps.usageMu.Unlock()
 
-	client := resty.New().SetAuthToken(ps.opts.ProxyToken()).SetTimeout(usageReportTimeout)
+	client, err := util.GetRestyClientWithPolicy(util.BestEffortRetryPolicy())
+	if err != nil {
+		log.Debug().Err(err).Msg("failed to build usage-reporting client; dropping batch")
+		return
+	}
+	client.SetAuthToken(ps.opts.ProxyToken()).SetTimeout(usageReportTimeout)
 	for serviceID := range snapshot {
 		if err := api.CallReportProxiedServiceUsage(client, serviceID); err != nil {
 			// Warn once: the usual cause is a missing Report Usage permission, which fails every attempt.

packages/cmd/agent.go

@@ -1684,15 +1684,14 @@ func (tm *AgentManager) RevokeCredentials() error {
 
 // Refreshes the existing access token
 func (tm *AgentManager) RefreshAccessToken(accessToken string) error {
-	httpClient, err := util.GetRestyClientWithCustomHeaders()
+	policy := util.AgentRetryPolicy()
+	policy.ReplaySafe = true // renewal extends the presented token, so a replay cannot double-apply
+
+	httpClient, err := util.GetRestyClientWithPolicy(policy)
 	if err != nil {
 		return err
 	}
 
-	httpClient.SetRetryCount(10000).
-		SetRetryMaxWaitTime(20 * time.Second).
-		SetRetryWaitTime(5 * time.Second)
-
 	response, err := api.CallMachineIdentityRefreshAccessToken(httpClient, api.UniversalAuthRefreshRequest{AccessToken: accessToken})
 	if err != nil {
 		return err

packages/cmd/agent_proxy.go

@@ -172,7 +172,11 @@ func runAgentProxyConnect(cmd *cobra.Command, args []string) {
 		Set("credentialSource", tokenSource).
 		Set("allowReadableBrokeredSecrets", allowReadableBrokered))
 
-	httpClient := resty.New().SetAuthToken(token.Token)
+	httpClient, err := util.GetRestyClientWithCustomHeaders()
+	if err != nil {
+		util.HandleError(err, "Failed to build the API client")
+	}
+	httpClient.SetAuthToken(token.Token)
 
 	caResp, err := api.CallGetAgentProxyCa(httpClient)
 	if err != nil {

packages/cmd/agent_proxy_run.go

@@ -85,7 +85,11 @@ func runAgentProxyRun(cmd *cobra.Command, args []string) {
 	// The single identity for the run: fetches config and secret values in the parent. The child gets none of it.
 	src := resolveDeveloperTokenSource(cmd)
 
-	httpClient := resty.New().SetAuthToken(src.token())
+	httpClient, err := util.GetRestyClientWithCustomHeaders()
+	if err != nil {
+		util.HandleError(err, "Failed to build the API client")
+	}
+	httpClient.SetAuthToken(src.token())
 	placeholders := fetchLocalProxiedServiceConfig(httpClient, projectID, environment, secretPath)
 
 	local := &agentproxy.LocalOptions{

packages/cmd/login.go

@@ -638,7 +638,6 @@ func getFreshUserCredentials(email string, password string) (*api.GetLoginV3Resp
 	if err != nil {
 		return nil, err
 	}
-	httpClient.SetRetryCount(5)
 
 	loginV3Response, err := api.CallLoginV3(httpClient, api.GetLoginV3Request{
 		Email:    email,
@@ -658,7 +657,6 @@ func getFreshUserCredentialsWithSrp(email string, password string) (*api.GetLogi
 	if err != nil {
 		return nil, nil, err
 	}
-	httpClient.SetRetryCount(5)
 
 	params := srp.GetParams(4096)
 	secret1 := srp.GenKey()

packages/pam/agent/run.go

@@ -54,7 +54,10 @@ type Options struct {
 
 // Run binds a proxy per account and launches the agent. It returns the child's exit code.
 func Run(opts Options) (int, error) {
-	httpClient := resty.New()
+	httpClient, err := util.GetRestyClientWithCustomHeaders()
+	if err != nil {
+		return 1, fmt.Errorf("failed to build the API client: %w", err)
+	}
 	httpClient.SetHeader("User-Agent", api.USER_AGENT)
 
 	// Read the token per request rather than fixing it once. Sessions are created lazily and ended at

packages/pam/local/access.go

@@ -68,7 +68,11 @@ func StartPAMAccess(accessToken, path, reason, durationStr, targetHost string, p
 	log.Info().Msgf("Starting PAM access for: %s", strings.TrimPrefix(displayPath, "/"))
 	log.Info().Msgf("Session duration: %s", durationStr)
 
-	httpClient := resty.New()
+	httpClient, err := util.GetRestyClientWithCustomHeaders()
+	if err != nil {
+		util.HandleError(err, "Failed to build the API client")
+		return
+	}
 	httpClient.SetAuthToken(accessToken)
 	httpClient.SetHeader("User-Agent", api.USER_AGENT)
 

packages/util/common.go

@@ -45,7 +45,14 @@ func ValidateInfisicalAPIConnection() (ok bool) {
 	return err == nil
 }
 
+// GetRestyClientWithCustomHeaders is the single place API clients are built, which is what applies
+// the retry policy everywhere. Do not construct resty clients directly; TestNoDirectRestyConstruction
+// enforces this.
 func GetRestyClientWithCustomHeaders() (*resty.Client, error) {
+	return GetRestyClientWithPolicy(DefaultRetryPolicy())
+}
+
+func GetRestyClientWithPolicy(policy RetryPolicy) (*resty.Client, error) {
 	httpClient := resty.New()
 	customHeaders := os.Getenv("INFISICAL_CUSTOM_HEADERS")
 	if customHeaders != "" {
@@ -56,7 +63,7 @@ func GetRestyClientWithCustomHeaders() (*resty.Client, error) {
 
 		httpClient.SetHeaders(headers)
 	}
-	return httpClient, nil
+	return applyRetryPolicy(httpClient, policy), nil
 }
 
 func GetInfisicalCustomHeadersMap() (map[string]string, error) {

packages/util/helper.go

@@ -334,10 +334,6 @@ func UniversalAuthLogin(clientId string, clientSecret string) (api.UniversalAuth
 		return api.UniversalAuthLoginResponse{}, err
 	}
 
-	httpClient.SetRetryCount(10000).
-		SetRetryMaxWaitTime(20 * time.Second).
-		SetRetryWaitTime(5 * time.Second)
-
 	tokenResponse, err := api.CallUniversalAuthLogin(httpClient, api.UniversalAuthLoginRequest{ClientId: clientId, ClientSecret: clientSecret})
 	if err != nil {
 		return api.UniversalAuthLoginResponse{}, err
@@ -347,16 +343,14 @@ func UniversalAuthLogin(clientId string, clientSecret string) (api.UniversalAuth
 }
 
 func RenewMachineIdentityAccessToken(accessToken string) (string, error) {
+	policy := DefaultRetryPolicy()
+	policy.ReplaySafe = true // renewal extends the presented token, so a replay cannot double-apply
 
-	httpClient, err := GetRestyClientWithCustomHeaders()
+	httpClient, err := GetRestyClientWithPolicy(policy)
 	if err != nil {
 		return "", err
 	}
 
-	httpClient.SetRetryCount(10000).
-		SetRetryMaxWaitTime(20 * time.Second).
-		SetRetryWaitTime(5 * time.Second)
-
 	request := api.UniversalAuthRefreshRequest{
 		AccessToken: accessToken,
 	}