Represent the agent's task queue explicitly

This commit is contained in:
Dominik Jain authored and Dominik Jain committed 2025-10-30 01:28:06 +01:00
1 parent eb4a740369
commit 6e26d08d5d
1 file changed
+34 -12
+34 -12
View File
@@ -2,6 +2,7 @@
The Serena Model Context Protocol (MCP) Server
"""
import concurrent.futures
import multiprocessing
import os
import platform
@@ -9,12 +10,15 @@ import sys
import threading
import webbrowser
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
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 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
@@ -153,7 +157,9 @@ class SerenaAgent:
# 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 = ThreadPoolExecutor(max_workers=1, thread_name_prefix="SerenaAgentExecutor")
self._task_executor_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
@@ -192,6 +198,29 @@ 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()
self._function = function
def _tostring_includes(self) -> list[str]:
return ["name"]
def execute(self):
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):
while True:
task: SerenaAgent.Task = self._task_executor_queue.get()
task.execute()
def get_language_server_manager(self) -> LanguageServerManager | None:
if self._active_project is not None:
return self._active_project.language_server_manager
@@ -376,17 +405,10 @@ class SerenaAgent:
with self._task_executor_lock:
task_name = f"Task-{self._task_executor_task_index}[{name or task.__name__}]"
self._task_executor_task_index += 1
def task_execution_wrapper() -> Any:
try:
with LogTime(task_name, logger=log):
return task()
except Exception as e:
log.error(f"Error during execution of {task_name}: {e}", exc_info=e)
raise
log.info(f"Scheduling {task_name}")
return self._task_executor.submit(task_execution_wrapper)
task = SerenaAgent.Task(function=task, name=task_name)
self._task_executor_queue.put(task)
return task.future
def execute_task(self, task: Callable[[], T], name: str | None = None) -> T:
"""