Repository navigation
fix: retry Atmos Pro uploads on transient 401/5xx with exponential backoff - #2255
Conversation
…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>
Dependency Review✅ No vulnerabilities or license issues found.Snapshot WarningsEnsure 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 FilesNone |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ 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
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
📝 WalkthroughWalkthroughAdds 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 ( Changes
Sequence DiagramsequenceDiagram
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
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
pkg/pro/api_client.go (1)
177-195: MakeRefreshTokenmockable.
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 inpkg/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
📒 Files selected for processing (11)
errors/errors.gointernal/exec/describe_affected.gointernal/exec/terraform_execute_helpers_exec.gopkg/list/list_instances.gopkg/pro/api_client.gopkg/pro/api_client_instance_status.gopkg/pro/api_client_instances.gopkg/pro/api_error.gopkg/pro/api_error_test.gopkg/pro/retry.gopkg/pro/retry_test.go
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>
There was a problem hiding this comment.
🧹 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/mockgenfor mock generation. For trivial interfaces liketokenRefresher, manual mocks are readable and maintainable. If you prefer strict consistency, a//go:generate mockgendirective 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
📒 Files selected for processing (2)
pkg/pro/retry.gopkg/pro/retry_test.go
|
These changes were released in v1.213.0-test.6. |
|
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>
8173753
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>
There was a problem hiding this comment.
🧹 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
📒 Files selected for processing (4)
pkg/pro/api_client.gopkg/pro/api_client_instances.gopkg/pro/retry.gopkg/pro/retry_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- pkg/pro/api_client_instances.go
|
These changes were released in v1.213.0. |
what
UploadInstanceStatus,UploadAffectedStacks, andUploadInstanceswhy
references
Summary by CodeRabbit
New Features
Bug Fixes
Tests