|
29 | 29 | from e2b_code_interpreter.exceptions import ( |
30 | 30 | format_execution_timeout_error, |
31 | 31 | format_request_timeout_error, |
| 32 | + format_sandbox_killed_error, |
32 | 33 | ) |
33 | 34 |
|
34 | 35 | logger = logging.getLogger(__name__) |
@@ -83,6 +84,23 @@ def _client(self) -> AsyncClient: |
83 | 84 | transport=get_transport(self.connection_config, http2=False), |
84 | 85 | ) |
85 | 86 |
|
| 87 | + async def _raise_if_sandbox_killed(self, err: Exception) -> None: |
| 88 | + """ |
| 89 | + Raises a descriptive exception if the connection error was caused by |
| 90 | + the sandbox being killed mid-request. If the sandbox is still running |
| 91 | + (or its state can't be determined), returns so the caller can re-raise |
| 92 | + the original error. |
| 93 | + """ |
| 94 | + try: |
| 95 | + running = await self.is_running() |
| 96 | + except Exception: |
| 97 | + # The state check itself failed, so we can't tell whether the |
| 98 | + # sandbox was killed — let the caller re-raise the original error |
| 99 | + # instead of wrongly claiming the sandbox is gone. |
| 100 | + return |
| 101 | + if not running: |
| 102 | + raise format_sandbox_killed_error() from err |
| 103 | + |
86 | 104 | @overload |
87 | 105 | async def run_code( |
88 | 106 | self, |
@@ -217,6 +235,9 @@ async def run_code( |
217 | 235 | raise format_execution_timeout_error() |
218 | 236 | except httpx.TimeoutException: |
219 | 237 | raise format_request_timeout_error() |
| 238 | + except (httpx.ReadError, httpx.RemoteProtocolError) as err: |
| 239 | + await self._raise_if_sandbox_killed(err) |
| 240 | + raise |
220 | 241 |
|
221 | 242 | async def create_code_context( |
222 | 243 | self, |
@@ -263,6 +284,9 @@ async def create_code_context( |
263 | 284 | return Context.from_json(data) |
264 | 285 | except httpx.TimeoutException: |
265 | 286 | raise format_request_timeout_error() |
| 287 | + except (httpx.ReadError, httpx.RemoteProtocolError) as err: |
| 288 | + await self._raise_if_sandbox_killed(err) |
| 289 | + raise |
266 | 290 |
|
267 | 291 | async def remove_code_context( |
268 | 292 | self, |
@@ -295,6 +319,9 @@ async def remove_code_context( |
295 | 319 | raise err |
296 | 320 | except httpx.TimeoutException: |
297 | 321 | raise format_request_timeout_error() |
| 322 | + except (httpx.ReadError, httpx.RemoteProtocolError) as err: |
| 323 | + await self._raise_if_sandbox_killed(err) |
| 324 | + raise |
298 | 325 |
|
299 | 326 | async def list_code_contexts(self) -> List[Context]: |
300 | 327 | """ |
@@ -323,6 +350,9 @@ async def list_code_contexts(self) -> List[Context]: |
323 | 350 | return [Context.from_json(context_data) for context_data in data] |
324 | 351 | except httpx.TimeoutException: |
325 | 352 | raise format_request_timeout_error() |
| 353 | + except (httpx.ReadError, httpx.RemoteProtocolError) as err: |
| 354 | + await self._raise_if_sandbox_killed(err) |
| 355 | + raise |
326 | 356 |
|
327 | 357 | async def restart_code_context( |
328 | 358 | self, |
@@ -354,3 +384,6 @@ async def restart_code_context( |
354 | 384 | raise err |
355 | 385 | except httpx.TimeoutException: |
356 | 386 | raise format_request_timeout_error() |
| 387 | + except (httpx.ReadError, httpx.RemoteProtocolError) as err: |
| 388 | + await self._raise_if_sandbox_killed(err) |
| 389 | + raise |
0 commit comments