e2b-code-interpreter 0.0.10a0__tar.gz → 0.0.11a0__tar.gz

{e2b_code_interpreter-0.0.10a0 → e2b_code_interpreter-0.0.11a0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: e2b-code-interpreter
-Version: 0.0.10a0
+Version: 0.0.11a0
 Summary: E2B Code Interpreter - Stateful code execution
 Home-page: https://e2b.dev/
 License: Apache-2.0
@@ -14,9 +14,11 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: e2b (>=0.17.1)
-Requires-Dist: pydantic (>1,<3)
+Requires-Dist: e2b (==0.17.2a20)
+Requires-Dist: pydantic
+Requires-Dist: requests (>=2.32.3,<3.0.0)
 Requires-Dist: websocket-client (>=1.7.0,<2.0.0)
+Requires-Dist: websockets (>=12.0,<13.0)
 Project-URL: Bug Tracker, https://github.com/e2b-dev/code-interpreter/issues
 Project-URL: Repository, https://github.com/e2b-dev/e2b-code-interpreter/tree/python
 Description-Content-Type: text/markdown

{e2b_code_interpreter-0.0.10a0 → e2b_code_interpreter-0.0.11a0}/e2b_code_interpreter/__init__.py

@@ -2,3 +2,4 @@ from e2b import *
 
 from .main import CodeInterpreter, JupyterExtension
 from .models import Execution, Error, Result, KernelException, MIMEType, Logs
+from .messaging import CellMessage

e2b_code_interpreter-0.0.11a0/e2b_code_interpreter/constants.py

@@ -0,0 +1 @@
+TIMEOUT = 60

e2b_code_interpreter-0.0.11a0/e2b_code_interpreter/future.py

@@ -0,0 +1,46 @@
+import asyncio
+
+from asyncio import AbstractEventLoop
+from concurrent.futures import Future, CancelledError
+from typing import Any, Awaitable, Callable, Generic, List, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+# Check if using event is not better for most use cases
+class DeferredFuture(Generic[T]):
+    def __init__(self, cleanup_list: Optional[List[Callable[[], Any]]] = None):
+        self._future = Future()
+        if cleanup_list is not None:
+            cleanup_list.append(self.cancel)
+
+    def __call__(self, result: T):
+        if not self._future.done():
+            self._future.set_result(result)
+
+    def result(self, timeout: Optional[float] = None) -> T:
+        try:
+            return self._future.result(timeout=timeout)
+        # The thread was interrupted while waiting, it's okay to release the thread
+        except CancelledError:
+            pass
+
+    def done(self) -> bool:
+        return self._future.done()
+
+    def cancel(self):
+        if not self._future.done():
+            self._future.cancel()
+
+    def reject(self, reason: Exception):
+        if not self._future.done():
+            self._future.set_exception(reason)
+
+
+def run_async_func_in_loop(loop: AbstractEventLoop, coro: Awaitable):
+    asyncio.set_event_loop(loop)
+
+    try:
+        return loop.run_until_complete(coro)
+    finally:
+        loop.close()

{e2b_code_interpreter-0.0.10a0 → e2b_code_interpreter-0.0.11a0}/e2b_code_interpreter/main.py

@@ -3,16 +3,14 @@ from __future__ import annotations
 import logging
 import threading
 import uuid
-
 import requests
 
 from concurrent.futures import Future
 from typing import Any, Callable, List, Optional, Dict
+from e2b import Sandbox
 
-from e2b import EnvVars, ProcessMessage, Sandbox
-from e2b.constants import TIMEOUT
-
-from e2b_code_interpreter.messaging import JupyterKernelWebSocket
+from e2b_code_interpreter.constants import TIMEOUT
+from e2b_code_interpreter.messaging import CellMessage, JupyterKernelWebSocket
 from e2b_code_interpreter.models import KernelException, Execution, Result
 
 logger = logging.getLogger(__name__)
@@ -23,42 +21,56 @@ class CodeInterpreter(Sandbox):
     E2B code interpreter sandbox extension.
     """
 
-    template = "code-interpreter-multikernel"
+    default_template = "code-interpreter-stateful"
 
     def __init__(
         self,
         template: Optional[str] = None,
+        timeout: Optional[int] = None,
+        metadata: Optional[Dict[str, str]] = None,
         api_key: Optional[str] = None,
-        cwd: Optional[str] = None,
-        env_vars: Optional[EnvVars] = None,
-        timeout: Optional[float] = TIMEOUT,
-        on_stdout: Optional[Callable[[ProcessMessage], Any]] = None,
-        on_stderr: Optional[Callable[[ProcessMessage], Any]] = None,
-        on_exit: Optional[Callable[[int], Any]] = None,
-        **kwargs,
+        domain: Optional[str] = None,
+        debug: Optional[bool] = None,
+        sandbox_id: Optional[str] = None,
+        request_timeout: float = TIMEOUT,
     ):
         super().__init__(
-            template=template or self.template,
-            api_key=api_key,
-            cwd=cwd,
-            env_vars=env_vars,
+            template=template,
             timeout=timeout,
-            on_stdout=on_stdout,
-            on_stderr=on_stderr,
-            on_exit=on_exit,
-            **kwargs,
+            metadata=metadata,
+            api_key=api_key,
+            domain=domain,
+            debug=debug,
+            sandbox_id=sandbox_id,
+            request_timeout=request_timeout,
         )
-        self.notebook = JupyterExtension(self, timeout=timeout)
-        # Close all the websocket connections when the interpreter is closed
-        self._process_cleanup.append(self.notebook.close)
+        self.notebook = JupyterExtension(self, request_timeout=request_timeout)
 
+    def close(self):
+        # Close all the websocket connections to the kernels
+        self.notebook.close()
+
+    def get_protocol(self, base_protocol: str = "http") -> str:
+        """
+        The function decides whether to use the secure or insecure protocol.
+
+        :param base_protocol: Specify the specific protocol you want to use. Do not include the `s` in `https` or `wss`.
+        :param secure: Specify whether you want to use the secure protocol or not.
+
+        :return: Protocol for the connection to the sandbox
+        """
+        return base_protocol if self._connection_config.debug else f"{base_protocol}s"
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+        super().__exit__(exc_type, exc_value, traceback)
 
-class JupyterExtension:
 
-    def __init__(self, sandbox: CodeInterpreter, timeout: Optional[float] = TIMEOUT):
+class JupyterExtension:
+    def __init__(self, sandbox: CodeInterpreter, request_timeout: float = TIMEOUT):
         self._sandbox = sandbox
         self._kernel_id_set = Future()
-        self._start_connecting_to_default_kernel(timeout=timeout)
+        self._start_connecting_to_default_kernel(request_timeout=request_timeout)
         self._connected_kernels: Dict[str, Future[JupyterKernelWebSocket]] = {}
         self._default_kernel_id: Optional[str] = None
 
@@ -66,10 +78,11 @@ class JupyterExtension:
         self,
         code: str,
         kernel_id: Optional[str] = None,
-        on_stdout: Optional[Callable[[ProcessMessage], Any]] = None,
-        on_stderr: Optional[Callable[[ProcessMessage], Any]] = None,
+        on_stdout: Optional[Callable[[CellMessage], Any]] = None,
+        on_stderr: Optional[Callable[[CellMessage], Any]] = None,
         on_result: Optional[Callable[[Result], Any]] = None,
-        timeout: Optional[float] = TIMEOUT,
+        timeout: float = TIMEOUT,
+        request_timeout: float = TIMEOUT,
     ) -> Execution:
         """
         Execute code in a notebook cell.
@@ -93,7 +106,11 @@ class JupyterExtension:
             ws = ws_future.result(timeout=timeout)
         else:
             logger.debug(f"Creating new websocket connection to kernel {kernel_id}")
-            ws = self._connect_to_kernel_ws(kernel_id, None, timeout=timeout)
+            ws = self._connect_to_kernel_ws(
+                kernel_id,
+                None,
+                request_timeout=request_timeout,
+            )
 
         message_id = ws.send_execution_message(code, on_stdout, on_stderr, on_result)
         logger.debug(
@@ -124,7 +141,7 @@ class JupyterExtension:
         self,
         cwd: str = "/home/user",
         kernel_name: Optional[str] = None,
-        timeout: Optional[float] = TIMEOUT,
+        request_timeout: Optional[float] = TIMEOUT,
     ) -> str:
         """
         Creates a new kernel, this can be useful if you want to have multiple independent code execution environments.
@@ -138,18 +155,23 @@ class JupyterExtension:
         :param kernel_name:
             Specifies which kernel should be used, useful if you have multiple kernel types.
             If not provided, the default kernel will be used.
-        :param timeout: Timeout for the kernel creation request.
+        :param request_timeout: Timeout for the kernel creation request.
         :return: Kernel id of the created kernel
         """
         kernel_name = kernel_name or "python3"
 
-        data = {"path": str(uuid.uuid4()), "kernel": {"name": kernel_name}, "type": "notebook", "name": str(uuid.uuid4())}
+        data = {
+            "path": str(uuid.uuid4()),
+            "kernel": {"name": kernel_name},
+            "type": "notebook",
+            "name": str(uuid.uuid4()),
+        }
         logger.debug(f"Creating kernel with data: {data}")
 
         response = requests.post(
-            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/sessions",
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_host(8888)}/api/sessions",
             json=data,
-            timeout=timeout,
+            timeout=request_timeout,
         )
         if not response.ok:
             raise KernelException(f"Failed to create kernel: {response.text}")
