From bc6b0f2ebaffe6b5b0866a5d5be1dbba67d9ffff Mon Sep 17 00:00:00 2001 From: Michael Panchenko Date: Sun, 15 Jun 2025 17:57:36 +0200 Subject: [PATCH] Replaced queue-based approach to worker with pipe-based one --- src/serena/process_isolated_agent.py | 166 +++++++++++++-------------- 1 file changed, 77 insertions(+), 89 deletions(-) diff --git a/src/serena/process_isolated_agent.py b/src/serena/process_isolated_agent.py index d8e3ec5..be6234b 100644 --- a/src/serena/process_isolated_agent.py +++ b/src/serena/process_isolated_agent.py @@ -9,12 +9,10 @@ This module provides: import logging import multiprocessing -import queue -import threading -import time import traceback from collections.abc import Callable from enum import StrEnum +from multiprocessing.connection import Connection from typing import Any, Literal, Self from mcp.server.fastmcp.utilities.func_metadata import FuncMetadata @@ -28,37 +26,40 @@ log = logging.getLogger(__name__) class SerenaAgentWorker: """Worker process that hosts the actual SerenaAgent.""" - def __init__(self, request_queue: multiprocessing.Queue, response_queue: multiprocessing.Queue): - self.request_queue = request_queue - self.response_queue = response_queue + def __init__(self, conn: Connection): + self.conn = conn self.agent: SerenaAgent | None = None - self.running = True def run(self) -> None: """Main worker loop - runs in separate process.""" try: log.info("SerenaAgent worker process started") - while self.running: + while True: try: - # Wait for requests with timeout to allow clean shutdown - request = self.request_queue.get(timeout=1.0) + # Block until request received or connection closed + request = self.conn.recv() if request is None: # Shutdown signal break response = self._handle_request(request) - self.response_queue.put(response) + self.conn.send(response) - except queue.Empty: - continue + except EOFError: + # Connection closed - parent process terminated + log.info("Connection closed, worker shutting down") + break except Exception as e: log.error(f"Error in worker process: {e}") - response = { - "id": getattr(request, "id", None), - "error": f"Worker process error: {e!s}", - "traceback": traceback.format_exc(), - } - self.response_queue.put(response) + try: + response = { + "error": f"Worker process error: {e!s}", + "traceback": traceback.format_exc(), + } + self.conn.send(response) + except (EOFError, BrokenPipeError): + # Connection is broken, can't send error response + break except Exception as e: log.error(f"Fatal error in worker process: {e}") @@ -71,28 +72,27 @@ class SerenaAgentWorker: try: method = request["method"] params = request.get("params", {}) - request_id = request["id"] match method: case self.RequestMethod.INITIALIZE: - return self._initialize(request_id, params) + return self._initialize(params) case self.RequestMethod.TOOL_CALL: - return self._tool_call(request_id, params) + return self._tool_call(params) case self.RequestMethod.GET_ACTIVE_TOOL_NAMES: - return self._get_active_tool_names(request_id) + return self._get_active_tool_names() case self.RequestMethod.IS_LANGUAGE_SERVER_RUNNING: - return self._is_language_server_running(request_id) + return self._is_language_server_running() case self.RequestMethod.RESET_LANGUAGE_SERVER: - return self._reset_language_server(request_id) + return self._reset_language_server() case self.RequestMethod.GET_EXPOSED_TOOL_NAMES: - return self._get_exposed_tool_names(request_id) + return self._get_exposed_tool_names() case _: - return {"id": request_id, "error": f"Unknown method: {method}"} + return {"error": f"Unknown method: {method}"} except Exception as e: - return {"id": request.get("id"), "error": str(e), "traceback": traceback.format_exc()} + return {"error": str(e), "traceback": traceback.format_exc()} - def _initialize(self, request_id: str, params: dict[str, Any]) -> dict[str, Any]: + def _initialize(self, params: dict[str, Any]) -> dict[str, Any]: """Initialize the SerenaAgent.""" try: # Extract all possible initialization parameters @@ -127,14 +127,14 @@ class SerenaAgentWorker: tool_timeout=tool_timeout, ) - return {"id": request_id, "result": "SerenaAgent initialized successfully"} + return {"result": "SerenaAgent initialized successfully"} except Exception as e: - return {"id": request_id, "error": f"Failed to initialize SerenaAgent: {e!s}", "traceback": traceback.format_exc()} + return {"error": f"Failed to initialize SerenaAgent: {e!s}", "traceback": traceback.format_exc()} - def _tool_call(self, request_id: str, params: dict[str, Any]) -> dict[str, Any]: + def _tool_call(self, params: dict[str, Any]) -> dict[str, Any]: """Execute a tool call.""" if self.agent is None: - return {"id": request_id, "error": "SerenaAgent not initialized"} + return {"error": "SerenaAgent not initialized"} try: tool_name = params["tool_name"] @@ -148,61 +148,61 @@ class SerenaAgentWorker: break if tool is None: - return {"id": request_id, "error": f"Tool '{tool_name}' not found or not active"} + return {"error": f"Tool '{tool_name}' not found or not active"} # Execute the tool result = tool.apply_ex(**tool_params) - return {"id": request_id, "result": result} + return {"result": result} except Exception as e: - return {"id": request_id, "error": str(e), "traceback": traceback.format_exc()} + return {"error": str(e), "traceback": traceback.format_exc()} - def _get_active_tool_names(self, request_id: str) -> dict[str, Any]: + def _get_active_tool_names(self) -> dict[str, Any]: """Get list of active tool names.""" if self.agent is None: - return {"id": request_id, "error": "SerenaAgent not initialized"} + return {"error": "SerenaAgent not initialized"} try: tool_names = self.agent.get_active_tool_names() - return {"id": request_id, "result": tool_names} + return {"result": tool_names} except Exception as e: - return {"id": request_id, "error": str(e), "traceback": traceback.format_exc()} + return {"error": str(e), "traceback": traceback.format_exc()} - def _is_language_server_running(self, request_id: str) -> dict[str, Any]: + def _is_language_server_running(self) -> dict[str, Any]: """Check if language server is running.""" if self.agent is None: - return {"id": request_id, "error": "SerenaAgent not initialized"} + return {"error": "SerenaAgent not initialized"} try: is_running = self.agent.is_language_server_running() - return {"id": request_id, "result": is_running} + return {"result": is_running} except Exception as e: - return {"id": request_id, "error": str(e), "traceback": traceback.format_exc()} + return {"error": str(e), "traceback": traceback.format_exc()} - def _reset_language_server(self, request_id: str) -> dict[str, Any]: + def _reset_language_server(self) -> dict[str, Any]: """Reset the language server.""" if self.agent is None: - return {"id": request_id, "error": "SerenaAgent not initialized"} + return {"error": "SerenaAgent not initialized"} try: self.agent.reset_language_server() - return {"id": request_id, "result": "Language server reset successfully"} + return {"result": "Language server reset successfully"} except Exception as e: - return {"id": request_id, "error": str(e), "traceback": traceback.format_exc()} + return {"error": str(e), "traceback": traceback.format_exc()} - def _get_exposed_tool_names(self, request_id: str) -> dict[str, Any]: + def _get_exposed_tool_names(self) -> dict[str, Any]: """Get exposed tool names for MCP tool creation.""" if self.agent is None: - return {"id": request_id, "error": "SerenaAgent not initialized"} + return {"error": "SerenaAgent not initialized"} try: tool_instances = self.agent.get_exposed_tool_instances() # Return only tool names - metadata will be reconstructed from ToolRegistry tool_names = [tool.get_name_from_cls() for tool in tool_instances] - return {"id": request_id, "result": tool_names} + return {"result": tool_names} except Exception as e: - return {"id": request_id, "error": str(e), "traceback": traceback.format_exc()} + return {"error": str(e), "traceback": traceback.format_exc()} def _cleanup(self) -> None: """Clean up resources.""" @@ -215,7 +215,7 @@ class SerenaAgentWorker: self.agent = None class RequestMethod(StrEnum): - """Enumeration of available request methods for type-safe IPC.""" + """Enumeration of available request methods.""" INITIALIZE = "initialize" TOOL_CALL = "tool_call" @@ -262,10 +262,7 @@ class ProcessIsolatedSerenaAgent: # Use the default config loader self.serena_config = SerenaConfig.from_config_file() self.process: multiprocessing.Process | None = None - self.request_queue: multiprocessing.Queue | None = None - self.response_queue: multiprocessing.Queue | None = None - self.request_id_counter = 0 - self._lock = threading.Lock() + self.conn: Connection | None = None def start(self) -> None: """Start the worker process.""" @@ -274,12 +271,12 @@ class ProcessIsolatedSerenaAgent: log.info("Starting process-isolated SerenaAgent") - # Create communication queues - self.request_queue = multiprocessing.Queue() - self.response_queue = multiprocessing.Queue() + # Create communication pipe + parent_conn, child_conn = multiprocessing.Pipe() + self.conn = parent_conn # Create and start worker process - worker = SerenaAgentWorker(self.request_queue, self.response_queue) + worker = SerenaAgentWorker(child_conn) self.process = multiprocessing.Process(target=worker.run) self.process.start() @@ -313,9 +310,9 @@ class ProcessIsolatedSerenaAgent: log.info("Stopping process-isolated SerenaAgent") try: - # Send shutdown signal - if self.request_queue is not None: - self.request_queue.put(None) + # Close connection to signal worker to shutdown + if self.conn is not None: + self.conn.close() # Wait for process to terminate if self.process.is_alive(): @@ -336,8 +333,7 @@ class ProcessIsolatedSerenaAgent: finally: self.process = None - self.request_queue = None - self.response_queue = None + self.conn = None log.info("Process-isolated SerenaAgent stopped") @@ -346,35 +342,27 @@ class ProcessIsolatedSerenaAgent: if self.process is None or not self.process.is_alive(): raise RuntimeError("Worker process is not running") - with self._lock: - self.request_id_counter += 1 - request_id = str(self.request_id_counter) + if self.conn is None: + raise RuntimeError("Connection is not initialized") - request = {"id": request_id, "method": method, "params": params or {}} + request = {"method": method, "params": params or {}} # Send request - if self.request_queue is not None: - self.request_queue.put(request) + try: + self.conn.send(request) + except (EOFError, BrokenPipeError) as e: + raise RuntimeError("Failed to send request: worker process may have crashed") from e - # Wait for response + # Wait for response with timeout timeout = self.serena_config.tool_timeout - start_time = time.time() - while time.time() - start_time < timeout: + if self.conn.poll(timeout): try: - if self.response_queue is not None: - response = self.response_queue.get(timeout=1.0) - else: - raise RuntimeError("Response queue is not initialized") - if response["id"] == request_id: - return response - else: - # Put back response for different request - if self.response_queue is not None: - self.response_queue.put(response) - except queue.Empty: - continue - - raise TimeoutError(f"Request {request_id} timed out after {timeout} seconds") + response = self.conn.recv() + return response + except (EOFError, BrokenPipeError) as e: + raise RuntimeError("Failed to receive response: worker process may have crashed") from e + else: + raise TimeoutError(f"Request {method} timed out after {timeout} seconds") def _make_request_with_result(self, method: SerenaAgentWorker.RequestMethod, params: dict[str, Any] | None = None) -> Any: """Make a request and return the result, raising an exception if there's an error."""