interrupt

Author: mishushakov

js/tests/interrupt.test.ts

@@ -0,0 +1,25 @@
+import { expect } from 'vitest'
+
+import { sandboxTest, wait } from './setup'
+
+sandboxTest(
+  'subsequent execution works after client timeout',
+  async ({ sandbox }) => {
+    // Start a long-running execution with a short timeout.
+    // This simulates a client disconnect: the SDK aborts the connection,
+    // which should trigger the server to interrupt the kernel (#213).
+    await expect(
+      sandbox.runCode('import time; time.sleep(30)', { timeoutMs: 3_000 })
+    ).rejects.toThrow()
+
+    // Wait for the server to detect the disconnect (via keepalive write
+    // failure) and interrupt the kernel.
+    await wait(5_000)
+
+    // Run a simple execution. Without the kernel interrupt fix, this would
+    // block behind the still-running sleep(30) and time out.
+    const result = await sandbox.runCode('1 + 1', { timeoutMs: 10_000 })
+    expect(result.text).toEqual('2')
+  },
+  60_000
+)

python/tests/async/test_async_interrupt.py

@@ -0,0 +1,25 @@
+import asyncio
+
+import pytest
+
+from e2b import TimeoutException
+from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
+
+
+async def test_subsequent_execution_works_after_client_timeout(
+    async_sandbox: AsyncSandbox,
+):
+    # Start a long-running execution with a short timeout.
+    # This simulates a client disconnect: the SDK closes the connection,
+    # which should trigger the server to interrupt the kernel (#213).
+    with pytest.raises(TimeoutException):
+        await async_sandbox.run_code("import time; time.sleep(30)", timeout=3)
+
+    # Wait for the server to detect the disconnect (via keepalive write
+    # failure) and interrupt the kernel.
+    await asyncio.sleep(5)
+
+    # Run a simple execution. Without the kernel interrupt fix, this would
+    # block behind the still-running sleep(30) and time out.
+    result = await async_sandbox.run_code("1 + 1", timeout=10)
+    assert result.text == "2"

python/tests/sync/test_interrupt.py

@@ -0,0 +1,23 @@
+import time
+
+import pytest
+
+from e2b import TimeoutException
+from e2b_code_interpreter.code_interpreter_sync import Sandbox
+
+
+def test_subsequent_execution_works_after_client_timeout(sandbox: Sandbox):
+    # Start a long-running execution with a short timeout.
+    # This simulates a client disconnect: the SDK closes the connection,
+    # which should trigger the server to interrupt the kernel (#213).
+    with pytest.raises(TimeoutException):
+        sandbox.run_code("import time; time.sleep(30)", timeout=3)
+
+    # Wait for the server to detect the disconnect (via keepalive write
+    # failure) and interrupt the kernel.
+    time.sleep(5)
+
+    # Run a simple execution. Without the kernel interrupt fix, this would
+    # block behind the still-running sleep(30) and time out.
+    result = sandbox.run_code("1 + 1", timeout=10)
+    assert result.text == "2"

template/server/messaging.py

@@ -4,6 +4,8 @@
 import uuid
 import asyncio
 
+import httpx
+
 from asyncio import Queue
 from typing import (
     Dict,
@@ -26,13 +28,15 @@
     OutputType,
     UnexpectedEndOfExecution,
 )
+from consts import JUPYTER_BASE_URL
 from errors import ExecutionError
 from envs import get_envs
 
 logger = logging.getLogger(__name__)
 
 MAX_RECONNECT_RETRIES = 3
 PING_TIMEOUT = 30
+KEEPALIVE_INTERVAL = 5  # seconds between keepalive pings during streaming
 
 
 class Execution:
@@ -97,6 +101,22 @@ async def connect(self):
             name="receive_message",
         )
 
