From 04c9a680ead34a0948689e4410b3415b798b33d6 Mon Sep 17 00:00:00 2001 From: Dominik Jain Date: Thu, 30 Oct 2025 21:36:26 +0100 Subject: [PATCH] Improve timeout handling of tasks, * allowing a timeout to be passed at task construction * using the given timeout (if any) for the task queue itself and for result retrieval --- src/serena/agent.py | 17 ++++++----- src/serena/task_executor.py | 58 +++++++++++++++++++++++++------------ 2 files changed, 50 insertions(+), 25 deletions(-) diff --git a/src/serena/agent.py b/src/serena/agent.py index 78adcdc3..aecc21d3 100644 --- a/src/serena/agent.py +++ b/src/serena/agent.py @@ -8,9 +8,8 @@ import platform import sys import webbrowser from collections.abc import Callable -from concurrent.futures import Future from logging import Logger -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 @@ -371,7 +370,9 @@ 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, logged: bool = True) -> 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. @@ -379,11 +380,12 @@ class SerenaAgent: :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 - :return: a Future object representing the execution of the task + :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 """ - return self._task_executor.issue_task(task, name=name, logged=logged) + return self._task_executor.issue_task(task, name=name, logged=logged, timeout=timeout) - def execute_task(self, task: Callable[[], T], name: str | None = None, logged: bool = True) -> 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. @@ -391,9 +393,10 @@ class SerenaAgent: :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 """ - return self._task_executor.execute_task(task, name=name, logged=logged) + 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 index 9ea20139..bae3d325 100644 --- a/src/serena/task_executor.py +++ b/src/serena/task_executor.py @@ -1,7 +1,3 @@ -""" -The Serena Model Context Protocol (MCP) Server -""" - import concurrent.futures import threading import time @@ -9,7 +5,7 @@ from collections.abc import Callable from concurrent.futures import Future from dataclasses import dataclass from threading import Thread -from typing import Any, Optional, TypeVar +from typing import Generic, Optional, TypeVar from sensai.util import logging from sensai.util.logging import LogTime @@ -28,16 +24,18 @@ class TaskExecutor: self._task_executor_task_index = 1 self._task_executor_current_task: Optional[TaskExecutor.Task] = None - class Task(ToStringMixin): - def __init__(self, function: Callable[[], Any], name: str, logged: bool = True): + 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]: @@ -72,12 +70,34 @@ class TaskExecutor: """ 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, or None to wait indefinitely - :return: True if the task is done, False if the timeout was reached + :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) @@ -103,7 +123,7 @@ class TaskExecutor: task.start() # wait for task completion - task.wait_until_done() + task.wait_until_done(timeout=task.timeout) with self._task_executor_lock: self._task_executor_current_task = None @@ -143,7 +163,7 @@ class TaskExecutor: tasks.append(self.TaskInfo.from_task(task, False)) return tasks - def issue_task(self, task: Callable[[], Any], name: str | None = None, logged: bool = True) -> Future: + 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. @@ -151,7 +171,8 @@ class TaskExecutor: :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 - :return: a Future object representing the execution of the task + :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: @@ -162,11 +183,11 @@ class TaskExecutor: task_name = f"{task_prefix_name}:{name or task.__name__}" if logged: log.info(f"Scheduling {task_name}") - task = self.Task(function=task, name=task_name, logged=logged) - self._task_executor_queue.append(task) - return task.future + 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) -> 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. @@ -174,7 +195,8 @@ class TaskExecutor: :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, logged=logged) - return future.result() + task_obj = self.issue_task(task, name=name, logged=logged, timeout=timeout) + return task_obj.result()