Preserve Retry-After on SDK rate limit errors

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

@@ -4,7 +4,12 @@ import type { components, paths } from './schema.gen'
 import { defaultHeaders } from './metadata'
 import { createApiFetch } from './http2'
 import { ConnectionConfig } from '../connectionConfig'
-import { AuthenticationError, RateLimitError, SandboxError } from '../errors'
+import {
+  AuthenticationError,
+  RateLimitError,
+  SandboxError,
+  parseRetryAfter,
+} from '../errors'
 import { createApiLogger } from '../logs'
 
 export function handleApiError(
@@ -34,11 +39,16 @@ 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 retryAfterHeader = response.response.headers?.get('Retry-After')
+    const retryAfter = parseRetryAfter(retryAfterHeader)
 
     if (content) {
-      return new RateLimitError(`${message} - ${content}`)
+      return new RateLimitError(`${message} - ${content}`, {
+        retryAfter,
+        retryAfterHeader,
+      })
     }
-    return new RateLimitError(message)
+    return new RateLimitError(message, { retryAfter, retryAfterHeader })
   }
 
   const message = response.error?.message ?? response.error

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

@@ -11,6 +11,8 @@ import {
   SandboxNotFoundError,
   formatSandboxTimeoutError,
   AuthenticationError,
+  RateLimitError,
+  parseRetryAfter,
 } from '../errors'
 import { StartResponse, ConnectResponse } from './process/process_pb'
 import { Code, ConnectError } from '@connectrpc/connect'
@@ -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),
 }
@@ -56,6 +56,17 @@ export async function handleEnvdApiError(
     return errorMap[res.response.status]?.(message)
   }
 
