Merge remote-tracking branch 'origin/cancelable-tasks' into dashboard_extesions

src/serena/agent.py

@@ -2,23 +2,17 @@
 The Serena Model Context Protocol (MCP) Server
 """
 
-import concurrent.futures
 import multiprocessing
 import os
 import platform
 import sys
-import threading
 import webbrowser
 from collections.abc import Callable
-from concurrent.futures import Future
 from logging import Logger
-from queue import Queue
-from threading import Thread
-from typing import TYPE_CHECKING, Any, Optional, TypeVar
+from typing import TYPE_CHECKING, Optional, TypeVar
 
 from sensai.util import logging
 from sensai.util.logging import LogTime
-from sensai.util.string import ToStringMixin
 
 from interprompt.jinja_template import JinjaTemplate
 from serena import serena_version
@@ -29,6 +23,7 @@ from serena.dashboard import SerenaDashboardAPI
 from serena.ls_manager import LanguageServerManager
 from serena.project import Project
 from serena.prompt_factory import SerenaPromptFactory
+from serena.task_executor import TaskExecutor
 from serena.tools import ActivateProjectTool, GetCurrentConfigTool, Tool, ToolMarker, ToolRegistry
 from serena.util.inspection import iter_subclasses
 from serena.util.logging import MemoryLogHandler
@@ -156,12 +151,8 @@ class SerenaAgent:
         log.info(f"Number of exposed tools: {len(self._exposed_tools)}")
 
         # create executor for starting the language server and running tools in another thread
-        # This executor is used to achieve linear task execution, so it is important to use a single-threaded executor.
-        self._task_executor_queue: Queue = Queue()
-        self._task_executor_thread = Thread(target=self._process_task_queue, name="SerenaAgentTaskExecutor", daemon=True)
-        self._task_executor_thread.start()
-        self._task_executor_lock = threading.Lock()
-        self._task_executor_task_index = 1
+        # This executor is used to achieve linear task execution
+        self._task_executor = TaskExecutor("SerenaAgentTaskExecutor")
 
         # Initialize the prompt factory
         self.prompt_factory = SerenaPromptFactory()
@@ -198,31 +189,14 @@ class SerenaAgent:
                 process.start()
                 process.join(timeout=1)
 
-    class Task(ToStringMixin):
-        def __init__(self, function: Callable[[], Any], name: str):
-            self.name = name
-            self.future: concurrent.futures.Future = concurrent.futures.Future()
-            self._function = function
+    def get_current_tasks(self) -> list[TaskExecutor.TaskInfo]:
+        """
+        Gets the list of tasks currently running or queued for execution.
+        The function returns a list of thread-safe TaskInfo objects (specifically created for the caller).
 
-        def _tostring_includes(self) -> list[str]:
-            return ["name"]
-
-        def execute(self) -> None:
-            """
-            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 _process_task_queue(self) -> None:
-        while True:
-            task: SerenaAgent.Task = self._task_executor_queue.get()
-            task.execute()
+        :return: the list of tasks in the execution order (running task first)
+        """
+        return self._task_executor.get_current_tasks()
 
     def get_language_server_manager(self) -> LanguageServerManager | None:
         if self._active_project is not None:
@@ -396,33 +370,33 @@ class SerenaAgent:
 
         log.info(f"Active tools ({len(self._active_tools)}): {', '.join(self.get_active_tool_names())}")
 
-    def issue_task(self, task: Callable[[], Any], name: str | None = None) -> Future:
+    def issue_task(
+        self, task: Callable[[], T], name: str | None = None, logged: bool = True, timeout: float | None = None
+    ) -> TaskExecutor.Task[T]:
         """
         Issue a task to the executor for asynchronous execution.
         It is ensured that tasks are executed in the order they are issued, one after another.
 
         :param task: the task to execute
         :param name: the name of the task for logging purposes; if None, use the task function's name
-        :return: a Future object representing the execution of the task
+        :param logged: whether to log management of the task; if False, only errors will be logged
+        :param timeout: the maximum time to wait for task completion in seconds, or None to wait indefinitely
+        :return: the task object, through which the task's future result can be accessed
         """
