From 03b77b2f65318d2b45761aa9f174045263659cfa Mon Sep 17 00:00:00 2001
From: winklemad <winklemad@outlook.com>
Date: Sun, 12 Jul 2026 12:25:23 +0530
Subject: [PATCH] Retry HTTP 503 responses

`retry_if_specific_error` listed the retryable status codes as
`{429, 500, 502, 504, 500}` -- 500 appears twice and 503 is missing, so a
transient HTTP 503 Service Unavailable is never retried even though its sibling
upstream errors 502 and 504 are. 503 is the canonical "try again later" status
(load balancer draining, deploy, overload), so a request that hits one is
returned to the caller as an error on the first attempt.

Replace the duplicate 500 with 503, making the set `{429, 500, 502, 503, 504}`.

Added a test asserting the transient statuses are retried and that non-transient
or non-`ESMProteinError` errors are not.
---
 esm/sdk/retry.py      |  2 +-
 esm/sdk/retry_test.py | 21 +++++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)
 create mode 100644 esm/sdk/retry_test.py

diff --git a/esm/sdk/retry.py b/esm/sdk/retry.py
index 302d6cf0..3c37c487 100644
--- a/esm/sdk/retry.py
+++ b/esm/sdk/retry.py
@@ -23,8 +23,8 @@ def retry_if_specific_error(exception):
         429,
         500,
         502,
+        503,
         504,
-        500,
     }
 
 
diff --git a/esm/sdk/retry_test.py b/esm/sdk/retry_test.py
new file mode 100644
index 00000000..882a807a
--- /dev/null
+++ b/esm/sdk/retry_test.py
@@ -0,0 +1,21 @@
+"""Tests for retry.py"""
+
+from esm.sdk.api import ESMProteinError
+from esm.sdk.retry import retry_if_specific_error
+
+
+def test_retry_if_specific_error_retries_transient_statuses():
+    # Rate-limit and transient server errors are retried. 503 Service
+    # Unavailable must be retried alongside its 500/502/504 siblings.
+    for code in [429, 500, 502, 503, 504]:
+        err = ESMProteinError(error_code=code, error_msg="transient")
+        assert retry_if_specific_error(err), f"HTTP {code} should be retryable"
+
+
+def test_retry_if_specific_error_skips_non_transient_and_non_esm():
+    # A non-transient client error is not retried.
+    assert not retry_if_specific_error(
+        ESMProteinError(error_code=404, error_msg="not found")
+    )
+    # Exceptions that are not ESMProteinError are not retried.
+    assert not retry_if_specific_error(ValueError("boom"))