+  if (res.response.status === 429) {
+    const retryAfterHeader = res.response.headers.get('Retry-After')
+    return new RateLimitError(
+      `${message}: The requests are being rate limited.`,
+      {
+        retryAfter: parseRetryAfter(retryAfterHeader),
+        retryAfterHeader,
+      }
+    )
+  }
+
   // 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,10 +134,44 @@ export class TemplateError extends SandboxError {
  * Thrown when the API rate limit is exceeded.
  */
 export class RateLimitError extends SandboxError {
-  constructor(message: string) {
-    super(message)
+  readonly retryAfter?: number
+  readonly retryAfterHeader?: string
+
+  constructor(
+    message: string,
+    opts: { retryAfter?: number; retryAfterHeader?: string | null } = {}
+  ) {
+    super(appendRetryAfter(message, opts.retryAfter))
     this.name = 'RateLimitError'
+    this.retryAfter = opts.retryAfter
+    this.retryAfterHeader = opts.retryAfterHeader ?? undefined
+  }
+}
+
+export function parseRetryAfter(retryAfterHeader?: string | null) {
+  if (!retryAfterHeader) {
+    return undefined
+  }
+
+  const trimmedRetryAfter = retryAfterHeader.trim()
+  if (/^-?\d+$/.test(trimmedRetryAfter)) {
+    const retryAfter = Number.parseInt(trimmedRetryAfter, 10)
+    return Math.max(retryAfter, 0)
+  }
+
+  const retryAt = Date.parse(trimmedRetryAfter)
+  if (Number.isNaN(retryAt)) {
+    return undefined
+  }
+
+  return Math.max(Math.floor((retryAt - Date.now()) / 1000), 0)
+}
+
+function appendRetryAfter(message: string, retryAfter?: number) {
+  if (retryAfter === undefined) {
+    return message
   }
+  return `${message} Retry after ${retryAfter} seconds.`
 }
 
 /**

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

@@ -9,14 +9,19 @@ import {
 function createMockResponse(
   status: number,
   error: unknown,
-  data?: unknown
+  data?: unknown,
+  headers: Record<string, string> = {}
 ): {
-  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: new Headers(headers),
+    },
     error,
     data,
   }
@@ -90,6 +95,20 @@ describe('handleApiError', () => {
       assert.instanceOf(err, RateLimitError)
       assert.include(err?.message, 'Rate limit')
     })
+
+    test('preserves Retry-After on 429', () => {
+      const res = createMockResponse(
+        429,
+        { message: 'Too many requests' },
+        undefined,
+        { 'Retry-After': '60' }
+      )
+      const err = handleApiError(res as any)
+      assert.instanceOf(err, RateLimitError)
+      assert.equal((err as RateLimitError).retryAfter, 60)
+      assert.equal((err as RateLimitError).retryAfterHeader, '60')
+      assert.include(err?.message, 'Retry after 60 seconds')
+    })
   })
 
   describe('success responses', () => {

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

@@ -0,0 +1,20 @@
+import { assert, describe, test } from 'vitest'
+import { handleEnvdApiError } from '../../src/envd/api'
+import { RateLimitError } from '../../src/errors'
+
+describe('handleEnvdApiError', () => {
+  test('preserves Retry-After on 429', async () => {
+    const err = await handleEnvdApiError({
+      error: { message: 'too many requests' },
+      response: new Response('', {
+        status: 429,
+        headers: { 'Retry-After': '45' },
+      }),
+    })
+
+    assert.instanceOf(err, RateLimitError)
+    assert.equal((err as RateLimitError).retryAfter, 45)
+    assert.equal((err as RateLimitError).retryAfterHeader, '45')
+    assert.include(err?.message, 'Retry after 45 seconds')
+  })
+})

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

@@ -16,6 +16,7 @@
     RateLimitException,
     SandboxException,
 )
+from e2b.rate_limit import append_retry_after, parse_retry_after
 
 logger = logging.getLogger(__name__)
 
@@ -55,7 +56,14 @@ 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", {})
+        retry_after_header = headers.get("Retry-After") or headers.get("retry-after")
+        retry_after = parse_retry_after(retry_after_header)
+        return RateLimitException(
+            append_retry_after(message, retry_after),
+            retry_after=retry_after,
+            retry_after_header=retry_after_header,
+        )
 
     if "message" in body:
         return default_exception_class(

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

@@ -5,12 +5,14 @@
 
 from e2b.exceptions import (
     SandboxException,
+    RateLimitException,
     NotFoundException,
     AuthenticationException,
     InvalidArgumentException,
     NotEnoughSpaceException,
     format_sandbox_timeout_exception,
 )
+from e2b.rate_limit import append_retry_after, parse_retry_after
 
 
 ENVD_API_FILES_ROUTE = "/files"
@@ -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,
+        retry_after_header=res.headers.get("Retry-After"),
+    )
 
 
 async def ahandle_envd_api_exception(
@@ -65,13 +69,19 @@ 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,
+        retry_after_header=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_header: Optional[str] = None,
 ):
     """Map an HTTP status code and message to the appropriate exception.
 
@@ -83,6 +93,18 @@ def format_envd_api_exception(
     if error_map and status_code in error_map:
         return error_map[status_code](message)
 
+    if status_code == 429:
+        retry_after = parse_retry_after(retry_after_header)
+        message = append_retry_after(
+            f"{message}: The requests are being rate limited.",
+            retry_after,
+        )
+        return RateLimitException(
+            message,
+            retry_after=retry_after,
+            retry_after_header=retry_after_header,
+        )
+
     if status_code in _DEFAULT_API_ERROR_MAP:
         return _DEFAULT_API_ERROR_MAP[status_code](message)
 

packages/python-sdk/e2b/exceptions.py

@@ -118,6 +118,16 @@ class RateLimitException(SandboxException):
     Raised when the API rate limit is exceeded.
     """
 
+    def __init__(
+        self,
+        message: str,
+        retry_after: int | None = None,
+        retry_after_header: str | None = None,
+    ):
+        super().__init__(message)
+        self.retry_after = retry_after
+        self.retry_after_header = retry_after_header
+
 
 class BuildException(Exception):
     """

packages/python-sdk/e2b/rate_limit.py

@@ -0,0 +1,30 @@
+from datetime import datetime, timezone
+from email.utils import parsedate_to_datetime
+from typing import Optional
+
+
+def parse_retry_after(retry_after_header: Optional[str]) -> Optional[int]:
+    if not retry_after_header:
+        return None
+
+    try:
+        retry_after = int(retry_after_header)
+        return max(retry_after, 0)
+    except ValueError:
+        pass
+
+    try:
+        retry_at = parsedate_to_datetime(retry_after_header)
+    except (TypeError, ValueError):
+        return None
+
+    if retry_at.tzinfo is None:
+        retry_at = retry_at.replace(tzinfo=timezone.utc)
+
+    return max(int((retry_at - datetime.now(timezone.utc)).total_seconds()), 0)
+
+
+def append_retry_after(message: str, retry_after: Optional[int]) -> str:
+    if retry_after is None:
+        return message
+    return f"{message} Retry after {retry_after} seconds."

packages/python-sdk/tests/test_retry_after_rate_limit.py

@@ -0,0 +1,35 @@
+import httpx
+
+from e2b.api import handle_api_exception
+from e2b.envd.api import handle_envd_api_exception
+from e2b.exceptions import RateLimitException
+
+
+class ApiError:
+    status_code = 429
+    content = b'{"message":"too many requests"}'
+    headers = {"Retry-After": "60"}
+
+
+def test_api_rate_limit_preserves_retry_after_header():
+    err = handle_api_exception(ApiError())
+
+    assert isinstance(err, RateLimitException)
+    assert err.retry_after == 60
+    assert err.retry_after_header == "60"
+    assert "Retry after 60 seconds" in str(err)
+
+
+def test_envd_rate_limit_preserves_retry_after_header():
+    res = httpx.Response(
+        429,
+        json={"message": "too many requests"},
+        headers={"Retry-After": "45"},
+    )
+
+    err = handle_envd_api_exception(res)
+
+    assert isinstance(err, RateLimitException)
+    assert err.retry_after == 45
+    assert err.retry_after_header == "45"
+    assert "Retry after 45 seconds" in str(err)