@@ -159,9 +181,9 @@ class JupyterExtension:
         kernel_id = session_data["kernel"]["id"]
 
         response = requests.patch(
-            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/sessions/{session_id}",
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_host(8888)}/api/sessions/{session_id}",
             json={"path": cwd},
-            timeout=timeout,
+            timeout=request_timeout,
         )
         if not response.ok:
             raise KernelException(f"Failed to create kernel: {response.text}")
@@ -169,12 +191,15 @@ class JupyterExtension:
         logger.debug(f"Created kernel {kernel_id}")
 
         threading.Thread(
-            target=self._connect_to_kernel_ws, args=(kernel_id, session_id, timeout)
+            target=self._connect_to_kernel_ws,
+            args=(kernel_id, session_id, request_timeout),
         ).start()
         return kernel_id
 
     def restart_kernel(
-        self, kernel_id: Optional[str] = None, timeout: Optional[float] = TIMEOUT
+        self,
+        kernel_id: Optional[str] = None,
+        request_timeout: Optional[float] = TIMEOUT,
     ) -> None:
         """
         Restarts an existing Jupyter kernel. This can be useful to reset the kernel's state or to recover from errors.
@@ -190,8 +215,8 @@ class JupyterExtension:
         logger.debug(f"Closed websocket connection to kernel {kernel_id}")
 
         response = requests.post(
-            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/kernels/{kernel_id}/restart",
-            timeout=timeout,
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_host(8888)}/api/kernels/{kernel_id}/restart",
+            timeout=request_timeout,
         )
         if not response.ok:
             raise KernelException(f"Failed to restart kernel {kernel_id}")
@@ -199,11 +224,13 @@ class JupyterExtension:
         logger.debug(f"Restarted kernel {kernel_id}")
 
         threading.Thread(
-            target=self._connect_to_kernel_ws, args=(kernel_id, None, timeout)
+            target=self._connect_to_kernel_ws, args=(kernel_id, None, request_timeout)
         ).start()
 
     def shutdown_kernel(
-        self, kernel_id: Optional[str] = None, timeout: Optional[float] = TIMEOUT
+        self,
+        kernel_id: Optional[str] = None,
+        request_timeout: Optional[float] = TIMEOUT,
     ) -> None:
         """
         Shuts down an existing Jupyter kernel. This method is used to gracefully terminate a kernel's process.
