Add option for proxying SDK requests to API and sandbox (#695)

PR that allows you to pass `proxy` (of the type ProxyTypes) when calling
methods in the SDK that send requests. Also, if you pass `proxy` when
creating a sandbox, all the sandbox requests (called on the
Sandbox/AsyncSandbox instance) will automatically use the same proxy.

The PR also contains a small bugfix—adding missing passing of limits and
passing these directly instead of through transport to be able to pass
the proxy parameter properly.

This change should also be compatible with the code interpreter—it uses
the client transport, which uses the passed proxy if there is any.

Author: ValentaTomas

.changeset/cuddly-buckets-tie.md

@@ -0,0 +1,5 @@
+---
+'@e2b/python-sdk': patch
+---
+
+Add support for passing proxy

packages/python-sdk/e2b/__init__.py

@@ -31,6 +31,7 @@
 )
 from .connection_config import (
     ConnectionConfig,
+    ProxyTypes,
 )
 from .exceptions import (
     SandboxException,

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

@@ -1,9 +1,9 @@
 import json
 import logging
+from typing import Optional
+from httpx import Limits
 from dataclasses import dataclass
 
-from typing import Optional, Union
-from httpx import HTTPTransport, AsyncHTTPTransport
 
 from e2b.api.client.client import AuthenticatedClient
 from e2b.connection_config import ConnectionConfig
@@ -50,7 +50,7 @@ def __init__(
         config: ConnectionConfig,
         require_api_key: bool = True,
         require_access_token: bool = False,
-        transport: Optional[Union[HTTPTransport, AsyncHTTPTransport]] = None,
+        limits: Optional[Limits] = None,
         *args,
         **kwargs,
     ):
@@ -97,7 +97,8 @@ def __init__(
                     "request": [self._log_request],
                     "response": [self._log_response],
                 },
-                "transport": transport,
+                "proxy": config.proxy,
+                "limits": limits,
             },
             headers=headers,
             token=token,

packages/python-sdk/e2b/connection_config.py

@@ -1,6 +1,7 @@
 import os
 
 from typing import Literal, Optional, Dict
+from httpx._types import ProxyTypes
 
 REQUEST_TIMEOUT: float = 30.0  # 30 seconds
 
@@ -37,12 +38,14 @@ def __init__(
         access_token: Optional[str] = None,
         request_timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
     ):
         self.domain = domain or ConnectionConfig._domain()
         self.debug = debug or ConnectionConfig._debug()
         self.api_key = api_key or ConnectionConfig._api_key()
         self.access_token = access_token or ConnectionConfig._access_token()
         self.headers = headers
+        self.proxy = proxy
 
         self.request_timeout = ConnectionConfig._get_request_timeout(
             REQUEST_TIMEOUT,

packages/python-sdk/e2b/sandbox_async/main.py

@@ -4,7 +4,7 @@
 from typing import Dict, Optional, TypedDict, overload
 from typing_extensions import Unpack
 
-from e2b.connection_config import ConnectionConfig
+from e2b.connection_config import ConnectionConfig, ProxyTypes
 from e2b.envd.api import ENVD_API_HEALTH_ROUTE, ahandle_envd_api_exception
 from e2b.exceptions import format_request_timeout_error
 from e2b.sandbox.main import SandboxSetup
@@ -106,7 +106,9 @@ def __init__(self, **opts: Unpack[AsyncSandboxOpts]):
         self._envd_api_url = f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(self.envd_port)}"
         self._envd_version = opts["envd_version"]
 
