diff --git a/src/serena/agent.py b/src/serena/agent.py index c919ada..8384896 100644 --- a/src/serena/agent.py +++ b/src/serena/agent.py @@ -45,6 +45,7 @@ from serena.util.file_system import scan_directory from serena.util.general import load_yaml, save_yaml from serena.util.inspection import determine_programming_language_composition, iter_subclasses from serena.util.shell import execute_shell_command +from serena.util.thread import ExecutionResult, execute_with_timeout if TYPE_CHECKING: from serena.gui_log_viewer import GuiLogViewerHandler @@ -53,6 +54,7 @@ log = logging.getLogger(__name__) LOG_FORMAT = "%(levelname)-5s %(asctime)-15s %(name)s:%(funcName)s:%(lineno)d - %(message)s" TTool = TypeVar("TTool", bound="Tool") SUCCESS_RESULT = "OK" +DEFAULT_TOOL_TIMEOUT: float = 240 def show_fatal_exception_safe(e: Exception) -> None: @@ -205,6 +207,7 @@ class SerenaConfigBase(ABC): gui_log_window_enabled: bool = False gui_log_window_level: int = logging.INFO web_dashboard: bool = True + tool_timeout: float = DEFAULT_TOOL_TIMEOUT @cached_property def project_paths(self) -> list[str]: @@ -353,6 +356,7 @@ class SerenaConfig(SerenaConfigBase): instance.gui_log_window_enabled = loaded_commented_yaml.get("gui_log_window", False) instance.gui_log_window_level = loaded_commented_yaml.get("gui_log_level", logging.INFO) instance.web_dashboard = loaded_commented_yaml.get("web_dashboard", True) + instance.tool_timeout = loaded_commented_yaml.get("tool_timeout", DEFAULT_TOOL_TIMEOUT) # re-save the configuration file if any migrations were performed if num_project_migrations > 0: @@ -968,8 +972,15 @@ class Tool(Component): if not self.agent.is_language_server_running(): log.info("Language server is not running. Starting it ...") self.agent.reset_language_server() - # apply the actual tool - result = apply_fn(**kwargs) + + # apply the actual tool with a timeout + execution_fn = lambda: apply_fn(**kwargs) + execution_result = execute_with_timeout(execution_fn, self.agent.serena_config.tool_timeout, self.get_name()) + if execution_result.status == ExecutionResult.Status.SUCCESS: + result = cast(str, execution_result.result_value) + else: + assert execution_result.exception is not None + raise execution_result.exception except Exception as e: if not catch_exceptions: diff --git a/src/serena/resources/serena_config.template.yml b/src/serena/resources/serena_config.template.yml index 33a9cd3..01bcdd7 100644 --- a/src/serena/resources/serena_config.template.yml +++ b/src/serena/resources/serena_config.template.yml @@ -23,6 +23,8 @@ web_dashboard: True gui_log_level: 20 # minimum log level for the GUI log window (10 = debug, 20 = info, 30 = warning, 40 = error) +tool_timeout: 240 +# timeout, in seconds, after which tool executions are terminated # MANAGED BY SERENA, KEEP AT THE BOTTOM OF THE YAML AND DON'T EDIT WITHOUT NEED # The list of registered projects. diff --git a/src/serena/util/thread.py b/src/serena/util/thread.py new file mode 100644 index 0000000..7b77918 --- /dev/null +++ b/src/serena/util/thread.py @@ -0,0 +1,69 @@ +import threading +from collections.abc import Callable +from enum import Enum +from typing import Generic, TypeVar + +from sensai.util.string import ToStringMixin + + +class TimeoutException(Exception): + def __init__(self, message: str, timeout: float) -> None: + super().__init__(message) + self.timeout = timeout + + +T = TypeVar("T") + + +class ExecutionResult(Generic[T], ToStringMixin): + + class Status(Enum): + SUCCESS = "success" + TIMEOUT = "timeout" + EXCEPTION = "error" + + def __init__(self) -> None: + self.result_value: T | None = None + self.status: ExecutionResult.Status | None = None + self.exception: Exception | None = None + + def set_result_value(self, value: T) -> None: + self.result_value = value + self.status = ExecutionResult.Status.SUCCESS + + def set_timed_out(self, exception: TimeoutException) -> None: + self.exception = exception + self.status = ExecutionResult.Status.TIMEOUT + + def set_exception(self, exception: Exception) -> None: + self.exception = exception + self.status = ExecutionResult.Status.EXCEPTION + + +def execute_with_timeout(func: Callable[[], T], timeout: float, function_name: str) -> ExecutionResult[T]: + """ + Executes the given function with a timeout + + :param func: the function to execute + :param timeout: the timeout in seconds + :param function_name: the name of the function (for error messages) + :returns: the execution result + """ + execution_result: ExecutionResult[T] = ExecutionResult() + + def target() -> None: + try: + value = func() + execution_result.set_result_value(value) + except Exception as e: + execution_result.set_exception(e) + + thread = threading.Thread(target=target, daemon=True) + thread.start() + thread.join(timeout=timeout) + + if thread.is_alive(): + timeout_exception = TimeoutException(f"Execution of '{function_name}' timed out after {timeout} seconds.", timeout) + execution_result.set_timed_out(timeout_exception) + + return execution_result