Skip to content

fix: retry Atmos Pro uploads on transient 401/5xx with exponential backoff - #2255

Merged
Andriy Knysh (aknysh) merged 6 commits into
mainfrom
osterman/robust-upload-retry
Mar 27, 2026
Merged

Andriy Knysh (aknysh) merged 6 commits into
mainfrom
osterman/robust-upload-retry

Conversation

@osterman

@osterman Erik Osterman (Cloud Posse) (osterman) commented Mar 26, 2026 •

Copy link
Copy Markdown
Member

what

  • Add retry logic with exponential backoff (1s, 2s, 4s) to all three Atmos Pro upload methods: UploadInstanceStatus, UploadAffectedStacks, and UploadInstances
  • On 401 errors, re-exchange the OIDC token before retrying (handles JWT secret mismatches across deployment instances)
  • On 5xx or network errors, retry with backoff without re-auth
  • On 400/403/404, fail immediately without retrying
  • Upload failures no longer cause non-zero exit codes — the terraform plan/apply result drives the command exit code, not telemetry

why

  • Intermittent 401 errors on upload PATCH calls even though the OIDC exchange succeeded seconds earlier, likely due to JWT signed by a different deployment instance with a mismatched secret
  • Happens most often when multiple Atmos commands run in parallel (e.g., 14 concurrent stack operations in a matrix workflow)
  • Retrying the same job usually works, indicating transient failures
  • Upload telemetry should never block or fail the primary terraform workflow

references

Summary by CodeRabbit

  • New Features

    • Automatic upload retries with exponential backoff and automatic API token refresh on auth failures.
    • More structured API error classification to improve retry and auth decisioning.
  • Bug Fixes

    • Uploads and status-reporting failures now log warnings and no longer cause commands to fail.
  • Tests

    • Comprehensive tests for retry logic, token refresh behavior, and API error handling.

…ckoff

Add retry logic to all three upload methods (UploadInstanceStatus,
UploadAffectedStacks, UploadInstances) to handle intermittent 401 errors
caused by JWT secret mismatches across deployment instances, and transient
5xx server errors during concurrent GitHub Actions matrix workflows.

On 401, re-exchange the OIDC token before retrying. On 5xx or network
errors, retry with exponential backoff (1s, 2s, 4s). Non-retryable
errors (400, 403, 404) fail immediately.

Upload failures no longer cause non-zero exit codes — the terraform
plan/apply result drives the exit code, not telemetry.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@github-actions github-actions Bot added the size/m Medium size PR label Mar 26, 2026
@github-actions

github-actions Bot commented Mar 26, 2026 •

Copy link
Copy Markdown

Dependency Review

✅ No vulnerabilities or license issues found.

Snapshot Warnings

⚠️: No snapshots were found for the head SHA 6afbde5.
Ensure that dependencies are being submitted on PR branches and consider enabling retry-on-snapshot-warnings. See the documentation for more information and troubleshooting advice.

Scanned Files

None

@codecov

codecov Bot commented Mar 26, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 82.14286% with 25 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.28%. Comparing base (2be2909) to head (6afbde5).
⚠️ Report is 2 commits behind head on main.

Files with missing lines Patch % Lines
pkg/pro/api_client.go 71.42% 10 Missing ⚠️
internal/exec/describe_affected.go 0.00% 5 Missing ⚠️
pkg/pro/api_client_instance_status.go 60.00% 2 Missing and 2 partials ⚠️
pkg/pro/api_client_instances.go 60.00% 2 Missing and 2 partials ⚠️
pkg/list/list_instances.go 0.00% 2 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #2255      +/-   ##
==========================================
+ Coverage   77.21%   77.28%   +0.06%     
==========================================
  Files        1018     1020       +2     
  Lines       96368    96476     +108     
