Add timeout flag to CLI commands for improved resilience (#194)

This pull request implements robust, configurable network retry logic
for the Envault CLI, improving resilience to transient API/network
failures. It introduces a unified `--timeout` flag (and
`ENVAULT_RETRY_MAX_DURATION` env var) for `pull` and `run` commands,
which controls the maximum duration for retrying failed API requests.
The retry logic uses exponential backoff with jitter and is thoroughly
tested for various scenarios.

**Network Resilience & Retry Logic:**

* Adds exponential backoff with jitter for retrying transient network
and server (5xx) errors in the API client, up to a configurable maximum
duration (`ENVAULT_RETRY_MAX_DURATION`, default 5 minutes).
* Refactors API client to respect context cancellation and to avoid
retrying on 4xx errors (except for 401, which triggers a token refresh).
**CLI Command Improvements:**

* Adds a `--timeout` flag to both `pull` and `run` commands, allowing
users to override the maximum retry duration from the command line.
* Both commands now use a unified timeout resolver, defaulting to 5
minutes if not set by flag or env var.

**Testing & Documentation:**

* Adds comprehensive tests for retry logic, including network errors,
5xx/4xx handling, token refresh, and timeout configuration.
* Updates CLI documentation for `pull` and `run` commands, and adds a
new section on network resilience and timeouts.

---
Co-authored-by: Rajat Patra <113469515+RawJat@users.noreply.github.com>

Author: DinanathDash

.gitignore

@@ -54,6 +54,7 @@ cli-go/envault
 cli-go/envault.json
 cli-go/out.txt
 cli-go/.env
+cli-go/.gocache
 
 # Generated source files (fumadocs-mdx)
 .source/

cli-go/cmd/network_timeout.go

@@ -0,0 +1,27 @@
+package cmd
+
+import (
+	"os"
+	"strings"
+	"time"
+)
+
+const defaultCLINetworkTimeout = 5 * time.Minute
+
+func resolveCLINetworkTimeout(flagTimeout time.Duration) time.Duration {
+	if flagTimeout > 0 {
+		return flagTimeout
+	}
+
+	raw := strings.TrimSpace(os.Getenv("ENVAULT_RETRY_MAX_DURATION"))
+	if raw == "" {
+		return defaultCLINetworkTimeout
+	}
+
+	d, err := time.ParseDuration(raw)
+	if err != nil || d <= 0 {
+		return defaultCLINetworkTimeout
+	}
+
+	return d
+}

cli-go/cmd/pull.go

@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"strings"
 	"syscall"
+	"time"
 
 	"github.com/AlecAivazis/survey/v2"
 	"github.com/DinanathDash/Envault/cli-go/internal/api"
@@ -40,14 +41,15 @@ type ProjectResponse struct {
 var forcePull bool
 var projectFlag string
 var fileFlag string
+var pullTimeoutFlag time.Duration
 
 var pullCmd = &cobra.Command{
 	Use:   "pull",
 	Short: "Fetch secrets and write to .env",
 	Run: func(cmd *cobra.Command, args []string) {
 		// Graceful cancellation: cancel the context on Ctrl+C / SIGTERM so that
 		// in-flight HTTP requests are aborted cleanly.
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithTimeout(context.Background(), resolveCLINetworkTimeout(pullTimeoutFlag))
 		defer cancel()
 		sigCh := make(chan os.Signal, 1)
 		signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
@@ -335,4 +337,5 @@ func init() {
 	pullCmd.Flags().BoolVarP(&forcePull, "force", "f", false, "Overwrite .env without confirmation")
 	pullCmd.Flags().StringVarP(&projectFlag, "project", "p", "", "Project ID")
 	pullCmd.Flags().StringVar(&fileFlag, "file", "", "Local .env file path override")
+	pullCmd.Flags().DurationVar(&pullTimeoutFlag, "timeout", 0, "Maximum duration for API retry/wait behavior (e.g. 30s, 2m). Overrides ENVAULT_RETRY_MAX_DURATION.")
 }

cli-go/cmd/run.go

@@ -1,13 +1,13 @@
 package cmd
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/url"
 	"os"
 	"os/exec"
 	"runtime"
-	"strconv"
 	"strings"
 	"time"
 
@@ -18,6 +18,8 @@ import (
 	"github.com/spf13/cobra"
 )
 
+var runTimeoutFlag time.Duration
+
 var runCmd = &cobra.Command{
 	Use:   "run -- <command>",
 	Short: "Run a command with secrets injected from Envault",
@@ -50,12 +52,14 @@ var runCmd = &cobra.Command{
 		}
 		client := api.NewClient()
 		path := fmt.Sprintf("/projects/%s/secrets?environment=%s", projectID, url.QueryEscape(targetEnv))
+		ctx, cancel := context.WithTimeout(context.Background(), resolveCLINetworkTimeout(runTimeoutFlag))
+		defer cancel()
 
 		var secretsResp SecretsResponse
 		var envSecrets []offlinecache.Secret
 		loader := ui.NewLoader(ui.LoaderThemeFetch, fmt.Sprintf("VaultPulse preparing runtime secrets (%s)...", targetEnv))
 		loader.Start()
-		respBytes, err := client.GetWithTimeout(path, resolveRunTimeout(client.BaseURL))
+		respBytes, err := client.GetWithContext(ctx, path)
 		loader.Stop()
 
 		usedOfflineCache := false
@@ -165,25 +169,6 @@ func humanizeDuration(d time.Duration) string {
 	return d.Round(time.Minute).String()
 }
 
-func resolveRunTimeout(baseURL string) time.Duration {
-	defaultSeconds := 10
-	if isLocalBaseURL(baseURL) {
-		defaultSeconds = 20
-	}
-
-	raw := strings.TrimSpace(os.Getenv("ENVAULT_RUN_TIMEOUT_SECONDS"))
-	if raw == "" {
-		return time.Duration(defaultSeconds) * time.Second
-	}
-
-	seconds, err := strconv.Atoi(raw)
-	if err != nil || seconds <= 0 {
-		return time.Duration(defaultSeconds) * time.Second
-	}
-
-	return time.Duration(seconds) * time.Second
-}
-
 func isLocalBaseURL(baseURL string) bool {
 	u, err := url.Parse(baseURL)
 	if err != nil {
@@ -221,4 +206,5 @@ func isLikelyDevCommand(runTarget string, runArgs []string) bool {
 func init() {
 	rootCmd.AddCommand(runCmd)
 	runCmd.Flags().StringVarP(&projectFlag, "project", "p", "", "Project ID")
+	runCmd.Flags().DurationVar(&runTimeoutFlag, "timeout", 0, "Maximum duration for API retry/wait behavior (e.g. 30s, 2m). Overrides ENVAULT_RETRY_MAX_DURATION.")
 }

cli-go/internal/api/client.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"math/rand"
 	"net"
 	"net/http"
 	"net/url"
@@ -34,6 +35,15 @@ type Client struct {
 	HTTP    *http.Client
 }
 
+var (
+	defaultRetryMaxDuration = 5 * time.Minute
+	retryBaseBackoff        = 2 * time.Second
+	retryMultiplier         = 2.0
+	sleepWithContextFn      = sleepWithContext
+	refreshTokenFn          = func(c *Client, httpClient *http.Client) error { return c.refreshToken(httpClient) }
+	nowFn                   = time.Now
+)
+
 func NewClient() *Client {
 	baseURL := os.Getenv("ENVAULT_CLI_URL")
 	if baseURL == "" {
@@ -189,62 +199,187 @@ func (c *Client) doReqCtx(ctx context.Context, method, path string, body interfa
 		httpClient = &http.Client{}
 	}
 
-	var bodyReader io.Reader
-
+	var reqBody []byte
+	var err error
 	if body != nil {
-		reqBody, err := json.Marshal(body)
+		reqBody, err = json.Marshal(body)
 		if err != nil {
 			return nil, err
 		}
-		bodyReader = bytes.NewBuffer(reqBody)
 	}
 
-	req, err := http.NewRequestWithContext(ctx, method, c.BaseURL+path, bodyReader)
-	if err != nil {
-		return nil, err
+	start := nowFn()
+	attempt := 1
+	tokenRefreshTried := false
+	maxDuration := resolveRetryMaxDuration()
+
+	for {
+		var bodyReader io.Reader
+		if reqBody != nil {
+			bodyReader = bytes.NewReader(reqBody)
+		}
+
+		req, err := http.NewRequestWithContext(ctx, method, c.BaseURL+path, bodyReader)
+		if err != nil {
+			return nil, err
+		}
+
+		req.Header.Set("Content-Type", "application/json")
+		if c.Token != "" {
+			req.Header.Set("Authorization", "Bearer "+c.Token)
+		}
+		if actorSource := strings.TrimSpace(os.Getenv("ENVAULT_CLI_ACTOR_SOURCE")); actorSource != "" {
+			req.Header.Set("X-Envault-Actor-Source", actorSource)
+		}
+
+		resp, reqErr := httpClient.Do(req)
+		if reqErr != nil {
+			if !canRetry || !isRetryableRequestError(reqErr) {
+				return nil, reqErr
+			}
+			wait, ok := computeRetryWait(start, attempt, maxDuration)
+			if !ok {
+				return nil, reqErr
+			}
+			logRetryWarning(method, path, attempt, reqErr, wait)
+			if err := sleepWithContextFn(ctx, wait); err != nil {
+				return nil, err
+			}
+			attempt++
+			continue
+		}
+
+		bodyBytes, readErr := io.ReadAll(resp.Body)
+		_ = resp.Body.Close()
+		if readErr != nil {
+			return nil, readErr
+		}
+
+		if resp.StatusCode == 401 && canRetry && !tokenRefreshTried {
+			if c.Token != "" && !strings.HasPrefix(c.Token, "envault_svc_") {
+				errRefresh := refreshTokenFn(c, httpClient)
+				if errRefresh == nil {
+					tokenRefreshTried = true
+					continue
+				}
+				return nil, fmt.Errorf("Refresh Token Exchange Failed: %v | (Original Auth Error: %s)", errRefresh, string(bodyBytes))
+			}
+		}
+
+		if isRetryableStatusCode(resp.StatusCode) && canRetry {
+			wait, ok := computeRetryWait(start, attempt, maxDuration)
+			if !ok {
+				return nil, &APIError{StatusCode: resp.StatusCode, Body: string(bodyBytes)}
+			}
+			logRetryWarning(method, path, attempt, fmt.Errorf("server returned %s", resp.Status), wait)
+			if err := sleepWithContextFn(ctx, wait); err != nil {
+				return nil, err
+			}
+			attempt++
+			continue
+		}
+
+		if resp.StatusCode >= 400 {
+			bodyStr := string(bodyBytes)
+			contentType := resp.Header.Get("Content-Type")
+			if strings.Contains(contentType, "text/html") {
+				bodyStr = "Server returned an HTML page (" + resp.Status + "). Ensure the API server is running."
+			}
+			return nil, &APIError{StatusCode: resp.StatusCode, Body: bodyStr}
+		}
+
+		contentType := resp.Header.Get("Content-Type")
+		if strings.Contains(contentType, "text/html") {
+			return nil, &APIError{StatusCode: resp.StatusCode, Body: "Server returned HTML instead of expected JSON API response. Ensure the API server is running."}
+		}
+
+		return bodyBytes, nil
 	}
+}
 
-	req.Header.Set("Content-Type", "application/json")
-	if c.Token != "" {
-		req.Header.Set("Authorization", "Bearer "+c.Token)
+func isRetryableStatusCode(statusCode int) bool {
+	return statusCode >= 500 && statusCode <= 599
+}
+
+func isRetryableRequestError(err error) bool {
+	if err == nil {
+		return false
 	}
-	if actorSource := strings.TrimSpace(os.Getenv("ENVAULT_CLI_ACTOR_SOURCE")); actorSource != "" {
-		req.Header.Set("X-Envault-Actor-Source", actorSource)
+	if errors.Is(err, context.Canceled) {
+		return false
 	}
+	return IsFallbackEligible(err)
+}
 
-	resp, err := httpClient.Do(req)
-	if err != nil {
-		return nil, err
+func computeRetryWait(start time.Time, attempt int, maxDuration time.Duration) (time.Duration, bool) {
+	if attempt < 1 {
+		attempt = 1
 	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode == 401 && canRetry {
-		if c.Token != "" && !strings.HasPrefix(c.Token, "envault_svc_") {
-			bodyBytes, _ := io.ReadAll(resp.Body)
-			errRefresh := c.refreshToken(httpClient)
-			if errRefresh == nil {
-				return c.doReqCtx(ctx, method, path, body, false, httpClient)
-			}
-			return nil, fmt.Errorf("Refresh Token Exchange Failed: %v | (Original Auth Error: %s)", errRefresh, string(bodyBytes))
+	exp := retryBaseBackoff
+	for i := 1; i < attempt; i++ {
+		next := time.Duration(float64(exp) * retryMultiplier)
+		if next <= exp {
+			break
 		}
+		exp = next
+	}
+	wait := jitterDuration(exp)
+	if wait <= 0 {
+		wait = retryBaseBackoff
 	}
 
-	if resp.StatusCode >= 400 {
-		bodyBytes, _ := io.ReadAll(resp.Body)
-		bodyStr := string(bodyBytes)
-		contentType := resp.Header.Get("Content-Type")
-		if strings.Contains(contentType, "text/html") {
-			bodyStr = "Server returned an HTML page (" + resp.Status + "). Ensure the API server is running."
+	elapsed := nowFn().Sub(start)
+	if elapsed >= maxDuration {
+		return 0, false
+	}
+	if elapsed+wait > maxDuration {
+		remaining := maxDuration - elapsed
+		if remaining > time.Second {
+			wait = remaining
+			return wait, true
 		}
-		return nil, &APIError{StatusCode: resp.StatusCode, Body: bodyStr}
+		return 0, false
+	}
+	return wait, true
+}
+
+func resolveRetryMaxDuration() time.Duration {
+	raw := strings.TrimSpace(os.Getenv("ENVAULT_RETRY_MAX_DURATION"))
+	if raw == "" {
+		return defaultRetryMaxDuration
 	}
 
-	contentType := resp.Header.Get("Content-Type")
-	if strings.Contains(contentType, "text/html") {
-		return nil, &APIError{StatusCode: resp.StatusCode, Body: "Server returned HTML instead of expected JSON API response. Ensure the API server is running."}
+	d, err := time.ParseDuration(raw)
+	if err != nil || d <= 0 {
+		return defaultRetryMaxDuration
 	}
 
-	return io.ReadAll(resp.Body)
+	return d
+}
+
+func jitterDuration(base time.Duration) time.Duration {
+	// Full-jitter style randomization in [base/2, base) to spread retries.
+	half := base / 2
+	if half <= 0 {
+		return base
+	}
+	return half + time.Duration(rand.Int63n(int64(half)))
+}
+
+func logRetryWarning(method, path string, attempt int, cause error, wait time.Duration) {
+	fmt.Fprintf(os.Stderr, "Warning: Envault API transient failure (%s %s, attempt %d). Retrying in %s. Cause: %v\n",
+		method, path, attempt, wait.Round(time.Millisecond), cause)
+}
+
+func sleepWithContext(ctx context.Context, d time.Duration) error {
+	timer := time.NewTimer(d)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return nil
+	}
 }
 
 func IsFallbackEligible(err error) bool {