+    async def interrupt(self):
+        """Interrupt the current kernel execution via the Jupyter REST API."""
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    f"{JUPYTER_BASE_URL}/api/kernels/{self.context_id}/interrupt"
+                )
+                if response.is_success:
+                    logger.info(f"Kernel {self.context_id} interrupted successfully")
+                else:
+                    logger.error(
+                        f"Failed to interrupt kernel {self.context_id}: {response.status_code}"
+                    )
+        except Exception as e:
+            logger.error(f"Error interrupting kernel {self.context_id}: {e}")
+
     def _get_execute_request(
         self, msg_id: str, code: Union[str, StrictStr], background: bool
     ) -> str:
@@ -239,7 +259,18 @@ async def _wait_for_result(self, message_id: str):
         queue = self._executions[message_id].queue
 
         while True:
-            output = await queue.get()
+            try:
+                output = await asyncio.wait_for(
+                    queue.get(), timeout=KEEPALIVE_INTERVAL
+                )
+            except asyncio.TimeoutError:
+                # Yield a keepalive so Starlette writes to the socket.
+                # If the client has disconnected, the write fails and
+                # uvicorn delivers http.disconnect, which cancels this
+                # generator via CancelledError.
+                yield {"type": "keepalive"}
+                continue
+
             if output.type == OutputType.END_OF_EXECUTION:
                 break
 
@@ -294,11 +325,6 @@ async def execute(
         if self._ws is None:
             raise Exception("WebSocket not connected")
 
-        # Lock only the setup + send phase, not the streaming phase.
-        # Results are read from a per-execution queue (keyed by message_id)
-        # so streaming doesn't need serialization. Releasing before streaming
-        # prevents client disconnect from holding the lock until the kernel
-        # finishes execution (see #213).
         async with self._lock:
             # Wait for any pending cleanup task to complete
             if self._cleanup_task and not self._cleanup_task.done():
@@ -367,22 +393,35 @@ async def execute(
                 )
                 await execution.queue.put(UnexpectedEndOfExecution())
 
-            # Schedule env var cleanup inside the lock so the next execution
-            # can wait for it. The task sends reset code to the kernel, which
-            # queues it after the current execution's code.
-            if env_vars:
+            # Stream the results.
+            # If the client disconnects (Starlette cancels the task), we
+            # interrupt the kernel so the next execution isn't blocked (#213).
+            client_disconnected = False
+            try:
+                async for item in self._wait_for_result(message_id):
+                    yield item
+            except (asyncio.CancelledError, GeneratorExit):
+                client_disconnected = True
+                logger.warning(
+                    f"Client disconnected during execution ({message_id}), interrupting kernel"
+                )
+                # Shield the interrupt from the ongoing cancellation so
+                # the HTTP request to the kernel actually completes.
+                try:
+                    await asyncio.shield(self.interrupt())
+                except asyncio.CancelledError:
+                    pass
+                raise
+            finally:
+                if message_id in self._executions:
+                    del self._executions[message_id]
+
+            # Clean up env vars in a separate request after the main code has run
+            if env_vars and not client_disconnected:
                 self._cleanup_task = asyncio.create_task(
                     self._cleanup_env_vars(env_vars)
                 )
 
-        # Stream the results without holding the lock
-        try:
-            async for item in self._wait_for_result(message_id):
-                yield item
-        finally:
-            if message_id in self._executions:
-                del self._executions[message_id]
-
     async def _receive_message(self):
         if not self._ws:
             logger.error("No WebSocket connection")
@@ -394,7 +433,8 @@ async def _receive_message(self):
         except Exception as e:
             logger.error(f"WebSocket received error while receiving messages: {str(e)}")
         finally:
-            # To prevent infinite hang, cancel all ongoing executions as results may be lost during reconnect.
+            # To prevent infinite hang, we need to cancel all ongoing execution as we could lost results during the reconnect
+            # Thanks to the locking, there can be either no ongoing execution or just one.
             for key, execution in self._executions.items():
                 await execution.queue.put(
                     Error(