==========================================
+ Hits        74415    74558     +143     
+ Misses      17755    17721      -34     
+ Partials     4198     4197       -1     
Flag Coverage Δ
unittests 77.28% <82.14%> (+0.06%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
errors/errors.go 100.00% <ø> (ø)
internal/exec/terraform_execute_helpers_exec.go 82.45% <100.00%> (ø)
pkg/pro/api_error.go 100.00% <100.00%> (ø)
pkg/pro/retry.go 100.00% <100.00%> (ø)
pkg/list/list_instances.go 82.78% <0.00%> (-0.35%) ⬇️
pkg/pro/api_client_instance_status.go 67.74% <60.00%> (+2.22%) ⬆️
pkg/pro/api_client_instances.go 69.09% <60.00%> (+1.16%) ⬆️
internal/exec/describe_affected.go 56.90% <0.00%> (-0.73%) ⬇️
pkg/pro/api_client.go 91.31% <71.42%> (-1.89%) ⬇️

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@coderabbitai

coderabbitai Bot commented Mar 26, 2026 •

Copy link
Copy Markdown
Contributor
📝 Walkthrough

Walkthrough

Adds structured Pro API error type and retry helper with exponential backoff and OIDC token-refresh; integrates retries into Atmos Pro upload/instance-status calls, surfaces two new sentinel errors (ErrUploadRetryExhausted, ErrTokenRefreshFailed), and makes several upload failures non-fatal in CLI commands. (47 words)

Changes

Cohort / File(s) Summary
Error Sentinels
errors/errors.go
Added two exported sentinel errors: ErrUploadRetryExhausted, ErrTokenRefreshFailed.
Pro API error & tests
pkg/pro/api_error.go, pkg/pro/api_error_test.go
New pro.APIError type (status, operation, wrapped error) with Error(), Unwrap(), IsRetryable(), IsAuthError() and unit tests validating formatting, unwrap, errors.Is/As, and status-code behavior.
Retry core & tests
pkg/pro/retry.go, pkg/pro/retry_test.go
New doWithRetry implementing exponential backoff, retry classification, token-refresh hook for 401, default retry constants, and comprehensive tests covering retry, auth-refresh, non-retryable errors, and backoff timing.
Atmos Pro client changes
pkg/pro/api_client.go
AtmosProAPIClient gains atmosConfig and useOIDC fields, RefreshToken() method; upload flows refactored to use doWithRetry, and response handling now produces structured APIError for classification.
Per-request retry integration
pkg/pro/api_client_instances.go, pkg/pro/api_client_instance_status.go
POST/PATCH upload flows now recreate authenticated requests per retry attempt and run inside doWithRetry, delegating transient/401 handling to retry logic.
Callers made non-blocking
internal/exec/describe_affected.go, internal/exec/terraform_execute_helpers_exec.go, pkg/list/list_instances.go
Upload/status errors are logged as warnings and do not fail the outer commands when Atmos Pro interactions fail.

Sequence Diagram

sequenceDiagram
    participant Client
    participant Retry as doWithRetry
    participant APIClient as AtmosProAPIClient
    participant API as Atmos Pro API
    participant OIDC as OIDC Provider

    Client->>Retry: invoke upload operation
    Retry->>APIClient: execute upload fn
    APIClient->>API: POST/PATCH with JWT
    alt 2xx Success
        API-->>APIClient: 2xx response
        APIClient-->>Retry: nil (success)
        Retry-->>Client: nil
    else 401 Unauthorized
        API-->>APIClient: 401
        APIClient-->>Retry: APIError(401)
        Retry->>APIClient: RefreshToken()
        APIClient->>OIDC: exchange OIDC token
        OIDC-->>APIClient: new JWT
        APIClient->>APIClient: update APIToken
        Retry->>Retry: backoff sleep then retry
    else 5xx Server Error
        API-->>APIClient: 5xx
        APIClient-->>Retry: APIError(5xx)
        Retry->>Retry: backoff sleep then retry
    else Non-retryable 4xx
        API-->>APIClient: 4xx (non-401)
        APIClient-->>Retry: APIError(4xx)
        Retry-->>Client: return error immediately
    end
    alt retries exhausted
        Retry-->>Client: ErrUploadRetryExhausted + last error
    end
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Possibly related PRs

Suggested labels

minor

Suggested reviewers

  • aknysh
🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 32.14% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately and specifically describes the main change: implementing retry logic with exponential backoff for Atmos Pro uploads on transient 401/5xx errors.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch osterman/robust-upload-retry

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (1)
pkg/pro/api_client.go (1)

177-195: Make RefreshToken mockable.

RefreshToken() hard-wires the OIDC fetch and token-exchange helpers, so the 401 recovery path can only be tested through globals or real HTTP servers. That is already showing up in pkg/pro/retry_test.go, where the auth case never asserts a real refresh. A tiny injected interface or function pair here would keep this path unit-testable.

As per coding guidelines, "Use interfaces for external dependencies to facilitate mocking" and "Prefer unit tests with mocks over integration tests. Use interfaces plus dependency injection for testability."

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@pkg/pro/api_client.go` around lines 177 - 195, RefreshToken currently calls
getGitHubOIDCToken and exchangeOIDCTokenForAtmosToken directly, making it hard
to mock; add injectable function fields on AtmosProAPIClient (e.g.,
OIDCTokenFetcher func(githubOIDC string) (string, error) and TokenExchanger
func(baseURL, baseAPIEndpoint, oidcToken, workspaceID string) (string, error)),
defaulting those fields to getGitHubOIDCToken and exchangeOIDCTokenForAtmosToken
when nil, then have RefreshToken call
c.OIDCTokenFetcher(c.atmosConfig.Settings.Pro.GithubOIDC) and
c.TokenExchanger(c.BaseURL, c.BaseAPIEndpoint, oidcToken, workspaceID) instead
of the globals so tests can inject mocks for token fetch/exchange while
production code keeps existing behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@pkg/pro/retry_test.go`:
- Around line 69-86: The test TestDoWithRetry_AuthErrorThenSuccess currently
uses newTestClient() which leaves useOIDC=false so RefreshToken() is never
exercised; update the test to create a client with useOIDC enabled (or adjust
newTestClient to accept an OIDC flag) and inject a mock/stub for RefreshToken so
you can assert it gets called on the 401 path and that doWithRetry retries after
a successful RefreshToken; also add a second test that simulates RefreshToken
returning an error to verify doWithRetry aborts the retry flow on refresh
failure. Ensure references to doWithRetry, TestDoWithRetry_AuthErrorThenSuccess,
newTestClient, useOIDC, and RefreshToken are used to locate and modify the code.

In `@pkg/pro/retry.go`:
- Around line 87-96: The current fallback treats any non-*APIError as retryable
(using lastErr and apiErr) which also retries deterministic local failures like
ErrFailedToCreateAuthRequest; change the logic in the retry routine (where
apiErr, lastErr, attempt, cfg.maxRetries are used) to only retry genuine
transport errors: check if lastErr implements net.Error with Temporary() or
Timeout(), or is a *url.Error, or introduce/recognize a RetryableError marker
(e.g., errors.Is(lastErr, ErrRetryable) or a custom interface) and only
log/return nil to retry in those cases; explicitly exclude known non-retryable
sentinels such as ErrFailedToCreateAuthRequest so tests like
TestDoWithRetry_NetworkErrorRetried still cover network retries but
deterministic local failures are not retried.

---

Nitpick comments:
In `@pkg/pro/api_client.go`:
- Around line 177-195: RefreshToken currently calls getGitHubOIDCToken and
exchangeOIDCTokenForAtmosToken directly, making it hard to mock; add injectable
function fields on AtmosProAPIClient (e.g., OIDCTokenFetcher func(githubOIDC
string) (string, error) and TokenExchanger func(baseURL, baseAPIEndpoint,
oidcToken, workspaceID string) (string, error)), defaulting those fields to
getGitHubOIDCToken and exchangeOIDCTokenForAtmosToken when nil, then have
RefreshToken call c.OIDCTokenFetcher(c.atmosConfig.Settings.Pro.GithubOIDC) and
c.TokenExchanger(c.BaseURL, c.BaseAPIEndpoint, oidcToken, workspaceID) instead
of the globals so tests can inject mocks for token fetch/exchange while
production code keeps existing behavior.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: fd089ae5-391c-44a5-9e4b-db10fac7f78a

📥 Commits

Reviewing files that changed from the base of the PR and between dbcba35 and 43dec1c.

📒 Files selected for processing (11)
  • errors/errors.go
  • internal/exec/describe_affected.go
  • internal/exec/terraform_execute_helpers_exec.go
  • pkg/list/list_instances.go
  • pkg/pro/api_client.go
  • pkg/pro/api_client_instance_status.go
  • pkg/pro/api_client_instances.go
  • pkg/pro/api_error.go
  • pkg/pro/api_error_test.go
  • pkg/pro/retry.go
  • pkg/pro/retry_test.go

Comment thread pkg/pro/retry_test.go
Comment thread pkg/pro/retry.go Outdated
@osterman Erik Osterman (Cloud Posse) (osterman) added the patch A minor, backward compatible change label Mar 26, 2026
Address CodeRabbit review comments:
- Extract tokenRefresher interface so tests can verify RefreshToken is
  called on 401 and that refresh failure aborts retries
- Guard against retrying deterministic errors like ErrFailedToCreateAuthRequest
- Add tests for refresh failure abort and non-retryable non-API errors

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (2)
pkg/pro/retry_test.go (2)

24-38: Manual mock is pragmatic here, but consider mockgen for consistency.

Guidelines suggest using go.uber.org/mock/mockgen for mock generation. For trivial interfaces like tokenRefresher, manual mocks are readable and maintainable. If you prefer strict consistency, a //go:generate mockgen directive could replace this.

Low priority given the simplicity.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@pkg/pro/retry_test.go` around lines 24 - 38, The test currently defines a
manual mock type mockRefresher with methods RefreshToken and helper
newMockRefresher; to align with project guidelines optionally replace or
supplement this manual mock with an auto-generated one via mockgen: add a
//go:generate mockgen directive targeting the tokenRefresher interface, run
mockgen to produce the mock (naming it e.g. MockTokenRefresher), and update
tests to use the generated MockTokenRefresher instead of or alongside
mockRefresher (or keep the manual mock if you prefer simplicity)—this preserves
readability while enabling consistent mock generation.

130-171: Consider consolidating into table-driven test.

Three nearly identical tests for 400/403/404. A table-driven approach would reduce duplication:

♻️ Optional consolidation
func TestDoWithRetry_NonRetryableStatusCodes(t *testing.T) {
	cases := []struct {
		name   string
		status int
	}{
		{"400 Bad Request", 400},
		{"403 Forbidden", 403},
		{"404 Not Found", 404},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			s := &fakeSleeper{}
			cfg := retryConfig{maxRetries: 3, baseDelay: time.Second, sleeper: s}

			callCount := 0
			err := doWithRetry("TestOp", func() error {
				callCount++
				return &APIError{StatusCode: tc.status, Operation: "TestOp", Err: fmt.Errorf("error")}
			}, newMockRefresher(), cfg)

			require.Error(t, err)
			assert.Equal(t, 1, callCount, "should not retry on %d", tc.status)
			assert.Empty(t, s.sleeps)
		})
	}
}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@pkg/pro/retry_test.go` around lines 130 - 171, Replace the three nearly
identical tests TestDoWithRetry_NonRetryable400/403/404 with a single
table-driven test (e.g., TestDoWithRetry_NonRetryableStatusCodes) that iterates
over a slice of cases containing the status codes and names, and inside a t.Run
for each case call doWithRetry with the same setup (fakeSleeper, retryConfig,
newMockRefresher) and assert require.Error, callCount == 1, and that s.sleeps is
empty; reference and reuse symbols doWithRetry, retryConfig, fakeSleeper,
APIError and newMockRefresher so behavior and assertions remain identical.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@pkg/pro/retry_test.go`:
- Around line 24-38: The test currently defines a manual mock type mockRefresher
with methods RefreshToken and helper newMockRefresher; to align with project
guidelines optionally replace or supplement this manual mock with an
auto-generated one via mockgen: add a //go:generate mockgen directive targeting
the tokenRefresher interface, run mockgen to produce the mock (naming it e.g.
MockTokenRefresher), and update tests to use the generated MockTokenRefresher
instead of or alongside mockRefresher (or keep the manual mock if you prefer
simplicity)—this preserves readability while enabling consistent mock
generation.
- Around line 130-171: Replace the three nearly identical tests
TestDoWithRetry_NonRetryable400/403/404 with a single table-driven test (e.g.,
TestDoWithRetry_NonRetryableStatusCodes) that iterates over a slice of cases
containing the status codes and names, and inside a t.Run for each case call
doWithRetry with the same setup (fakeSleeper, retryConfig, newMockRefresher) and
assert require.Error, callCount == 1, and that s.sleeps is empty; reference and
reuse symbols doWithRetry, retryConfig, fakeSleeper, APIError and
newMockRefresher so behavior and assertions remain identical.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: af2a7609-2460-49bd-b91c-cf644a124783

📥 Commits

Reviewing files that changed from the base of the PR and between 43dec1c and e574a5f.

📒 Files selected for processing (2)
  • pkg/pro/retry.go
  • pkg/pro/retry_test.go

coderabbitai[bot]
coderabbitai Bot previously approved these changes Mar 26, 2026
@github-actions

Copy link
Copy Markdown

These changes were released in v1.213.0-test.6.

@github-actions

Copy link
Copy Markdown

These changes were released in v1.213.0-test.8.

Both features (chunked payload splitting from main, retry with backoff from
this branch) modify the same upload methods. Resolution keeps main's chunked
architecture and adds retry logic inside each send*Request helper so every
chunk is individually retried on transient 401/5xx failures.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Replace the hand-rolled retry implementation in pkg/pro/retry.go with the
generic pkg/retry package. The pro layer now builds a schema.RetryConfig and
delegates to retry.WithPredicate, keeping only the domain-specific error
classification (401 → token refresh, 5xx → retry, 400/403/404 → abort) as
a predicate function.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
pkg/pro/retry_test.go (1)

184-197: Backoff timing assertion has reasonable tolerance.

The test expects ~70ms minimum (10+20+40) but asserts ≥50ms to account for timing variability. Consider using a slightly tighter bound or asserting an upper bound too to catch potential issues where delays are too long.

Optional: add upper-bound check for completeness
 	// Exponential backoff: 10ms + 20ms + 40ms = 70ms minimum.
 	assert.GreaterOrEqual(t, elapsed.Milliseconds(), int64(50), "should have waited at least ~70ms (with tolerance)")
+	// Upper bound sanity check - shouldn't take more than 500ms even with jitter/scheduling delays.
+	assert.Less(t, elapsed.Milliseconds(), int64(500), "should complete within reasonable time")
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@pkg/pro/retry_test.go` around lines 184 - 197, Update
TestDoWithRetry_ExponentialBackoff to tighten the timing assertions: in the test
that constructs retryConfig and calls doWithRetry, change the lower-bound
assertion on elapsed to a slightly higher threshold (e.g.,
assert.GreaterOrEqual(..., int64(60))) to better reflect the expected ~70ms
backoff, and add an upper-bound assertion (e.g., assert.LessOrEqual(...,
int64(200))) to detect regressions where delays become excessively long;
reference the test TestDoWithRetry_ExponentialBackoff, the retryConfig used, and
the doWithRetry call when making the edits.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@pkg/pro/retry_test.go`:
- Around line 184-197: Update TestDoWithRetry_ExponentialBackoff to tighten the
timing assertions: in the test that constructs retryConfig and calls
doWithRetry, change the lower-bound assertion on elapsed to a slightly higher
threshold (e.g., assert.GreaterOrEqual(..., int64(60))) to better reflect the
expected ~70ms backoff, and add an upper-bound assertion (e.g.,
assert.LessOrEqual(..., int64(200))) to detect regressions where delays become
excessively long; reference the test TestDoWithRetry_ExponentialBackoff, the
retryConfig used, and the doWithRetry call when making the edits.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: f3920f45-4887-4f1b-aea6-49b1eab17469

📥 Commits

Reviewing files that changed from the base of the PR and between 85a18fe and 0edc361.

📒 Files selected for processing (4)
  • pkg/pro/api_client.go
  • pkg/pro/api_client_instances.go
  • pkg/pro/retry.go
  • pkg/pro/retry_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
  • pkg/pro/api_client_instances.go

@aknysh
Andriy Knysh (aknysh) merged commit 9c2a42f into main Mar 27, 2026
56 checks passed
@aknysh
Andriy Knysh (aknysh) deleted the osterman/robust-upload-retry branch March 27, 2026 03:09
@github-actions

Copy link
Copy Markdown

These changes were released in v1.213.0.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

patch A minor, backward compatible change size/m Medium size PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants