diff --git a/.changeset/handle-aborted-status.md b/.changeset/handle-aborted-status.md new file mode 100644 index 00000000..7088dea0 --- /dev/null +++ b/.changeset/handle-aborted-status.md @@ -0,0 +1,5 @@ +--- +'@e2b/code-interpreter-template': patch +--- + +Handle aborted status diff --git a/template/server/messaging.py b/template/server/messaging.py index 7248a5b4..8653ed73 100644 --- a/template/server/messaging.py +++ b/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( diff --git a/tests/test_interrupt.py b/tests/test_interrupt.py index cda1de46..91c39914 100644 --- a/tests/test_interrupt.py +++ b/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"