From 343c00f6052a228c970e69d8b0c7b4d9ec2ffd2e Mon Sep 17 00:00:00 2001 From: Michael Panchenko Date: Sun, 15 Jun 2025 18:31:17 +0200 Subject: [PATCH] Fix process termination issues and allow terminating through dashboard --- src/serena/dashboard.py | 22 +++++-- src/serena/mcp.py | 21 +++++++ src/serena/process_isolated_agent.py | 61 +++++++++++++++++++ .../test_make_tool_process_isolation.py | 2 +- 4 files changed, 101 insertions(+), 5 deletions(-) diff --git a/src/serena/dashboard.py b/src/serena/dashboard.py index 622bef1..549c003 100644 --- a/src/serena/dashboard.py +++ b/src/serena/dashboard.py @@ -3,6 +3,7 @@ import queue import socket import sys import threading +from collections.abc import Callable import uvicorn from fastapi import FastAPI @@ -68,9 +69,12 @@ class ResponseToolNames(BaseModel): class SerenaDashboardAPI: log = logging.getLogger(__qualname__) - def __init__(self, memory_log_handler: MemoryLogHandler, tool_names: list[str]) -> None: + def __init__( + self, memory_log_handler: MemoryLogHandler, tool_names: list[str], shutdown_callback: Callable[[], None] | None = None + ) -> None: self._memory_log_handler = memory_log_handler self._tool_names = tool_names + self._shutdown_callback = shutdown_callback self._app = FastAPI(title="Serena Dashboard") self._setup_routes() @@ -92,9 +96,19 @@ class SerenaDashboardAPI: async def _shutdown(self) -> None: print("Shutdown initiated by dashbaord ...", file=sys.stderr) log.info("Shutting down Serena") - # noinspection PyUnresolvedReferences - # noinspection PyProtectedMember - os._exit(0) + if self._shutdown_callback: + self._shutdown_callback() + else: + # Try to use the global shutdown function from process_isolated_agent + try: + from serena.process_isolated_agent import call_global_shutdown + + call_global_shutdown() + except ImportError: + # Fallback to the old behavior if not in process-isolated mode + # noinspection PyUnresolvedReferences + # noinspection PyProtectedMember + os._exit(0) @staticmethod def _find_first_free_port(start_port: int) -> int: diff --git a/src/serena/mcp.py b/src/serena/mcp.py index e8f4ec6..a5ade13 100644 --- a/src/serena/mcp.py +++ b/src/serena/mcp.py @@ -183,14 +183,35 @@ def create_mcp_server_and_agent( async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[None]: """Manage server startup and shutdown lifecycle.""" nonlocal process_agent + import asyncio + mark_used(mcp_server) + async def monitor_worker_process() -> None: + """Monitor the worker process and shutdown server if it exits.""" + while True: + if process_agent.process is None or not process_agent.process.is_alive(): + log.info("Worker process has exited, shutting down MCP server") + # Trigger server shutdown + import os + + os._exit(0) + await asyncio.sleep(1) + try: process_agent.start() # Update tools now that the process agent is running update_tools() + # Start monitoring task + monitor_task = asyncio.create_task(monitor_worker_process()) yield finally: + if "monitor_task" in locals(): + monitor_task.cancel() + import contextlib + + with contextlib.suppress(asyncio.CancelledError): + await monitor_task process_agent.stop() mcp_settings = Settings(lifespan=server_lifespan, host=host, port=port) diff --git a/src/serena/process_isolated_agent.py b/src/serena/process_isolated_agent.py index be6234b..6ff4468 100644 --- a/src/serena/process_isolated_agent.py +++ b/src/serena/process_isolated_agent.py @@ -23,6 +23,36 @@ from serena.config import SerenaAgentContext, SerenaAgentMode log = logging.getLogger(__name__) +# Global shutdown function that can be called by the dashboard +class _GlobalShutdownRegistry: + """Registry for global shutdown function.""" + + def __init__(self) -> None: + self.shutdown_func: Callable[[], None] | None = None + + def set_shutdown_func(self, func: Callable[[], None]) -> None: + """Set the global shutdown function.""" + self.shutdown_func = func + + def call_shutdown(self) -> None: + """Call the global shutdown function if it exists.""" + if self.shutdown_func: + self.shutdown_func() + + +_shutdown_registry = _GlobalShutdownRegistry() + + +def set_global_shutdown_func(func: Callable[[], None]) -> None: + """Set the global shutdown function.""" + _shutdown_registry.set_shutdown_func(func) + + +def call_global_shutdown() -> None: + """Call the global shutdown function if it exists.""" + _shutdown_registry.call_shutdown() + + class SerenaAgentWorker: """Worker process that hosts the actual SerenaAgent.""" @@ -86,6 +116,8 @@ class SerenaAgentWorker: return self._reset_language_server() case self.RequestMethod.GET_EXPOSED_TOOL_NAMES: return self._get_exposed_tool_names() + case self.RequestMethod.SHUTDOWN: + return self._shutdown() case _: return {"error": f"Unknown method: {method}"} @@ -127,6 +159,16 @@ class SerenaAgentWorker: tool_timeout=tool_timeout, ) + # Set up global shutdown function for dashboard + def shutdown_worker() -> None: + log.info("Global shutdown function called") + self._cleanup() + import sys + + sys.exit(0) + + set_global_shutdown_func(shutdown_worker) + return {"result": "SerenaAgent initialized successfully"} except Exception as e: return {"error": f"Failed to initialize SerenaAgent: {e!s}", "traceback": traceback.format_exc()} @@ -204,6 +246,20 @@ class SerenaAgentWorker: except Exception as e: return {"error": str(e), "traceback": traceback.format_exc()} + def _shutdown(self) -> dict[str, Any]: + """Handle shutdown request from dashboard.""" + try: + log.info("Shutdown requested from dashboard") + # Clean up resources before exiting + self._cleanup() + # Exit the worker process - this will cause the main process to detect the termination + # and shut down gracefully through the server lifespan manager + import sys + + sys.exit(0) + except Exception as e: + return {"error": str(e), "traceback": traceback.format_exc()} + def _cleanup(self) -> None: """Clean up resources.""" if self.agent is not None: @@ -223,6 +279,7 @@ class SerenaAgentWorker: IS_LANGUAGE_SERVER_RUNNING = "is_language_server_running" RESET_LANGUAGE_SERVER = "reset_language_server" GET_EXPOSED_TOOL_NAMES = "get_exposed_tool_names" + SHUTDOWN = "shutdown" class ProcessIsolatedSerenaAgent: @@ -394,6 +451,10 @@ class ProcessIsolatedSerenaAgent: """Reset the language server.""" self._make_request_with_result(SerenaAgentWorker.RequestMethod.RESET_LANGUAGE_SERVER) + def shutdown_from_dashboard(self) -> None: + """Request shutdown from dashboard.""" + self._make_request_with_result(SerenaAgentWorker.RequestMethod.SHUTDOWN) + def get_exposed_tool_names(self) -> list[str]: """Get tool names for MCP tool creation.""" return self._make_request_with_result(SerenaAgentWorker.RequestMethod.GET_EXPOSED_TOOL_NAMES) diff --git a/test/serena/test_make_tool_process_isolation.py b/test/serena/test_make_tool_process_isolation.py index 3404aa1..fb13f86 100644 --- a/test/serena/test_make_tool_process_isolation.py +++ b/test/serena/test_make_tool_process_isolation.py @@ -24,7 +24,7 @@ def regular_agent(in_memory_config): @pytest.fixture def process_isolated_agent(in_memory_config): """Create a ProcessIsolatedSerenaAgent for comparison.""" - agent = ProcessIsolatedSerenaAgent(in_memory_config) + agent = ProcessIsolatedSerenaAgent(serena_config=in_memory_config) agent.start() yield agent agent.stop()