release lock on client disconnect

Author: mishushakov

js/tests/timeout.test.ts

@@ -0,0 +1,28 @@
+import { expect } from 'vitest'
+import { sandboxTest } from './setup'
+
+// Regression test for issue #213: asyncio.Lock not released on client disconnect.
+// Uses its own context to avoid blocking other parallel tests with the sleep.
+sandboxTest(
+  'execution after timeout is not blocked',
+  async ({ sandbox }) => {
+    const context = await sandbox.createCodeContext()
+
+    // sleep(5) with 2s timeout: client disconnects at 2s, kernel finishes at 5s.
+    await expect(
+      sandbox.runCode('import time; time.sleep(5)', {
+        context,
+        timeoutMs: 2_000,
+      })
+    ).rejects.toThrow()
+
+    // With the fix (lock released after send), this sends immediately and
+    // succeeds once the kernel finishes the sleep. Without the fix, this
+    // blocks on the server lock indefinitely.
+    const result = await sandbox.runCode('x = 1; x', {
+      context,
+      timeoutMs: 15_000,
+    })
+    expect(result.text).toEqual('1')
+  },
+)

python/tests/async/test_async_timeout.py

@@ -0,0 +1,26 @@
+from e2b import TimeoutException
+from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
+
+import pytest
+
+
+async def test_execution_after_timeout_is_not_blocked(async_sandbox: AsyncSandbox):
+    """After a client-side timeout, subsequent executions should not be blocked
+    behind an orphaned lock. Regression test for issue #213.
+
+    Uses its own context to avoid blocking other parallel tests with the sleep.
+    """
+
+    context = await async_sandbox.create_code_context()
+
+    # sleep(5) with 2s timeout: client disconnects at 2s, kernel finishes at 5s.
+    with pytest.raises(TimeoutException):
+        await async_sandbox.run_code(
+            "import time; time.sleep(5)", context=context, timeout=2
+        )
+
+    # With the fix (lock released after send), this sends immediately and
+    # succeeds once the kernel finishes the sleep. Without the fix, this
+    # blocks on the server lock indefinitely.
+    result = await async_sandbox.run_code("x = 1; x", context=context, timeout=15)
+    assert result.text == "1"

python/tests/sync/test_timeout.py

@@ -0,0 +1,24 @@
+from e2b import TimeoutException
+from e2b_code_interpreter.code_interpreter_sync import Sandbox
+
+import pytest
+
+
+def test_execution_after_timeout_is_not_blocked(sandbox: Sandbox):
+    """After a client-side timeout, subsequent executions should not be blocked
+    behind an orphaned lock. Regression test for issue #213.
+
+    Uses its own context to avoid blocking other parallel tests with the sleep.
+    """
+
+    context = sandbox.create_code_context()
+
+    # sleep(5) with 2s timeout: client disconnects at 2s, kernel finishes at 5s.
+    with pytest.raises(TimeoutException):
+        sandbox.run_code("import time; time.sleep(5)", context=context, timeout=2)
+
+    # With the fix (lock released after send), this sends immediately and
+    # succeeds once the kernel finishes the sleep. Without the fix, this
+    # blocks on the server lock indefinitely.
+    result = sandbox.run_code("x = 1; x", context=context, timeout=15)
+    assert result.text == "1"

template/server/messaging.py

@@ -294,6 +294,14 @@ async def execute(
         if self._ws is None:
             raise Exception("WebSocket not connected")
 
+        message_id = str(uuid.uuid4())
+        execution = Execution()
+
+        # 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():
@@ -327,8 +335,6 @@ async def execute(
                 )
                 complete_code = f"{indented_env_code}\n{complete_code}"
 
-            message_id = str(uuid.uuid4())
-            execution = Execution()
             self._executions[message_id] = execution
 
             # Send the code for execution
@@ -362,17 +368,19 @@ async def execute(
                 )
                 await execution.queue.put(UnexpectedEndOfExecution())
 
-            # Stream the results
+        # 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]
 
-            del self._executions[message_id]
-
-            # Clean up env vars in a separate request after the main code has run
-            if env_vars:
-                self._cleanup_task = asyncio.create_task(
-                    self._cleanup_env_vars(env_vars)
-                )
+        # Clean up env vars in a separate request after the main code has run
+        if env_vars:
+            self._cleanup_task = asyncio.create_task(
+                self._cleanup_env_vars(env_vars)
+            )
 
     async def _receive_message(self):
         if not self._ws:
@@ -385,8 +393,7 @@ 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, 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.
+            # To prevent infinite hang, cancel all ongoing executions as results may be lost during reconnect.
             for key, execution in self._executions.items():
                 await execution.queue.put(
                     Error(