From 399bab7905200e5c88d47ebaafa57c46fa328c1f Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Fri, 31 May 2024 16:33:06 +0200
Subject: [PATCH 1/3] Add max retries to websocket reconnect

---
 .../e2b/sandbox/websocket_client.py           | 48 +++++++++++++++++--
 1 file changed, 45 insertions(+), 3 deletions(-)

diff --git a/packages/python-sdk/e2b/sandbox/websocket_client.py b/packages/python-sdk/e2b/sandbox/websocket_client.py
index 65ead1136..5fdf00cd8 100644
--- a/packages/python-sdk/e2b/sandbox/websocket_client.py
+++ b/packages/python-sdk/e2b/sandbox/websocket_client.py
@@ -2,15 +2,18 @@
 
 import asyncio
 import logging
+import random
 
 from queue import Queue
 from threading import Event
-from typing import Any, Callable, List, Optional
+from typing import Any, Callable, List, Optional, AsyncIterator
 
-from websockets.legacy.client import WebSocketClientProtocol, connect
+from websockets.legacy.client import WebSocketClientProtocol, Connect
 from websockets.exceptions import ConnectionClosed
 from websockets.typing import Data
 
+from e2b.sandbox.exception import SandboxException
+
 logger = logging.getLogger(__name__)
 
 
@@ -73,7 +76,7 @@ async def _connect(self):
         ws_logger = logger.getChild("websockets.client")
         ws_logger.setLevel(logging.ERROR)
 
-        websocket_connector = connect(
+        websocket_connector = E2BConnect(
             self.url,
             max_size=None,
             max_queue=None,
@@ -118,3 +121,42 @@ async def close(self):
 
         if self._ws:
             await self._ws.close()
+
+
+class E2BConnect(Connect):
+    async def __aiter__(self) -> AsyncIterator[WebSocketClientProtocol]:
+        tries = 0
+        max_tries = 10
+        backoff_delay = 0.2
+        while True:
+            try:
+                async with self as protocol:
+                    yield protocol
+            except Exception:
+                tries += 1
+                if tries >= max_tries:
+                    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.2:
+                    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.
+                backoff_delay = backoff_delay * 1.1
+                backoff_delay = min(backoff_delay, 10)
+                continue
+            else:
+                # Connection succeeded - reset backoff delay
+                backoff_delay = 0.2

From 510d7fd57a7e37c0530383a35bd1ae36d2bbfe0f Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Fri, 31 May 2024 16:51:44 +0200
Subject: [PATCH 2/3] Make first retries faster

---
 .../python-sdk/e2b/sandbox/websocket_client.py  | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/packages/python-sdk/e2b/sandbox/websocket_client.py b/packages/python-sdk/e2b/sandbox/websocket_client.py
index 5fdf00cd8..62bc79a5e 100644
--- a/packages/python-sdk/e2b/sandbox/websocket_client.py
+++ b/packages/python-sdk/e2b/sandbox/websocket_client.py
@@ -125,20 +125,20 @@ async def close(self):
 
 class E2BConnect(Connect):
     async def __aiter__(self) -> AsyncIterator[WebSocketClientProtocol]:
-        tries = 0
-        max_tries = 10
-        backoff_delay = 0.2
+        retries = 0
+        max_retries = 12
+        backoff_delay = 0.1
         while True:
             try:
                 async with self as protocol:
                     yield protocol
             except Exception:
-                tries += 1
-                if tries >= max_tries:
+                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.2:
+                if backoff_delay == 0.1:
                     initial_delay = random.random()
                     self.logger.info(
                         "! connect failed; reconnecting in %.1f seconds",
@@ -154,9 +154,10 @@ async def __aiter__(self) -> AsyncIterator[WebSocketClientProtocol]:
                     )
                     await asyncio.sleep(int(backoff_delay))
                 # Increase delay with truncated exponential backoff.
-                backoff_delay = backoff_delay * 1.1
+                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.2
+                backoff_delay = 0.1

From f0cbe766146c6ab3b55b79f97eed1e9f9260fb10 Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Fri, 31 May 2024 16:52:34 +0200
Subject: [PATCH 3/3] Add changeset

---
 .changeset/small-oranges-roll.md | 5 +++++
 1 file changed, 5 insertions(+)
 create mode 100644 .changeset/small-oranges-roll.md

diff --git a/.changeset/small-oranges-roll.md b/.changeset/small-oranges-roll.md
new file mode 100644
index 000000000..b2a0c6894
--- /dev/null
+++ b/.changeset/small-oranges-roll.md
@@ -0,0 +1,5 @@
+---
+"@e2b/python-sdk": patch
+---
+
+Add max retries for websocket connection