diff --git a/src/serena/agent.py b/src/serena/agent.py index e4e1199d..fc4ab315 100644 --- a/src/serena/agent.py +++ b/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: """ diff --git a/src/serena/task_executor.py b/src/serena/task_executor.py new file mode 100644 index 00000000..bae3d325 --- /dev/null +++ b/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() diff --git a/src/serena/tools/tools_base.py b/src/serena/tools/tools_base.py index c4b0f300..a79f2911 100644 --- a/src/serena/tools/tools_base.py +++ b/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: diff --git a/test/serena/test_task_executor.py b/test/serena/test_task_executor.py new file mode 100644 index 00000000..25c31c9d --- /dev/null +++ b/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"