diff --git a/.changeset/wild-donkeys-hunt.md b/.changeset/wild-donkeys-hunt.md new file mode 100644 index 0000000000..b320923dbb --- /dev/null +++ b/.changeset/wild-donkeys-hunt.md @@ -0,0 +1,5 @@ +--- +'@e2b/code-interpreter-python': patch +--- + +Abort Jupyter requests at the deadline they were given instead of twice it. The SDK's pyqwest-backed transport derives a single whole-request deadline by summing httpx's `read` and `write` phases, and every request set both to the requested timeout — so `run_code(timeout=1)` ran for up to 2 seconds, and the context requests (`request_timeout`) likewise bought twice their value. diff --git a/packages/code-interpreter-python/e2b_code_interpreter/_jupyter_timeout.py b/packages/code-interpreter-python/e2b_code_interpreter/_jupyter_timeout.py new file mode 100644 index 0000000000..3d32e7d32c --- /dev/null +++ b/packages/code-interpreter-python/e2b_code_interpreter/_jupyter_timeout.py @@ -0,0 +1,22 @@ +from typing import Optional + +import httpx + + +def _jupyter_timeout( + timeout: Optional[float], *, connect: Optional[float] = None +) -> httpx.Timeout: + """Build the `httpx.Timeout` every Jupyter request derives its deadline from. + + The SDK's pyqwest-backed transport has no per-phase timeouts: it collapses + them into a single whole-request deadline, derived as the sum of the + `read` and `write` phases (`connect` and `pool` are ignored). The whole + budget therefore goes on `read` and `write` stays at zero, which makes the + deadline `timeout` and not `2 * timeout`. This matches the JS SDK, which + aborts the execution on a `timeout`-long timer. A `None` timeout disables + the deadline entirely and sets no phase at all — the transport still bounds + connect on its own. + """ + if timeout is None: + return httpx.Timeout(None) + return httpx.Timeout(timeout, write=0, connect=connect) diff --git a/packages/code-interpreter-python/e2b_code_interpreter/code_interpreter_async.py b/packages/code-interpreter-python/e2b_code_interpreter/code_interpreter_async.py index 50c8e42869..bf76265e1a 100644 --- a/packages/code-interpreter-python/e2b_code_interpreter/code_interpreter_async.py +++ b/packages/code-interpreter-python/e2b_code_interpreter/code_interpreter_async.py @@ -10,6 +10,7 @@ ) from e2b.api.client_async import get_transport +from e2b_code_interpreter._jupyter_timeout import _jupyter_timeout from e2b_code_interpreter.constants import ( DEFAULT_TEMPLATE, JUPYTER_PORT, @@ -217,21 +218,8 @@ async def run_code( }, headers=headers, # `timeout` bounds the execution, `request_timeout` only the - # connect. Every non-connect phase must carry `timeout`: the - # SDK's pyqwest-backed transport collapses the per-phase - # timeouts into a single whole-request deadline and takes the - # longest of them, so leaving `request_timeout` on the write - # and pool phases would raise the floor to - # `max(timeout, request_timeout)` and silently ignore any - # `timeout` shorter than it. This matches the JS SDK, which - # aborts the execution on a `timeout`-long timer. - # `timeout=0` disables the deadline entirely; the transport - # still bounds connect on its own. - timeout=( - httpx.Timeout(timeout, connect=request_timeout) - if timeout is not None - else httpx.Timeout(None) - ), + # connect; `_jupyter_timeout` holds the transport details. + timeout=_jupyter_timeout(timeout, connect=request_timeout), ) as response: err = await aextract_exception(response) if err: @@ -294,7 +282,9 @@ async def create_code_context( f"{self._jupyter_url}/contexts", headers=headers, json=data, - timeout=request_timeout or self.connection_config.request_timeout, + timeout=_jupyter_timeout( + self.connection_config.get_request_timeout(request_timeout) + ), ) err = await aextract_exception(response) @@ -334,7 +324,7 @@ async def remove_code_context( response = await self._client.delete( f"{self._jupyter_url}/contexts/{context_id}", headers=headers, - timeout=self.connection_config.request_timeout, + timeout=_jupyter_timeout(self.connection_config.request_timeout), ) err = await aextract_exception(response) @@ -364,7 +354,7 @@ async def list_code_contexts(self) -> List[Context]: response = await self._client.get( f"{self._jupyter_url}/contexts", headers=headers, - timeout=self.connection_config.request_timeout, + timeout=_jupyter_timeout(self.connection_config.request_timeout), ) err = await aextract_exception(response) @@ -403,7 +393,7 @@ async def restart_code_context( response = await self._client.post( f"{self._jupyter_url}/contexts/{context_id}/restart", headers=headers, - timeout=self.connection_config.request_timeout, + timeout=_jupyter_timeout(self.connection_config.request_timeout), ) err = await aextract_exception(response) diff --git a/packages/code-interpreter-python/e2b_code_interpreter/code_interpreter_sync.py b/packages/code-interpreter-python/e2b_code_interpreter/code_interpreter_sync.py index 29d43f634e..ef4072e64e 100644 --- a/packages/code-interpreter-python/e2b_code_interpreter/code_interpreter_sync.py +++ b/packages/code-interpreter-python/e2b_code_interpreter/code_interpreter_sync.py @@ -6,6 +6,7 @@ from e2b import Sandbox as BaseSandbox, InvalidArgumentException from e2b.api.client_sync import get_transport +from e2b_code_interpreter._jupyter_timeout import _jupyter_timeout from e2b_code_interpreter.constants import ( DEFAULT_TEMPLATE, JUPYTER_PORT, @@ -214,21 +215,8 @@ def run_code( }, headers=headers, # `timeout` bounds the execution, `request_timeout` only the - # connect. Every non-connect phase must carry `timeout`: the - # SDK's pyqwest-backed transport collapses the per-phase - # timeouts into a single whole-request deadline and takes the - # longest of them, so leaving `request_timeout` on the write - # and pool phases would raise the floor to - # `max(timeout, request_timeout)` and silently ignore any - # `timeout` shorter than it. This matches the JS SDK, which - # aborts the execution on a `timeout`-long timer. - # `timeout=0` disables the deadline entirely; the transport - # still bounds connect on its own. - timeout=( - httpx.Timeout(timeout, connect=request_timeout) - if timeout is not None - else httpx.Timeout(None) - ), + # connect; `_jupyter_timeout` holds the transport details. + timeout=_jupyter_timeout(timeout, connect=request_timeout), ) as response: err = extract_exception(response) if err: @@ -291,7 +279,9 @@ def create_code_context( f"{self._jupyter_url}/contexts", json=data, headers=headers, - timeout=request_timeout or self.connection_config.request_timeout, + timeout=_jupyter_timeout( + self.connection_config.get_request_timeout(request_timeout) + ), ) err = extract_exception(response) @@ -331,7 +321,7 @@ def remove_code_context( response = self._client.delete( f"{self._jupyter_url}/contexts/{context_id}", headers=headers, - timeout=self.connection_config.request_timeout, + timeout=_jupyter_timeout(self.connection_config.request_timeout), ) err = extract_exception(response) @@ -361,7 +351,7 @@ def list_code_contexts(self) -> List[Context]: response = self._client.get( f"{self._jupyter_url}/contexts", headers=headers, - timeout=self.connection_config.request_timeout, + timeout=_jupyter_timeout(self.connection_config.request_timeout), ) err = extract_exception(response) @@ -401,7 +391,7 @@ def restart_code_context( response = self._client.post( f"{self._jupyter_url}/contexts/{context_id}/restart", headers=headers, - timeout=self.connection_config.request_timeout, + timeout=_jupyter_timeout(self.connection_config.request_timeout), ) err = extract_exception(response) diff --git a/packages/code-interpreter-python/tests/test_execute_timeout.py b/packages/code-interpreter-python/tests/test_execute_timeout.py index 0ea523d93c..d6b554a177 100644 --- a/packages/code-interpreter-python/tests/test_execute_timeout.py +++ b/packages/code-interpreter-python/tests/test_execute_timeout.py @@ -1,16 +1,21 @@ """The execution timeout must reach the transport as the request deadline. Regression test for the SDK's pyqwest-backed transport collapsing httpx's -per-phase timeouts into a single whole-request deadline equal to the longest -phase. With the previous `(request_timeout, timeout, request_timeout, -request_timeout)` tuple that deadline came out as `max(timeout, -request_timeout)`, so any `timeout` shorter than `request_timeout` was silently -ignored and an execution ran on until `request_timeout` instead. +per-phase timeouts into a single whole-request deadline, which it derives as +the sum of the `read` and `write` phases (`connect` and `pool` are ignored). +Setting both phases to `timeout` therefore gave a deadline of `2 * timeout`, +so an execution was aborted at twice the `timeout` it asked for. """ +import socket +import threading +import time + import httpx import pytest +from pyqwest.httpx._transport import convert_timeout +from e2b import TimeoutException from e2b.connection_config import ConnectionConfig from e2b_code_interpreter.code_interpreter_async import AsyncSandbox @@ -34,12 +39,16 @@ def stream(self, *args, **kwargs): self._captured.update(kwargs) raise _Stop + def post(self, *args, **kwargs): + self._captured.update(kwargs) + raise _Stop -def _sandbox(cls, captured: dict): + +def _sandbox(cls, captured=None, connection_config=None, url="http://127.0.0.1:9"): class _Fake(cls): @property def connection_config(self): - return ConnectionConfig( + return connection_config or ConnectionConfig( api_key="x", domain="e2b.app", request_timeout=REQUEST_TIMEOUT ) @@ -57,10 +66,12 @@ def traffic_access_token(self): @property def _jupyter_url(self): - return "http://127.0.0.1:9" + return url @property def _client(self): + if captured is None: + return super()._client return _CapturingClient(captured) return _Fake.__new__(_Fake) @@ -69,12 +80,60 @@ def _client(self): def _captured_timeout(captured: dict) -> httpx.Timeout: timeout = captured["timeout"] assert isinstance(timeout, httpx.Timeout), ( - "a plain tuple lets the transport pick the longest phase as the deadline; " - f"got {timeout!r}" + f"expected an httpx.Timeout carrying the deadline; got {timeout!r}" ) return timeout +def _derived_deadline(timeout: httpx.Timeout) -> float | None: + """The single operation deadline the transport derives for a request.""" + return convert_timeout({"timeout": timeout.as_dict()}) + + +class _StallingJupyterServer: + """Plays a Jupyter server that sends the headers, then stalls. + + The body only follows `stall` seconds later, so the deadline the transport + derives from the request timeout is the only thing that can end the request + before then. + """ + + RESULT_LINE = b'{"type":"result","is_main_result":true,"text":"42"}\n' + + def __init__(self, stall: float): + self._stall = stall + self._socket = socket.socket() + self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self._socket.bind(("127.0.0.1", 0)) + self._socket.listen(1) + self.url = f"http://127.0.0.1:{self._socket.getsockname()[1]}" + + def __enter__(self): + threading.Thread(target=self._serve, daemon=True).start() + return self + + def __exit__(self, *_): + self._socket.close() + + def _serve(self): + try: + connection, _ = self._socket.accept() + except OSError: + return + with connection: + connection.recv(65536) + connection.sendall( + b"HTTP/1.1 200 OK\r\n" + b"Content-Type: application/x-ndjson\r\n" + b"Connection: close\r\n\r\n" + ) + time.sleep(self._stall) + try: + connection.sendall(self.RESULT_LINE) + except OSError: + pass + + @pytest.mark.parametrize("timeout", [3, 10, 300]) def test_execution_timeout_is_the_request_deadline(timeout): captured: dict = {} @@ -82,12 +141,12 @@ def test_execution_timeout_is_the_request_deadline(timeout): _sandbox(Sandbox, captured).run_code("1 + 1", timeout=timeout) tmo = _captured_timeout(captured) - # Every non-connect phase carries `timeout`, so the deadline the transport - # derives is `timeout` and not `max(timeout, request_timeout)`. + # The transport derives its deadline by summing `read` and `write`, so the + # whole budget sits on one phase and the other contributes nothing. assert tmo.read == timeout - assert tmo.write == timeout - assert tmo.pool == timeout + assert tmo.write == 0 assert tmo.connect == REQUEST_TIMEOUT + assert _derived_deadline(tmo) == timeout @pytest.mark.parametrize("timeout", [3, 10, 300]) @@ -98,9 +157,9 @@ async def test_async_execution_timeout_is_the_request_deadline(timeout): tmo = _captured_timeout(captured) assert tmo.read == timeout - assert tmo.write == timeout - assert tmo.pool == timeout + assert tmo.write == 0 assert tmo.connect == REQUEST_TIMEOUT + assert _derived_deadline(tmo) == timeout def test_zero_timeout_disables_the_deadline(): @@ -121,3 +180,100 @@ async def test_async_zero_timeout_disables_the_deadline(): tmo = _captured_timeout(captured) assert (tmo.read, tmo.write, tmo.pool, tmo.connect) == (None, None, None, None) + + +# Long enough that a deadline of twice the requested timeout also elapses +# before the server sends the body. +STALL = 4.0 + + +def test_execution_is_aborted_after_timeout_seconds(): + with _StallingJupyterServer(STALL) as server: + started = time.monotonic() + with pytest.raises(TimeoutException, match="Execution timed out"): + _sandbox(Sandbox, url=server.url).run_code("1 + 1", timeout=1) + elapsed = time.monotonic() - started + + assert elapsed < 1.5, f"run_code(timeout=1) was aborted after {elapsed:.2f}s" + + +async def test_async_execution_is_aborted_after_timeout_seconds(): + with _StallingJupyterServer(STALL) as server: + started = time.monotonic() + with pytest.raises(TimeoutException, match="Execution timed out"): + await _sandbox(AsyncSandbox, url=server.url).run_code("1 + 1", timeout=1) + elapsed = time.monotonic() - started + + assert elapsed < 1.5, f"run_code(timeout=1) was aborted after {elapsed:.2f}s" + + +def test_execution_within_timeout_is_not_aborted(): + with _StallingJupyterServer(1.0) as server: + execution = _sandbox(Sandbox, url=server.url).run_code("1 + 1", timeout=5) + + assert execution.text == "42" + + +def test_context_request_is_aborted_after_request_timeout_seconds(): + config = ConnectionConfig(api_key="x", domain="e2b.app", request_timeout=1) + with _StallingJupyterServer(STALL) as server: + started = time.monotonic() + with pytest.raises(TimeoutException, match="Request timed out"): + _sandbox( + Sandbox, connection_config=config, url=server.url + ).list_code_contexts() + elapsed = time.monotonic() - started + + assert elapsed < 1.5, f"list_code_contexts() was aborted after {elapsed:.2f}s" + + +async def test_async_context_request_is_aborted_after_request_timeout_seconds(): + config = ConnectionConfig(api_key="x", domain="e2b.app", request_timeout=1) + with _StallingJupyterServer(STALL) as server: + started = time.monotonic() + with pytest.raises(TimeoutException, match="Request timed out"): + await _sandbox( + AsyncSandbox, connection_config=config, url=server.url + ).list_code_contexts() + elapsed = time.monotonic() - started + + assert elapsed < 1.5, f"list_code_contexts() was aborted after {elapsed:.2f}s" + + +@pytest.mark.parametrize( + "request_timeout, expected_read, expected_write", + [(0, None, None), (7, 7, 0), (None, REQUEST_TIMEOUT, 0)], +) +def test_context_request_timeout_resolution( + request_timeout, expected_read, expected_write +): + captured: dict = {} + with pytest.raises(_Stop): + _sandbox(Sandbox, captured).create_code_context(request_timeout=request_timeout) + + tmo = _captured_timeout(captured) + # An explicit `request_timeout=0` disables the deadline — it must not + # fall back to the configured default — while `None` keeps the default. + # A disabled deadline sets no phase at all, matching the execution path. + assert tmo.read == expected_read + assert tmo.write == expected_write + assert _derived_deadline(tmo) == expected_read + + +@pytest.mark.parametrize( + "request_timeout, expected_read, expected_write", + [(0, None, None), (7, 7, 0), (None, REQUEST_TIMEOUT, 0)], +) +async def test_async_context_request_timeout_resolution( + request_timeout, expected_read, expected_write +): + captured: dict = {} + with pytest.raises(_Stop): + await _sandbox(AsyncSandbox, captured).create_code_context( + request_timeout=request_timeout + ) + + tmo = _captured_timeout(captured) + assert tmo.read == expected_read + assert tmo.write == expected_write + assert _derived_deadline(tmo) == expected_read