-        self._transport = AsyncTransportWithLogger(limits=self._limits)
+        self._transport = AsyncTransportWithLogger(
+            limits=self._limits, proxy=self._connection_config.proxy
+        )
         self._envd_api = httpx.AsyncClient(
             base_url=self.envd_api_url,
             transport=self._transport,
@@ -177,6 +179,7 @@ async def create(
         domain: Optional[str] = None,
         debug: Optional[bool] = None,
         request_timeout: Optional[float] = None,
+        proxy: Optional[ProxyTypes] = None,
     ):
         """
         Create a new sandbox.
@@ -189,6 +192,7 @@ async def create(
         :param envs: Custom environment variables for the sandbox
         :param api_key: E2B API Key to use for authentication, defaults to `E2B_API_KEY` environment variable
         :param request_timeout: Timeout for the request in **seconds**
+        :param proxy: Proxy to use for the request and for the **requests made to the returned sandbox**
 
         :return: sandbox instance for the new sandbox
 
@@ -199,6 +203,7 @@ async def create(
             domain=domain,
             debug=debug,
             request_timeout=request_timeout,
+            proxy=proxy,
         )
 
         if connection_config.debug:
@@ -214,6 +219,7 @@ async def create(
                 debug=debug,
                 request_timeout=request_timeout,
                 env_vars=envs,
+                proxy=proxy,
             )
             sandbox_id = response.sandbox_id
             envd_version = response.envd_version
@@ -231,13 +237,15 @@ async def connect(
         api_key: Optional[str] = None,
         domain: Optional[str] = None,
         debug: Optional[bool] = None,
+        proxy: Optional[ProxyTypes] = None,
     ):
         """
         Connect to an existing sandbox.
         With sandbox ID you can connect to the same sandbox from different places or environments (serverless functions, etc).
 
         :param sandbox_id: Sandbox ID
         :param api_key: E2B API Key to use for authentication, defaults to `E2B_API_KEY` environment variable
+        :param proxy: Proxy to use for the request and for the **requests made to the returned sandbox**
 
         :return: sandbox instance for the existing sandbox
 
@@ -253,6 +261,7 @@ async def connect(
             api_key=api_key,
             domain=domain,
             debug=debug,
+            proxy=proxy,
         )
 
         return cls(
@@ -286,20 +295,25 @@ async def kill(
         domain: Optional[str] = None,
         debug: Optional[bool] = None,
         request_timeout: Optional[float] = None,
+        proxy: Optional[ProxyTypes] = None,
     ) -> bool:
         """
         Kill the sandbox specified by sandbox ID.
 
         :param sandbox_id: Sandbox ID
         :param api_key: E2B API Key to use for authentication, defaults to `E2B_API_KEY` environment variable
         :param request_timeout: Timeout for the request in **seconds**
+        :param proxy: Proxy to use for the request
 
         :return: `True` if the sandbox was killed, `False` if the sandbox was not found
         """
         ...
 
     @class_method_variant("_cls_kill")
-    async def kill(self, request_timeout: Optional[float] = None) -> bool:  # type: ignore
+    async def kill(
+        self,
+        request_timeout: Optional[float] = None,
+    ) -> bool:  # type: ignore
         config_dict = self.connection_config.__dict__
         config_dict.pop("access_token", None)
         config_dict.pop("api_url", None)
@@ -309,7 +323,7 @@ async def kill(self, request_timeout: Optional[float] = None) -> bool:  # type:
 
         await SandboxApi._cls_kill(
             sandbox_id=self.sandbox_id,
-            **self.connection_config.__dict__,
+            **config_dict,
         )
 
     @overload
@@ -339,6 +353,7 @@ async def set_timeout(
         domain: Optional[str] = None,
         debug: Optional[bool] = None,
         request_timeout: Optional[float] = None,
+        proxy: Optional[ProxyTypes] = None,
     ) -> None:
         """
         Set the timeout of the specified sandbox.
@@ -350,6 +365,7 @@ async def set_timeout(
         :param sandbox_id: Sandbox ID
         :param timeout: Timeout for the sandbox in **seconds**
         :param request_timeout: Timeout for the request in **seconds**
+        :param proxy: Proxy to use for the request
         """
         ...
 
@@ -369,7 +385,7 @@ async def set_timeout(  # type: ignore
         await SandboxApi._cls_set_timeout(
             sandbox_id=self.sandbox_id,
             timeout=timeout,
-            **self.connection_config.__dict__,
+            **config_dict,
         )
 
     async def get_info(  # type: ignore
@@ -391,5 +407,5 @@ async def get_info(  # type: ignore
 
         return await SandboxApi.get_info(
             sandbox_id=self.sandbox_id,
-            **self.connection_config.__dict__,
+            **config_dict,
         )

packages/python-sdk/e2b/sandbox_async/sandbox_api.py

@@ -1,7 +1,9 @@
 import urllib.parse
+
 from typing import Optional, Dict, List
 from packaging.version import Version
 
+
 from e2b.sandbox.sandbox_api import SandboxInfo, SandboxApiBase, SandboxQuery
 from e2b.exceptions import TemplateException
 from e2b.api import AsyncApiClient, SandboxCreateResponse
@@ -13,7 +15,7 @@
     delete_sandboxes_sandbox_id,
     post_sandboxes,
 )
-from e2b.connection_config import ConnectionConfig
+from e2b.connection_config import ConnectionConfig, ProxyTypes
 from e2b.api import handle_api_exception
 
 
@@ -27,6 +29,7 @@ async def list(
         debug: Optional[bool] = None,
         request_timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
     ) -> List[SandboxInfo]:
         """
         List all running sandboxes.
@@ -37,6 +40,7 @@ async def list(
         :param debug: Enable debug mode, all requested are then sent to localhost
         :param request_timeout: Timeout for the request in **seconds**
         :param headers: Additional headers to send with the request
+        :param proxy: Proxy to use for the request
 
         :return: List of running sandboxes
         """
@@ -46,6 +50,7 @@ async def list(
             debug=debug,
             request_timeout=request_timeout,
             headers=headers,
+            proxy=proxy,
         )
 
         # Convert filters to the format expected by the API
@@ -58,7 +63,10 @@ async def list(
                 }
                 metadata = urllib.parse.urlencode(quoted_metadata)
 
-        async with AsyncApiClient(config) as api_client:
+        async with AsyncApiClient(
+            config,
+            limits=SandboxApiBase._limits,
+        ) as api_client:
             res = await get_sandboxes.asyncio_detailed(
                 client=api_client,
                 metadata=metadata,
@@ -96,6 +104,7 @@ async def get_info(
         debug: Optional[bool] = None,
         request_timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
     ) -> SandboxInfo:
         """
         Get the sandbox info.
@@ -105,6 +114,7 @@ async def get_info(
         :param debug: Debug mode, defaults to `E2B_DEBUG` environment variable
         :param request_timeout: Timeout for the request in **seconds**
         :param headers: Additional headers to send with the request
+        :param proxy: Proxy to use for the request
 
         :return: Sandbox info
         """
@@ -114,9 +124,13 @@ async def get_info(
             debug=debug,
             request_timeout=request_timeout,
             headers=headers,
+            proxy=proxy,
         )
 
-        async with AsyncApiClient(config) as api_client:
+        async with AsyncApiClient(
+            config,
+            limits=SandboxApiBase._limits,
+        ) as api_client:
             res = await get_sandboxes_sandbox_id.asyncio_detailed(
                 sandbox_id,
                 client=api_client,
@@ -151,20 +165,25 @@ async def _cls_kill(
         debug: Optional[bool] = None,
         request_timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
     ) -> bool:
         config = ConnectionConfig(
             api_key=api_key,
             domain=domain,
             debug=debug,
             request_timeout=request_timeout,
             headers=headers,
+            proxy=proxy,
         )
 
         if config.debug:
             # Skip killing the sandbox in debug mode
             return True
 
-        async with AsyncApiClient(config) as api_client:
+        async with AsyncApiClient(
+            config,
+            limits=SandboxApiBase._limits,
+        ) as api_client:
             res = await delete_sandboxes_sandbox_id.asyncio_detailed(
                 sandbox_id,
                 client=api_client,
@@ -188,20 +207,25 @@ async def _cls_set_timeout(
         debug: Optional[bool] = None,
         request_timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
     ) -> None:
         config = ConnectionConfig(
             api_key=api_key,
             domain=domain,
             debug=debug,
             request_timeout=request_timeout,
             headers=headers,
+            proxy=proxy,
         )
 
         if config.debug:
             # Skip setting the timeout in debug mode
             return
 
-        async with AsyncApiClient(config) as api_client:
+        async with AsyncApiClient(
+            config,
+            limits=SandboxApiBase._limits,
+        ) as api_client:
             res = await post_sandboxes_sandbox_id_timeout.asyncio_detailed(
                 sandbox_id,
                 client=api_client,
@@ -223,16 +247,21 @@ async def _create_sandbox(
         debug: Optional[bool] = None,
         request_timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
     ) -> SandboxCreateResponse:
         config = ConnectionConfig(
             api_key=api_key,
             domain=domain,
             debug=debug,
             request_timeout=request_timeout,
             headers=headers,
+            proxy=proxy,
         )
 
-        async with AsyncApiClient(config) as api_client:
+        async with AsyncApiClient(
+            config,
+            limits=SandboxApiBase._limits,
+        ) as api_client:
             res = await post_sandboxes.asyncio_detailed(
                 body=NewSandbox(
                     template_id=template,
@@ -268,7 +297,3 @@ async def _create_sandbox(
                 ),
                 envd_version=res.parsed.envd_version,
             )
-
-    @staticmethod
-    def _get_sandbox_id(sandbox_id: str, client_id: str) -> str:
-        return f"{sandbox_id}-{client_id}"