fix: preserve Retry-After on rate limit errors

.changeset/rate-limit-retry-after.md

@@ -0,0 +1,6 @@
+---
+"e2b": patch
+"@e2b/python-sdk": patch
+---
+
+Expose `Retry-After` on rate-limit errors.

packages/js-sdk/src/api/index.ts

@@ -5,6 +5,7 @@ import { defaultHeaders } from './metadata'
 import { ConnectionConfig } from '../connectionConfig'
 import { AuthenticationError, RateLimitError, SandboxError } from '../errors'
 import { createApiLogger } from '../logs'
+import { parseRetryAfter } from '../utils'
 
 export function handleApiError(
   response: FetchResponse<any, any, any>,
@@ -33,11 +34,14 @@ export function handleApiError(
   if (response.response.status === 429) {
     const message = 'Rate limit exceeded, please try again later'
     const content = response.error?.message ?? response.error
+    const retryAfter = parseRetryAfter(
+      response.response.headers?.get('Retry-After')
+    )
 
     if (content) {
-      return new RateLimitError(`${message} - ${content}`)
+      return new RateLimitError(`${message} - ${content}`, retryAfter)
     }
-    return new RateLimitError(message)
+    return new RateLimitError(message, retryAfter)
   }
 
   const message = response.error?.message ?? response.error

packages/js-sdk/src/envd/api.ts

@@ -11,7 +11,9 @@ import {
   SandboxNotFoundError,
   formatSandboxTimeoutError,
   AuthenticationError,
+  RateLimitError,
 } from '../errors'
+import { parseRetryAfter } from '../utils'
 import { StartResponse, ConnectResponse } from './process/process_pb'
 import { Code, ConnectError } from '@connectrpc/connect'
 import { WatchDirResponse } from './filesystem/filesystem_pb'
@@ -22,8 +24,6 @@ const DEFAULT_ERROR_MAP: Record<number, (message: string) => Error> = {
   400: (message) => new InvalidArgumentError(message),
   401: (message) => new AuthenticationError(message),
   404: (message) => new NotFoundError(message),
-  429: (message) =>
-    new SandboxError(`${message}: The requests are being rate limited.`),
   502: formatSandboxTimeoutError,
   507: (message) => new NotEnoughSpaceError(message),
 }
@@ -42,7 +42,7 @@ export async function handleEnvdApiError(
   },
   errorMap?: Record<number, (message: string) => Error>
 ) {
-  if (!res.error) {
+  if (res.error === undefined) {
     return
   }
 
@@ -56,6 +56,14 @@ export async function handleEnvdApiError(
     return errorMap[res.response.status]?.(message)
   }
 
+  if (res.response.status === 429) {
+    const retryAfter = parseRetryAfter(res.response.headers.get('Retry-After'))
+    return new RateLimitError(
+      `${message}: The requests are being rate limited.`,
+      retryAfter
+    )
+  }
+
   // Check if there is a default error mapping for this error code
   if (res.response.status in DEFAULT_ERROR_MAP) {
     return DEFAULT_ERROR_MAP[res.response.status]?.(message)

packages/js-sdk/src/errors.ts

@@ -134,9 +134,15 @@ export class TemplateError extends SandboxError {
  * Thrown when the API rate limit is exceeded.
  */
 export class RateLimitError extends SandboxError {
-  constructor(message: string) {
+  /**
+   * Number of seconds the caller should wait before retrying, if provided by the server.
+   */
+  readonly retryAfter?: number
+
+  constructor(message: string, retryAfter?: number) {
     super(message)
     this.name = 'RateLimitError'
+    this.retryAfter = retryAfter
   }
 }
 

packages/js-sdk/src/utils.ts

@@ -67,6 +67,26 @@ export function timeoutToSeconds(timeout: number): number {
   return Math.ceil(timeout / 1000)
 }
 
+export function parseRetryAfter(retryAfter: string | null | undefined) {
+  if (!retryAfter) {
+    return
+  }
+
+  const trimmed = retryAfter.trim()
+  if (/^\d+$/.test(trimmed)) {
+    return Number(trimmed)
+  }
+
+  if (/^[+-]?\d/.test(trimmed)) {
+    return
+  }
+
+  const retryAt = Date.parse(trimmed)
+  if (!Number.isNaN(retryAt)) {
+    return Math.max(0, Math.ceil((retryAt - Date.now()) / 1000))
+  }
+}
+
 export function dynamicRequire<T>(module: string): T {
   if (runtime === 'browser') {
     throw new Error('Browser runtime is not supported for require')

packages/js-sdk/tests/api/handleApiError.test.ts

@@ -9,14 +9,15 @@ import {
 function createMockResponse(
   status: number,
   error: unknown,
-  data?: unknown
+  data?: unknown,
+  headers?: Headers
 ): {
-  response: { status: number; ok: boolean }
+  response: { status: number; ok: boolean; headers?: Headers }
   error: unknown
   data: unknown
 } {
   return {
-    response: { status, ok: status >= 200 && status < 300 },
+    response: { status, ok: status >= 200 && status < 300, headers },
     error,
     data,
   }
@@ -90,6 +91,30 @@ describe('handleApiError', () => {
       assert.instanceOf(err, RateLimitError)
       assert.include(err?.message, 'Rate limit')
     })
+
+    test('preserves Retry-After for 429 rate limits', () => {
+      const res = createMockResponse(
+        429,
+        { message: 'Too many requests' },
+        undefined,
+        new Headers({ 'Retry-After': ' 60 ' })
+      )
+      const err = handleApiError(res as any)
+      assert.instanceOf(err, RateLimitError)
+      assert.equal((err as any).retryAfter, 60)
+    })
+
+    test('ignores invalid Retry-After values for 429 rate limits', () => {
+      const res = createMockResponse(
+        429,
+        { message: 'Too many requests' },
+        undefined,
+        new Headers({ 'Retry-After': '-1' })
+      )
+      const err = handleApiError(res as any)
+      assert.instanceOf(err, RateLimitError)
+      assert.isUndefined((err as any).retryAfter)
+    })
   })
 
   describe('success responses', () => {

packages/js-sdk/tests/envd/handleEnvdApiError.test.ts

@@ -0,0 +1,31 @@
+import { assert, describe, test } from 'vitest'
+import { handleEnvdApiError } from '../../src/envd/api'
+import { RateLimitError } from '../../src/errors'
+
+describe('handleEnvdApiError', () => {
+  test('preserves Retry-After for 429 rate limits', async () => {
+    const err = await handleEnvdApiError({
+      error: { message: 'Too many requests' },
+      response: new Response('', {
+        status: 429,
+        headers: { 'Retry-After': '60' },
+      }),
+    })
+
+    assert.instanceOf(err, RateLimitError)
+    assert.equal((err as any).retryAfter, 60)
+  })
+
+  test('handles 429 rate limits with empty error body', async () => {
+    const err = await handleEnvdApiError({
+      error: '',
+      response: new Response('Too many requests', {
+        status: 429,
+        headers: { 'Retry-After': '60' },
+      }),
+    })
+
+    assert.instanceOf(err, RateLimitError)
+    assert.equal((err as any).retryAfter, 60)
+  })
+})

packages/python-sdk/e2b/api/__init__.py

@@ -15,6 +15,7 @@
     AuthenticationException,
     RateLimitException,
     SandboxException,
+    parse_retry_after,
 )
 
 logger = logging.getLogger(__name__)
@@ -55,7 +56,11 @@ def handle_api_exception(
         message = f"{e.status_code}: Rate limit exceeded, please try again later."
         if body.get("message"):
             message += f" - {body['message']}"
-        return RateLimitException(message)
+        headers = getattr(e, "headers", {})
+        return RateLimitException(
+            message,
+            retry_after=parse_retry_after(headers.get("Retry-After")),
+        )
 
     if "message" in body:
         return default_exception_class(

packages/python-sdk/e2b/envd/api.py

@@ -9,7 +9,9 @@
     AuthenticationException,
     InvalidArgumentException,
     NotEnoughSpaceException,
+    RateLimitException,
     format_sandbox_timeout_exception,
+    parse_retry_after,
 )
 
 
@@ -20,9 +22,6 @@
     400: InvalidArgumentException,
     401: AuthenticationException,
     404: NotFoundException,
-    429: lambda message: SandboxException(
-        f"{message}: The requests are being rate limited."
-    ),
     502: format_sandbox_timeout_exception,
     507: NotEnoughSpaceException,
 }
@@ -52,7 +51,12 @@ def handle_envd_api_exception(
 
     res.read()
 
-    return format_envd_api_exception(res.status_code, get_message(res), error_map)
+    return format_envd_api_exception(
+        res.status_code,
+        get_message(res),
+        error_map,
+        parse_retry_after(res.headers.get("Retry-After")),
+    )
 
 
 async def ahandle_envd_api_exception(
@@ -65,24 +69,37 @@ async def ahandle_envd_api_exception(
 
     await res.aread()
 
-    return format_envd_api_exception(res.status_code, get_message(res), error_map)
+    return format_envd_api_exception(
+        res.status_code,
+        get_message(res),
+        error_map,
+        parse_retry_after(res.headers.get("Retry-After")),
+    )
 
 
 def format_envd_api_exception(
     status_code: int,
     message: str,
     error_map: Optional[dict[int, Callable[[str], Exception]]] = None,
+    retry_after: Optional[int] = None,
 ):
     """Map an HTTP status code and message to the appropriate exception.
 
     :param status_code: The HTTP status code.
     :param message: The error message from the response body.
     :param error_map: Optional map of HTTP status codes to exception factories that override the defaults.
+    :param retry_after: Optional number of seconds to wait before retrying rate-limited requests.
     :return: The corresponding exception.
     """
     if error_map and status_code in error_map:
         return error_map[status_code](message)
 
+    if status_code == 429:
+        return RateLimitException(
+            f"{message}: The requests are being rate limited.",
+            retry_after=retry_after,
+        )
+
     if status_code in _DEFAULT_API_ERROR_MAP:
         return _DEFAULT_API_ERROR_MAP[status_code](message)
 

packages/python-sdk/e2b/exceptions.py

@@ -1,3 +1,32 @@
+from datetime import datetime, timezone
+from email.utils import parsedate_to_datetime
+from math import ceil
+from typing import Optional
+
+
+def parse_retry_after(retry_after: Optional[str]) -> Optional[int]:
+    if not retry_after:
+        return None
+
+    retry_after = retry_after.strip()
+
+    if retry_after.isdecimal():
+        return int(retry_after)
+
+    if retry_after[:1] in ("+", "-") or retry_after[:1].isdigit():
+        return None
+
+    try:
+        retry_at = parsedate_to_datetime(retry_after)
+    except (TypeError, ValueError):
+        return None
+
+    if retry_at.tzinfo is None:
+        retry_at = retry_at.replace(tzinfo=timezone.utc)
+
+    return max(0, ceil((retry_at - datetime.now(timezone.utc)).total_seconds()))
+
+
 def format_sandbox_timeout_exception(message: str):
     return TimeoutException(
         f"{message}: This error is likely due to sandbox timeout. You can modify the sandbox timeout by passing 'timeout' when starting the sandbox or calling '.set_timeout' on the sandbox with the desired timeout."
@@ -118,6 +147,12 @@ class RateLimitException(SandboxException):
     Raised when the API rate limit is exceeded.
     """
 
+    retry_after: Optional[int]
+
+    def __init__(self, *args, retry_after: Optional[int] = None):
+        super().__init__(*args)
+        self.retry_after = retry_after
+
 
 class BuildException(Exception):
     """

packages/python-sdk/tests/test_rate_limit_exceptions.py

@@ -0,0 +1,41 @@
+import httpx
+
+from e2b.api import handle_api_exception
+from e2b.envd.api import handle_envd_api_exception
+from e2b.exceptions import RateLimitException
+
+
+class ApiErrorResponse:
+    status_code = 429
+    content = b'{"message": "Too many requests"}'
+    headers = {"Retry-After": " 60 "}
+
+
+def test_api_rate_limit_exception_preserves_retry_after():
+    err = handle_api_exception(ApiErrorResponse())
+
+    assert isinstance(err, RateLimitException)
+    assert err.retry_after == 60
+
+
+def test_api_rate_limit_exception_ignores_invalid_retry_after():
+    response = ApiErrorResponse()
+    response.headers = {"Retry-After": "-1"}
+
+    err = handle_api_exception(response)
+
+    assert isinstance(err, RateLimitException)
+    assert err.retry_after is None
+
+
+def test_envd_rate_limit_exception_preserves_retry_after():
+    response = httpx.Response(
+        429,
+        json={"message": "Too many requests"},
+        headers={"Retry-After": "60"},
+    )
+
+    err = handle_envd_api_exception(response)
+
+    assert isinstance(err, RateLimitException)
+    assert err.retry_after == 60