-        with self._task_executor_lock:
-            task_name = f"Task-{self._task_executor_task_index}[{name or task.__name__}]"
-            self._task_executor_task_index += 1
-            log.info(f"Scheduling {task_name}")
-            task = SerenaAgent.Task(function=task, name=task_name)
-            self._task_executor_queue.put(task)
-            return task.future
+        return self._task_executor.issue_task(task, name=name, logged=logged, timeout=timeout)
 
-    def execute_task(self, task: Callable[[], T], name: str | None = None) -> T:
+    def execute_task(self, task: Callable[[], T], name: str | None = None, logged: bool = True, timeout: float | None = None) -> T:
         """
         Executes the given task synchronously via the agent's task executor.
         This is useful for tasks that need to be executed immediately and whose results are needed right away.
 
         :param task: the task to execute
+        :param name: the name of the task for logging purposes; if None, use the task function's name
+        :param logged: whether to log management of the task; if False, only errors will be logged
+        :param timeout: the maximum time to wait for task completion in seconds, or None to wait indefinitely
         :return: the result of the task execution
         """
-        future = self.issue_task(task, name=name)
-        return future.result()
+        return self._task_executor.execute_task(task, name=name, logged=logged, timeout=timeout)
 
     def is_using_language_server(self) -> bool:
         """

src/serena/task_executor.py

@@ -0,0 +1,202 @@
+import concurrent.futures
+import threading
+import time
+from collections.abc import Callable
+from concurrent.futures import Future
+from dataclasses import dataclass
+from threading import Thread
+from typing import Generic, Optional, TypeVar
+
+from sensai.util import logging
+from sensai.util.logging import LogTime
+from sensai.util.string import ToStringMixin
+
+log = logging.getLogger(__name__)
+T = TypeVar("T")
+
+
+class TaskExecutor:
+    def __init__(self, name: str):
+        self._task_executor_lock = threading.Lock()
+        self._task_executor_queue: list[TaskExecutor.Task] = []
+        self._task_executor_thread = Thread(target=self._process_task_queue, name=name, daemon=True)
+        self._task_executor_thread.start()
+        self._task_executor_task_index = 1
+        self._task_executor_current_task: Optional[TaskExecutor.Task] = None
+
+    class Task(ToStringMixin, Generic[T]):
+        def __init__(self, function: Callable[[], T], name: str, logged: bool = True, timeout: Optional[float] = None):
+            """
+            :param function: the function representing the task to execute
+            :param name: the name of the task
+            :param logged: whether to log management of the task; if False, only errors will be logged
+            :param timeout: the maximum time to wait for task completion in seconds, or None to wait indefinitely
+            """
+            self.name = name
+            self.future: concurrent.futures.Future = concurrent.futures.Future()
+            self.logged = logged
+            self.timeout = timeout
+            self._function = function
+
+        def _tostring_includes(self) -> list[str]:
+            return ["name"]
+
+        def start(self) -> None:
+            """
+            Executes the task in a separate thread, setting the result or exception on the future.
+            """
+
+            def run_task() -> None:
+                try:
+                    if self.future.done():
+                        if self.logged:
+                            log.info(f"Task {self.name} was already completed/cancelled; skipping execution")
+                        return
+                    with LogTime(self.name, logger=log, enabled=self.logged):
+                        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 is_done(self) -> bool:
+            """
+            :return: whether the task has completed (either successfully, with failure, or via cancellation)
+            """
+            return self.future.done()
+
+        def result(self, timeout: float | None = None) -> T:
+            """
+            Blocks until the task is done or the timeout is reached, and returns the result.
+            If an exception occurred during task execution, it is raised here.
+            If the timeout is reached, a TimeoutError is raised (but the task is not cancelled).
+            If the task is cancelled, a CancelledError is raised.
+
+            :param timeout: the maximum time to wait in seconds; if None, use the task's own timeout
+                (which may be None to wait indefinitely)
+            :return: True if the task is done, False if the timeout was reached
+            """
+            return self.future.result(timeout=timeout)
+
+        def cancel(self) -> None:
+            """
+            Cancels the task. If it has not yet started, it will not be executed.
+            If it has already started, its future will be marked as cancelled and will raise a CancelledError
+            when its result is requested.
+            """
+            self.future.cancel()
+
+        def wait_until_done(self, timeout: float | None = None) -> None:
+            """
+            Waits until the task is done or the timeout is reached.
+            The task is done if it either completed successfully, failed with an exception, or was cancelled.
+
+            :param timeout: the maximum time to wait in seconds; if None, use the task's own timeout
+                (which may be None to wait indefinitely)
+            """
+            try:
+                self.future.result(timeout=timeout)
+            except:
+                pass
+
+    def _process_task_queue(self) -> None:
+        while True:
+            # obtain task from the queue
+            task: TaskExecutor.Task | None = None
+            with self._task_executor_lock:
+                if len(self._task_executor_queue) > 0:
+                    task = self._task_executor_queue.pop(0)
+            if task is None:
+                time.sleep(0.1)
+                continue
+
+            # start task execution asynchronously
+            with self._task_executor_lock:
+                self._task_executor_current_task = task
+            if task.logged:
+                log.info("Starting execution of %s", task.name)
+            task.start()
+
+            # wait for task completion
+            task.wait_until_done(timeout=task.timeout)
+            with self._task_executor_lock:
+                self._task_executor_current_task = None
+
+    @dataclass
+    class TaskInfo:
+        name: str
+        is_running: bool
+        future: Future
+        """
+        future for accessing the task's result
+        """
+        task_id: int
+        """
+        unique identifier of the task
+        """
+
+        @staticmethod
+        def from_task(task: "TaskExecutor.Task", is_running: bool) -> "TaskExecutor.TaskInfo":
+            return TaskExecutor.TaskInfo(name=task.name, is_running=is_running, future=task.future, task_id=id(task))
+
+        def cancel(self) -> None:
+            self.future.cancel()
+
+    def get_current_tasks(self) -> list["TaskExecutor.TaskInfo"]:
+        """
+        Gets the list of tasks currently running or queued for execution.
+        The function returns a list of thread-safe TaskInfo objects (specifically created for the caller).
+
+        :return: the list of tasks in the execution order (running task first)
+        """
+        tasks = []
+        with self._task_executor_lock:
+            if self._task_executor_current_task is not None:
+                tasks.append(self.TaskInfo.from_task(self._task_executor_current_task, True))
+            for task in self._task_executor_queue:
+                if not task.is_done():
+                    tasks.append(self.TaskInfo.from_task(task, False))
+        return tasks
+
+    def issue_task(self, task: Callable[[], T], name: str | None = None, logged: bool = True, timeout: float | None = None) -> Task[T]:
+        """
+        Issue a task to the executor for asynchronous execution.
+        It is ensured that tasks are executed in the order they are issued, one after another.
+
+        :param task: the task to execute
+        :param name: the name of the task for logging purposes; if None, use the task function's name
+        :param logged: whether to log management of the task; if False, only errors will be logged
+        :param timeout: the maximum time to wait for task completion in seconds, or None to wait indefinitely
+        :return: the task object, through which the task's future result can be accessed
+        """
+        with self._task_executor_lock:
+            if logged:
+                task_prefix_name = f"Task-{self._task_executor_task_index}"
+                self._task_executor_task_index += 1
+            else:
+                task_prefix_name = "BackgroundTask"
+            task_name = f"{task_prefix_name}:{name or task.__name__}"
+            if logged:
+                log.info(f"Scheduling {task_name}")
+            task_obj = self.Task(function=task, name=task_name, logged=logged, timeout=timeout)
+            self._task_executor_queue.append(task_obj)
+            return task_obj
+
+    def execute_task(self, task: Callable[[], T], name: str | None = None, logged: bool = True, timeout: float | None = None) -> T:
+        """
+        Executes the given task synchronously via the agent's task executor.
+        This is useful for tasks that need to be executed immediately and whose results are needed right away.
+
+        :param task: the task to execute
+        :param name: the name of the task for logging purposes; if None, use the task function's name
+        :param logged: whether to log management of the task; if False, only errors will be logged
+        :param timeout: the maximum time to wait for task completion in seconds, or None to wait indefinitely
+        :return: the result of the task execution
+        """
+        task_obj = self.issue_task(task, name=name, logged=logged, timeout=timeout)
+        return task_obj.result()

src/serena/tools/tools_base.py

@@ -286,8 +286,14 @@ class Tool(Component):
 
             return result
 
-        future = self.agent.issue_task(task, name=self.__class__.__name__)
-        return future.result(timeout=self.agent.serena_config.tool_timeout)
+        # execute the tool in the agent's task executor, with timeout
+        try:
+            task_exec = self.agent.issue_task(task, name=self.__class__.__name__)
+            return task_exec.result(timeout=self.agent.serena_config.tool_timeout)
+        except Exception as e:  # typically TimeoutError (other exceptions caught in task)
+            msg = f"Error: {e.__class__.__name__} - {e}"
+            log.error(msg)
+            return msg
 
 
 class EditedFileContext:

test/serena/test_task_executor.py

@@ -0,0 +1,119 @@
+import time
+
+import pytest
+
+from serena.task_executor import TaskExecutor
+
+
+@pytest.fixture
+def executor():
+    """
+    Fixture for a basic SerenaAgent without a project
+    """
+    return TaskExecutor("TestExecutor")
+
+
+class Task:
+    def __init__(self, delay: float, exception: bool = False):
+        self.delay = delay
+        self.exception = exception
+        self.did_run = False
+
+    def run(self):
+        self.did_run = True
+        time.sleep(self.delay)
+        if self.exception:
+            raise ValueError("Task failed")
+        return True
+
+
+def test_task_executor_sequence(executor):
+    """
+    Tests that a sequence of tasks is executed correctly
+    """
+    future1 = executor.issue_task(Task(1).run, name="task1")
+    future2 = executor.issue_task(Task(1).run, name="task2")
+    assert future1.result() is True
+    assert future2.result() is True
+
+
+def test_task_executor_exception(executor):
+    """
+    Tests that tasks that raise exceptions are handled correctly, i.e. that
+      * the exception is propagated,
+      * subsequent tasks are still executed.
+    """
+    future1 = executor.issue_task(Task(1, exception=True).run, name="task1")
+    future2 = executor.issue_task(Task(1).run, name="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_executor_cancel_current(executor):
+    """
+    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 = executor.issue_task(Task(10).run, name="task1")
+    future2 = executor.issue_task(Task(1).run, name="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
+
+
+def test_task_executor_cancel_future(executor):
+    """
+    Tests that when a future task is cancelled, it is never run at all
+    """
+    task1 = Task(10)
+    task2 = Task(1)
+    future1 = executor.issue_task(task1.run, name="task1")
+    future2 = executor.issue_task(task2.run, name="task2")
+    time.sleep(1)
+    future2.cancel()
+    future1.cancel()
+    try:
+        future2.result()
+    except:
+        pass
+    assert task1.did_run
+    assert not task2.did_run
+
+
+def test_task_executor_cancellation_via_task_info(executor):
+    start_time = time.time()
+    executor.issue_task(Task(10).run, "task1")
+    executor.issue_task(Task(1).run, "task2")
+    task_infos = executor.get_current_tasks()
+    task_infos2 = executor.get_current_tasks()
+
+    # test expected tasks
+    assert len(task_infos) == 2
+    assert "task1" in task_infos[0].name
+    assert "task2" in task_infos[1].name
+
+    # test task identifiers being stable
+    assert task_infos2[0].task_id == task_infos[0].task_id
+
+    # test cancellation
+    task_infos[0].cancel()
+    assert task_infos[1].future.result()
+    end_time = time.time()
+    assert (end_time - start_time) < 9, "Cancelled task did not stop in time"