fix: handle aborted status

Author: oskrim

.changeset/handle-aborted-status.md

@@ -0,0 +1,5 @@
+---
+'@e2b/code-interpreter-template': patch
+---
+
+Handle aborted status

template/server/messaging.py

@@ -569,7 +569,7 @@ async def _process_message(self, data: dict):
                         traceback="".join(data["content"].get("traceback", [])),
                     )
                 )
-            elif data["content"]["status"] == "abort":
+            elif data["content"]["status"] in ("abort", "aborted"):
                 logger.debug(f"Execution {parent_msg_ig} was aborted")
                 await queue.put(
                     Error(

tests/test_interrupt.py

@@ -3,7 +3,7 @@
 import httpx
 import pytest
 
-from harness import CodeInterpreter
+from harness import CodeInterpreter, Execution
 
 
 def test_subsequent_execution_works_after_client_timeout(client: CodeInterpreter):
@@ -18,3 +18,21 @@ def test_subsequent_execution_works_after_client_timeout(client: CodeInterpreter
 
     execution = client.run_code("1 + 1", timeout=10)
     assert execution.text == "2"
+
+
+def test_execution_right_after_client_timeout_does_not_hang(client: CodeInterpreter):
+    with pytest.raises(httpx.TimeoutException):
+        client.run_code("import time; time.sleep(300)", timeout=3)
+
+    execution = Execution()
+    started = time.monotonic()
+    for event in client.stream_code("1 + 1", timeout=10):
+        execution.add(event)
+        if time.monotonic() - started > 15:
+            break
+
+    assert execution.completed
+    if execution.error is not None:
+        assert execution.error.name == "ExecutionAborted"
+    else:
+        assert execution.text == "2"