@@ -219,15 +246,15 @@ class JupyterExtension:
         logger.debug(f"Closed websocket connection to kernel {kernel_id}")
 
         response = requests.delete(
-            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/kernels/{kernel_id}",
-            timeout=timeout,
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_host(8888)}/api/kernels/{kernel_id}",
+            timeout=request_timeout,
         )
         if not response.ok:
             raise KernelException(f"Failed to shutdown kernel {kernel_id}")
 
         logger.debug(f"Shutdown kernel {kernel_id}")
 
-    def list_kernels(self, timeout: Optional[float] = TIMEOUT) -> List[str]:
+    def list_kernels(self, request_timeout: Optional[float] = TIMEOUT) -> List[str]:
         """
         Lists all available Jupyter kernels.
 
@@ -238,8 +265,8 @@ class JupyterExtension:
         :return: List of kernel ids
         """
         response = requests.get(
-            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/kernels",
-            timeout=timeout,
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_host(8888)}/api/kernels",
+            timeout=request_timeout,
         )
 
         if not response.ok:
@@ -256,7 +283,10 @@ class JupyterExtension:
             ws.result().close()
 
     def _connect_to_kernel_ws(
-        self, kernel_id: str, session_id: Optional[str], timeout: Optional[float] = TIMEOUT
+        self,
+        kernel_id: str,
+        session_id: Optional[str],
+        request_timeout: float = TIMEOUT,
     ) -> JupyterKernelWebSocket:
         """
         Establishes a WebSocket connection to a specified Jupyter kernel.
@@ -272,18 +302,18 @@ class JupyterExtension:
 
         session_id = session_id or str(uuid.uuid4())
         ws = JupyterKernelWebSocket(
-            url=f"{self._sandbox.get_protocol('ws')}://{self._sandbox.get_hostname(8888)}/api/kernels/{kernel_id}/channels",
-            session_id=session_id
+            url=f"{self._sandbox.get_protocol('ws')}://{self._sandbox.get_host(8888)}/api/kernels/{kernel_id}/channels",
+            session_id=session_id,
         )
 
-        ws.connect(timeout=timeout)
+        ws.connect(timeout=request_timeout)
         logger.debug(f"Connected to kernel's ({kernel_id}) websocket.")
 
         future.set_result(ws)
         return ws
 
     def _start_connecting_to_default_kernel(
-        self, timeout: Optional[float] = TIMEOUT
+        self, request_timeout: float = TIMEOUT
     ) -> None:
         """
         Start connecting to the default kernel in a separate thread to avoid blocking the main thread.
