From c239d1050cecc1ba6ccd581dd01de6f7ba243270 Mon Sep 17 00:00:00 2001 From: Ritwij Aryan Parmar Date: Tue, 26 May 2026 19:54:46 -0400 Subject: [PATCH] Preserve retry-after on rate limits --- packages/js-sdk/src/api/index.ts | 16 ++++++-- packages/js-sdk/src/envd/api.ts | 15 +++++++- packages/js-sdk/src/errors.ts | 38 ++++++++++++++++++- .../js-sdk/tests/api/handleApiError.test.ts | 25 ++++++++++-- packages/js-sdk/tests/envd/api.test.ts | 20 ++++++++++ packages/python-sdk/e2b/api/__init__.py | 10 ++++- packages/python-sdk/e2b/envd/api.py | 32 +++++++++++++--- packages/python-sdk/e2b/exceptions.py | 10 +++++ packages/python-sdk/e2b/rate_limit.py | 30 +++++++++++++++ .../tests/test_retry_after_rate_limit.py | 35 +++++++++++++++++ 10 files changed, 215 insertions(+), 16 deletions(-) create mode 100644 packages/js-sdk/tests/envd/api.test.ts create mode 100644 packages/python-sdk/e2b/rate_limit.py create mode 100644 packages/python-sdk/tests/test_retry_after_rate_limit.py diff --git a/packages/js-sdk/src/api/index.ts b/packages/js-sdk/src/api/index.ts index 7b6db519e4..26a13d7221 100644 --- a/packages/js-sdk/src/api/index.ts +++ b/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 diff --git a/packages/js-sdk/src/envd/api.ts b/packages/js-sdk/src/envd/api.ts index b0bd1704c8..7b6b7763c6 100644 --- a/packages/js-sdk/src/envd/api.ts +++ b/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 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) diff --git a/packages/js-sdk/src/errors.ts b/packages/js-sdk/src/errors.ts index e66ea07c61..841fc6d12f 100644 --- a/packages/js-sdk/src/errors.ts +++ b/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.` } /** diff --git a/packages/js-sdk/tests/api/handleApiError.test.ts b/packages/js-sdk/tests/api/handleApiError.test.ts index 1dc1e855e2..257806de6a 100644 --- a/packages/js-sdk/tests/api/handleApiError.test.ts +++ b/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 = {} ): { - 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', () => { diff --git a/packages/js-sdk/tests/envd/api.test.ts b/packages/js-sdk/tests/envd/api.test.ts new file mode 100644 index 0000000000..91fdb8bc25 --- /dev/null +++ b/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') + }) +}) diff --git a/packages/python-sdk/e2b/api/__init__.py b/packages/python-sdk/e2b/api/__init__.py index c30b037e96..eb167bcf5f 100644 --- a/packages/python-sdk/e2b/api/__init__.py +++ b/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( diff --git a/packages/python-sdk/e2b/envd/api.py b/packages/python-sdk/e2b/envd/api.py index 7e2a59c131..e64b16d599 100644 --- a/packages/python-sdk/e2b/envd/api.py +++ b/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) diff --git a/packages/python-sdk/e2b/exceptions.py b/packages/python-sdk/e2b/exceptions.py index 14587a5c0a..5291bb4e27 100644 --- a/packages/python-sdk/e2b/exceptions.py +++ b/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): """ diff --git a/packages/python-sdk/e2b/rate_limit.py b/packages/python-sdk/e2b/rate_limit.py new file mode 100644 index 0000000000..026f1ff4e6 --- /dev/null +++ b/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." diff --git a/packages/python-sdk/tests/test_retry_after_rate_limit.py b/packages/python-sdk/tests/test_retry_after_rate_limit.py new file mode 100644 index 0000000000..81b1101ed2 --- /dev/null +++ b/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)