From 3fded9b6e0f853661cf96ce4bced3931edb7da9e Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Fri, 12 Jan 2024 18:23:15 +0800
Subject: [PATCH 1/2] fix timeout

---
 metagpt/actions/execute_code.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/metagpt/actions/execute_code.py b/metagpt/actions/execute_code.py
index 8355d3aca..c75711e75 100644
--- a/metagpt/actions/execute_code.py
+++ b/metagpt/actions/execute_code.py
@@ -4,6 +4,7 @@
 @Author  :   orange-crow
 @File    :   code_executor.py
 """
+import asyncio
 import re
 import traceback
 from abc import ABC, abstractmethod
@@ -81,6 +82,9 @@ class ExecutePyCode(ExecuteCode, Action):
     async def reset(self):
         """reset NotebookClient"""
         await self.terminate()
+
+        # sleep 1s to wait for the kernel to be cleaned up completely
+        await asyncio.sleep(1)
         await self.build()
         self.nb_client = NotebookClient(self.nb, timeout=self.timeout)
 
@@ -181,7 +185,11 @@ class ExecutePyCode(ExecuteCode, Action):
             await self.nb_client.async_execute_cell(cell, cell_index)
             return True, ""
         except CellTimeoutError:
-            return False, "TimeoutError"
+            assert self.nb_client.km is not None
+            await self.nb_client.km.interrupt_kernel()
+            await asyncio.sleep(1)
+            error_msg = "Cell execution timed out: Execution exceeded the time limit and was stopped; consider optimizing your code for better performance."
+            return False, error_msg
         except DeadKernelError:
             await self.reset()
             return False, "DeadKernelError"

From 932a26cbb3be1e12a52b6d4e1656b7c6edb0e35e Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Mon, 15 Jan 2024 10:50:08 +0800
Subject: [PATCH 2/2] update unittest

---
 tests/metagpt/actions/test_execute_code.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/metagpt/actions/test_execute_code.py b/tests/metagpt/actions/test_execute_code.py
index 8340272e4..904cc3c58 100644
--- a/tests/metagpt/actions/test_execute_code.py
+++ b/tests/metagpt/actions/test_execute_code.py
@@ -96,4 +96,4 @@ async def test_run_with_timeout():
     code = "import time; time.sleep(2)"
     message, success = await pi.run(code)
     assert not success
-    assert message == "TimeoutError"
+    assert message.startswith("Cell execution timed out")