From d59c748929666446daf9896f28ce6a082dfca305 Mon Sep 17 00:00:00 2001
From: Dominik Jain <dominik.jain@oraios-ai.de>
Date: Thu, 30 Oct 2025 17:05:15 +0100
Subject: [PATCH] Allow to cancel a task in the task queue via the Future,
 executing every task in a separate thread

---
 src/serena/agent.py              | 27 +++++++++----
 test/serena/test_serena_agent.py | 67 ++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 7 deletions(-)

diff --git a/src/serena/agent.py b/src/serena/agent.py
index e4e1199d..d9015f55 100644
--- a/src/serena/agent.py
+++ b/src/serena/agent.py
@@ -211,18 +211,31 @@ class SerenaAgent:
             """
             Execute the task, setting the result or exception on the future.
             """
-            try:
-                with LogTime(self.name, logger=log):
-                    result = self._function()
-                    self.future.set_result(result)
-            except Exception as e:
-                log.error(f"Error during execution of {self.name}: {e}", exc_info=e)
-                self.future.set_exception(e)
+
+            def run_task():
+                try:
+                    with LogTime(self.name, logger=log):
+                        result = self._function()
+                        if not self.future.done():
+                            self.future.set_result(result)
+                except Exception as e:
+                    if not self.future.done():
+                        log.error(f"Error during execution of {self.name}: {e}", exc_info=e)
+                        self.future.set_exception(e)
+
+            thread = Thread(target=run_task, name=self.name)
+            thread.start()
 
     def _process_task_queue(self) -> None:
         while True:
+            log.debug("Waiting for task ...")
             task: SerenaAgent.Task = self._task_executor_queue.get()
             task.execute()
+            try:
+                # wait for task completion
+                task.future.result()
+            except:
+                pass
 
     def get_language_server_manager(self) -> LanguageServerManager | None:
         if self._active_project is not None:
diff --git a/test/serena/test_serena_agent.py b/test/serena/test_serena_agent.py
index 882b45be..90c053a0 100644
--- a/test/serena/test_serena_agent.py
+++ b/test/serena/test_serena_agent.py
@@ -61,6 +61,14 @@ def serena_agent(request: pytest.FixtureRequest, serena_config):
     return SerenaAgent(project=project_name, serena_config=serena_config)
 
 
+@pytest.fixture
+def basic_serena_agent(serena_config):
+    """
+    Fixture for a basic SerenaAgent without a project
+    """
+    return SerenaAgent(project=None, serena_config=serena_config)
+
+
 class TestSerenaAgent:
     @pytest.mark.parametrize(
         "serena_agent,symbol_name,expected_kind,expected_file",
@@ -284,3 +292,62 @@ class TestSerenaAgent:
 
         symbols = json.loads(result)
         assert not symbols, f"Expected to find no symbols for {name_path}. Symbols found: {symbols}"
+
+    @staticmethod
+    def _create_task(delay: float, name: str, exception: bool = False):
+        def fn():
+            time.sleep(delay)
+            if exception:
+                raise ValueError(f"Task {name} failed")
+            return True
+
+        fn.__name__ = name
+        return fn
+
+    def test_task_queue_sequence(self, basic_serena_agent):
+        """
+        Tests that a sequence of tasks is executed correctly
+        """
+        future1 = basic_serena_agent.issue_task(self._create_task(1, "task1"))
+        future2 = basic_serena_agent.issue_task(self._create_task(1, "task2"))
+        assert future1.result() is True
+        assert future2.result() is True
+
+    def test_task_queue_exception(self, basic_serena_agent):
+        """
+        Tests that tasks that raise exceptions are handled correctly, i.e. that
+          * the exception is propagated,
+          * subsequent tasks are still executed.
+        """
+        future1 = basic_serena_agent.issue_task(self._create_task(1, "task1", exception=True))
+        future2 = basic_serena_agent.issue_task(self._create_task(1, "task2"))
+        have_exception = False
+        try:
+            assert future1.result()
+        except Exception as e:
+            assert isinstance(e, ValueError)
+            have_exception = True
+        assert have_exception
+        assert future2.result() is True
+
+    def test_task_queue_cancel(self, basic_serena_agent):
+        """
+        Tests that tasks that are cancelled are handled correctly, i.e. that
+          * subsequent tasks are executed as soon as cancellation ensues.
+          * the cancelled task raises CancelledError when result() is called.
+        """
+        start_time = time.time()
+        future1 = basic_serena_agent.issue_task(self._create_task(10, "task1"))
+        future2 = basic_serena_agent.issue_task(self._create_task(1, "task2"))
+        time.sleep(1)
+        future1.cancel()
+        assert future2.result() is True
+        end_time = time.time()
+        assert (end_time - start_time) < 9, "Cancelled task did not stop in time"
+        have_cancelled_error = False
+        try:
+            future1.result()
+        except Exception as e:
+            assert e.__class__.__name__ == "CancelledError"
+            have_cancelled_error = True
+        assert have_cancelled_error