@@ -292,17 +322,19 @@ class JupyterExtension:
         logger.debug("Starting to connect to the default kernel")
 
         def setup_default_kernel():
-            kernel_id = self._sandbox.filesystem.read(
-                "/root/.jupyter/kernel_id", timeout=timeout
+            kernel_id = self._sandbox.files.read(
+                "/root/.jupyter/kernel_id", request_timeout=request_timeout
             )
 
-            if kernel_id is None and not self._sandbox.is_open:
+            if kernel_id is None and not self._sandbox.is_running(
+                request_timeout=request_timeout
+            ):
                 return
 
             kernel_id = kernel_id.strip()
 
             logger.debug(f"Default kernel id: {kernel_id}")
-            self._connect_to_kernel_ws(kernel_id, None, timeout=timeout)
+            self._connect_to_kernel_ws(kernel_id, None, request_timeout=request_timeout)
             self._kernel_id_set.set_result(kernel_id)
 
         threading.Thread(target=setup_default_kernel).start()

{e2b_code_interpreter-0.0.10a0 → e2b_code_interpreter-0.0.11a0}/e2b_code_interpreter/messaging.py

@@ -3,21 +3,38 @@ import logging
 import threading
 import time
 import uuid
+
 from concurrent.futures import Future
 from queue import Queue
 from typing import Callable, Dict, Any, Optional
+from pydantic import BaseModel
+from e2b import TimeoutException
 
-from e2b import ProcessMessage
-from e2b.constants import TIMEOUT
-from e2b.sandbox import TimeoutException
-from e2b.sandbox.websocket_client import WebSocket
-from e2b.utils.future import DeferredFuture
-
+from e2b_code_interpreter.websocket_client import WebSocket
+from e2b_code_interpreter.future import DeferredFuture
 from e2b_code_interpreter.models import Execution, Result, Error
+from e2b_code_interpreter.constants import TIMEOUT
+
 
 logger = logging.getLogger(__name__)
 
 
+class CellMessage(BaseModel):
+    """
+    A message from a process.
+    """
+
+    line: str
+    error: bool = False
+    timestamp: int
+    """
+    Unix epoch in nanoseconds
+    """
+
+    def __str__(self):
+        return self.line
+
+
 class CellExecution:
     """
     Represents the execution of a cell in the Jupyter kernel.
@@ -26,14 +43,14 @@ class CellExecution:
 
     input_accepted: bool = False
 
-    on_stdout: Optional[Callable[[ProcessMessage], Any]] = None
-    on_stderr: Optional[Callable[[ProcessMessage], Any]] = None
+    on_stdout: Optional[Callable[[CellMessage], Any]] = None
+    on_stderr: Optional[Callable[[CellMessage], Any]] = None
     on_result: Optional[Callable[[Result], Any]] = None
 
     def __init__(
         self,
-        on_stdout: Optional[Callable[[ProcessMessage], Any]] = None,
-        on_stderr: Optional[Callable[[ProcessMessage], Any]] = None,
+        on_stdout: Optional[Callable[[CellMessage], Any]] = None,
+        on_stderr: Optional[Callable[[CellMessage], Any]] = None,
         on_result: Optional[Callable[[Result], Any]] = None,
     ):
         self.partial_result = Execution()
@@ -44,7 +61,6 @@ class CellExecution:
 
 
 class JupyterKernelWebSocket:
-
     def __init__(self, url: str, session_id: str):
         self.url = url
         self.session_id = session_id
@@ -78,7 +94,7 @@ class JupyterKernelWebSocket:
                 queue_in=self._queue_in,
                 queue_out=self._queue_out,
                 started=started,
-                stopped=self._stopped
+                stopped=self._stopped,
             ).run,
             daemon=True,
             name="e2b-code-interpreter-websocket",
@@ -128,8 +144,8 @@ class JupyterKernelWebSocket:
     def send_execution_message(
         self,
         code: str,
-        on_stdout: Optional[Callable[[ProcessMessage], Any]] = None,
-        on_stderr: Optional[Callable[[ProcessMessage], Any]] = None,
+        on_stdout: Optional[Callable[[CellMessage], Any]] = None,
+        on_stderr: Optional[Callable[[CellMessage], Any]] = None,
         on_result: Optional[Callable[[Result], Any]] = None,
     ) -> str:
         message_id = str(uuid.uuid4())
@@ -187,7 +203,7 @@ class JupyterKernelWebSocket:
                 execution.logs.stdout.append(data["content"]["text"])
                 if cell.on_stdout:
                     cell.on_stdout(
-                        ProcessMessage(
+                        CellMessage(
                             line=data["content"]["text"],
                             timestamp=time.time_ns(),
                         )
@@ -197,7 +213,7 @@ class JupyterKernelWebSocket:
                 execution.logs.stderr.append(data["content"]["text"])
                 if cell.on_stderr:
                     cell.on_stderr(
-                        ProcessMessage(
+                        CellMessage(
                             line=data["content"]["text"],
                             error=True,
                             timestamp=time.time_ns(),

{e2b_code_interpreter-0.0.10a0 → e2b_code_interpreter-0.0.11a0}/e2b_code_interpreter/models.py

@@ -1,6 +1,6 @@
 import copy
 from typing import List, Optional, Iterable, Dict
-from pydantic import BaseModel, field_serializer
+from pydantic import BaseModel
 
 
 class Error(BaseModel):
@@ -211,6 +211,19 @@ class Logs(BaseModel):
     "List of strings printed to stderr by prints, subprocesses, etc."
 
 
+def serialize_results(results: List[Result]) -> List[Dict[str, str]]:
+    """
+    Serializes the results to JSON.
+    This method is used by the Pydantic JSON encoder.
+    """
+    serialized = []
+    for result in results:
+        serialized_dict = {key: result[key] for key in result.formats()}
+        serialized_dict["text"] = result.text
+        serialized.append(serialized_dict)
+    return serialized
+
+
 class Execution(BaseModel):
     """
     Represents the result of a cell execution.
@@ -218,6 +231,9 @@ class Execution(BaseModel):
 
     class Config:
         arbitrary_types_allowed = True
+        json_encoders = {
+            List[Result]: serialize_results,
+        }
 
     results: List[Result] = []
     "List of the result of the cell (interactively interpreted last line), display calls (e.g. matplotlib plots)."
@@ -245,19 +261,6 @@ class Execution(BaseModel):
         """
         return self.model_dump_json(exclude_none=True)
 
-    @field_serializer("results", when_used="json")
-    def serialize_results(results: List[Result]) -> List[Dict[str, str]]:
-        """
-        Serializes the results to JSON.
-        This method is used by the Pydantic JSON encoder.
-        """
-        serialized = []
-        for result in results:
-            serialized_dict = {key: result[key] for key in result.formats()}
-            serialized_dict["text"] = result.text
-            serialized.append(serialized_dict)
-        return serialized
-
 
 class KernelException(Exception):
     """

e2b_code_interpreter-0.0.11a0/e2b_code_interpreter/websocket_client.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import random
+
+from queue import Queue
+from threading import Event
+from typing import Any, Callable, List, Optional, AsyncIterator
+
+from websockets.legacy.client import WebSocketClientProtocol, Connect
+from websockets.exceptions import ConnectionClosed
+from websockets.typing import Data
+
+from e2b import SandboxException
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocket:
+    def __init__(
+        self,
+        url: str,
+        started: Event,
+        stopped: Event,
+        queue_in: Queue[dict],
+        queue_out: Queue[Data],
+    ):
+        self._ws: Optional[WebSocketClientProtocol] = None
+        self.url = url
+        self.started = started
+        self.stopped = stopped
+        self._process_cleanup: List[Callable[[], Any]] = []
+        self._queue_in = queue_in
+        self._queue_out = queue_out
+
+    def run(self):
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        loop.run_until_complete(self.async_run())
+
+    async def async_run(self):
+        await self._connect()
+        await self.close()
+
+    async def _send_message(self):
+        logger.debug("WebSocket starting to send messages")
+        while True:
+            if self._queue_in.empty():
+                await asyncio.sleep(0)
+                continue
+            message = self._queue_in.get()
+            logger.debug(f"WebSocket message to send: {message}")
+            if self._ws:
+                await self._ws.send(message)
+                logger.debug(f"WebSocket message sent: {message}")
+                self._queue_in.task_done()
+            else:
+                logger.error("No WebSocket connection")
+
+    async def _receive_message(self):
+        try:
+            if not self._ws:
+                logger.error("No WebSocket connection")
+                return
+            async for message in self._ws:
+                logger.debug(f"WebSocket received message: {message}".strip())
+                self._queue_out.put(message)
+        except Exception as e:
+            logger.error(f"WebSocket received error while receiving messages: {e}")
+
+    async def _connect(self):
+        logger.debug(f"WebSocket connecting to {self.url}")
+
+        ws_logger = logger.getChild("websockets.client")
+        ws_logger.setLevel(logging.ERROR)
+
+        websocket_connector = E2BConnect(
+            self.url,
+            max_size=None,
+            max_queue=None,
+            logger=ws_logger,
+        )
+
+        websocket_connector.BACKOFF_MIN = 1
+        websocket_connector.BACKOFF_FACTOR = 1
+        websocket_connector.BACKOFF_INITIAL = 0.2  # type: ignore
+
+        async for websocket in websocket_connector:
+            try:
+                self._ws = websocket
+                self.started.set()
+                logger.info(f"WebSocket connected to {self.url}")
+
+                send_task = asyncio.create_task(
+                    self._send_message(), name="send_message"
+                )
+                self._process_cleanup.append(send_task.cancel)
+
+                receive_task = asyncio.create_task(
+                    self._receive_message(), name="receive_message"
+                )
+                self._process_cleanup.append(receive_task.cancel)
+
+                while not self.stopped.is_set():
+                    await asyncio.sleep(0)
+
+                logger.info("WebSocket stopped")
+                break
+            except ConnectionClosed:
+                logger.warning("WebSocket disconnected, it will try to reconnect")
+                if self.stopped.is_set():
+                    break
+
+    async def close(self):
+        for cancel in self._process_cleanup:
+            cancel()
+
+        self._process_cleanup.clear()
+
+        if self._ws:
+            await self._ws.close()
+
+
+class E2BConnect(Connect):
+    async def __aiter__(self) -> AsyncIterator[WebSocketClientProtocol]:
+        retries = 0
+        max_retries = 12
+        backoff_delay = 0.1
+        while True:
+            try:
+                async with self as protocol:
+                    yield protocol
+            except Exception:
+                retries += 1
+                if retries >= max_retries:
+                    raise SandboxException("Failed to connect to the server")
+                # Add a random initial delay between 0 and 5 seconds.
+                # See 7.2.3. Recovering from Abnormal Closure in RFC 6544.
+                if backoff_delay == 0.1:
+                    initial_delay = random.random()
+                    self.logger.info(
+                        "! connect failed; reconnecting in %.1f seconds",
+                        initial_delay,
+                        exc_info=True,
+                    )
+                    await asyncio.sleep(initial_delay)
+                else:
+                    self.logger.info(
+                        "! connect failed again; retrying in %d seconds",
+                        int(backoff_delay),
+                        exc_info=True,
+                    )
+                    await asyncio.sleep(int(backoff_delay))
+                # Increase delay with truncated exponential backoff.
+                if retries > 4:
+                    backoff_delay = backoff_delay * 1.2
+                backoff_delay = min(backoff_delay, 10)
+                continue
+            else:
+                # Connection succeeded - reset backoff delay
+                backoff_delay = 0.1

{e2b_code_interpreter-0.0.10a0 → e2b_code_interpreter-0.0.11a0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "e2b-code-interpreter"
-version = "0.0.10a0"
+version = "0.0.11a0"
 description = "E2B Code Interpreter - Stateful code execution"
 authors = ["e2b <hello@e2b.dev>"]
 license = "Apache-2.0"
@@ -12,9 +12,11 @@ packages = [{ include = "e2b_code_interpreter" }]
 [tool.poetry.dependencies]
 python = "^3.8"
 
-pydantic = ">1, <3"
+pydantic = "*"
 websocket-client = "^1.7.0"
-e2b = ">=0.17.1"
+websockets = "^12.0"
+e2b = "0.17.2a20"
+requests = "^2.32.3"
 
 [tool.poetry.group.dev.dependencies]
 black = "^24.3.0"