diff --git a/src/serena/dashboard.py b/src/serena/dashboard.py index 63bf79a..22e3495 100644 --- a/src/serena/dashboard.py +++ b/src/serena/dashboard.py @@ -1,7 +1,6 @@ import os import queue import socket -import sys import threading from collections.abc import Callable @@ -98,21 +97,16 @@ class SerenaDashboardAPI: return ResponseToolNames(tool_names=self._tool_names) async def _shutdown(self) -> None: - print("Shutdown initiated by dashbaord ...", file=sys.stderr) log.info("Shutting down Serena") 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 + from serena.process_isolated_agent import request_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) + request_global_shutdown() + # 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 ff4d4c8..3551ec1 100644 --- a/src/serena/mcp.py +++ b/src/serena/mcp.py @@ -2,8 +2,13 @@ The Serena Model Context Protocol (MCP) Server """ +import asyncio import contextlib +import os +import signal import sys +import threading +import time from collections.abc import AsyncIterator, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass @@ -11,7 +16,7 @@ from logging import Formatter, Logger, StreamHandler from pathlib import Path from typing import Any, Literal -import click # Add click import +import click import docstring_parser from mcp.server.fastmcp import server from mcp.server.fastmcp.server import FastMCP, Settings @@ -22,7 +27,13 @@ from sensai.util.helper import mark_used from serena.agent import SerenaAgent, ToolInterface, create_serena_config, show_fatal_exception_safe from serena.config import SerenaAgentContext, SerenaAgentMode from serena.constants import DEFAULT_CONTEXT, DEFAULT_MODES -from serena.process_isolated_agent import ProcessIsolatedDashboard, ProcessIsolatedSerenaAgent +from serena.process_isolated_agent import ( + ProcessIsolatedDashboard, + ProcessIsolatedSerenaAgent, + ProcessIsolatedTool, + global_shutdown_event, + request_global_shutdown, +) log = logging.getLogger(__name__) LOG_FORMAT = "%(levelname)-5s %(asctime)-15s %(name)s:%(funcName)s:%(lineno)d - %(message)s" @@ -103,7 +114,7 @@ def create_mcp_server_and_agent( log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] | None = None, trace_lsp_communication: bool | None = None, tool_timeout: float | None = None, -) -> tuple[FastMCP, "ProcessIsolatedSerenaAgent"]: +) -> tuple[FastMCP, ProcessIsolatedSerenaAgent]: """ Create an MCP server with process-isolated SerenaAgent to prevent asyncio contamination. @@ -127,7 +138,6 @@ def create_mcp_server_and_agent( modes_instances = [SerenaAgentMode.load(mode) for mode in modes] try: - # Create configuration without instantiating a full SerenaAgent serena_config = create_serena_config( project=project, context=context_instance, @@ -138,9 +148,12 @@ def create_mcp_server_and_agent( trace_lsp_communication=trace_lsp_communication, tool_timeout=tool_timeout, ) - serena_agent_process = ProcessIsolatedSerenaAgent(serena_config=serena_config) + # Start process-isolated dashboard if enabled + serena_dashboard_process = None + if serena_config.web_dashboard: + serena_dashboard_process = ProcessIsolatedDashboard(tool_names=[]) except Exception as e: show_fatal_exception_safe(e) raise @@ -149,25 +162,13 @@ def create_mcp_server_and_agent( """Update the tools in the MCP server - adapted for process isolation.""" nonlocal mcp, serena_agent_process - # Check if process agent is running - if serena_agent_process.process is None or not serena_agent_process.process.is_alive(): - log.debug("Process agent not running yet, skipping tool update") - return - # Get tool names from process-isolated agent # Tools may change as a result of project activation. # NOTE: While we could pass updated tool information on to the MCP server via the callback, Claude Desktop does not, # unfortunately, query for changed tools. It only queries for changed resources and prompts regularly, # so we need to register all tools at startup, unfortunately. - try: - tool_names = serena_agent_process.get_exposed_tool_names() - except Exception as e: - log.error(f"Failed to get tool names from process agent: {e}") - return - + tool_names = serena_agent_process.get_exposed_tool_names() if mcp is not None: - from serena.process_isolated_agent import ProcessIsolatedTool - mcp._tool_manager._tools = {} for tool_name in tool_names: process_isolated_tool = ProcessIsolatedTool(process_agent=serena_agent_process, tool_name=tool_name) @@ -177,163 +178,69 @@ def create_mcp_server_and_agent( @asynccontextmanager async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[None]: """Manage server startup and shutdown lifecycle.""" - nonlocal serena_agent_process - import asyncio - import webbrowser - mark_used(mcp_server) - serena_dashboard_process: ProcessIsolatedDashboard | None = None - shutdown_event = asyncio.Event() + def signal_handler(signum: int, frame: Any) -> None: + log.info(f"Received signal {signum} in main process") + request_global_shutdown() - async def monitor_worker_process() -> None: - """Monitor the worker process and shutdown server if it exits.""" - while not shutdown_event.is_set(): - try: - if serena_agent_process.process is None or not serena_agent_process.process.is_alive(): - log.info("Worker process has exited, shutting down MCP server") - shutdown_event.set() - break - await asyncio.sleep(1) - except asyncio.CancelledError: - break - except Exception as e: - log.error(f"Error monitoring worker process: {e}") - await asyncio.sleep(1) + def force_exit() -> None: + time.sleep(2.0) # Wait 2 seconds for graceful shutdown + log.warning("Forcing exit after timeout") + os._exit(1) - async def monitor_dashboard_shutdown() -> None: - """Monitor the dashboard shutdown queue.""" - if serena_dashboard_process is None or serena_dashboard_process.shutdown_queue is None: - return + threading.Thread(target=force_exit, daemon=True).start() - while not shutdown_event.is_set(): - try: - # Check if dashboard sent shutdown signal - if not serena_dashboard_process.shutdown_queue.empty(): - serena_dashboard_process.shutdown_queue.get_nowait() - log.info("Dashboard shutdown signal received via queue") - shutdown_event.set() - # Force immediate exit since MCP server doesn't respond properly - import os - import sys + # Install signal handlers + sigint_singal = signal.signal(signal.SIGINT, signal_handler) + sigterm_signal = signal.signal(signal.SIGTERM, signal_handler) - log.info("Forcing immediate process exit from dashboard signal") - sys.stdout.flush() - sys.stderr.flush() - os._exit(0) - await asyncio.sleep(0.1) # Check more frequently - except asyncio.CancelledError: - break - except Exception as e: - log.error(f"Error monitoring dashboard shutdown: {e}") - await asyncio.sleep(1) + if serena_dashboard_process is not None: + log.info("Starting dashboard process") + serena_dashboard_process.start() + log.info("Starting serena agent process") + serena_agent_process.start() + update_tools() - def shutdown_server() -> None: - """Shutdown the server.""" - log.info("Shutdown initiated by dashboard") - shutdown_event.set() - # Force immediate exit since the MCP server doesn't respond to shutdown events - import os - import sys + async def monitor_global_shutdown() -> None: + """Monitor the global shutdown event and trigger local shutdown.""" + while not global_shutdown_event.is_set(): + # Poll the multiprocessing Event in async context + await asyncio.sleep(0.1) + continue + log.info("Global shutdown event detected, initiating server shutdown") + request_global_shutdown() + # Send SIGTERM to self to trigger graceful shutdown + os.kill(os.getpid(), signal.SIGTERM) - log.info("Dashboard shutdown - forcing process exit immediately") - sys.stdout.flush() - sys.stderr.flush() - os._exit(0) + # Start monitoring task + monitor_task = asyncio.create_task(monitor_global_shutdown()) try: - - # Set up global shutdown function for dashboard - from serena.process_isolated_agent import set_global_shutdown_func - - set_global_shutdown_func(shutdown_server) - - # Start process-isolated dashboard if enabled - if serena_config.web_dashboard: - try: - serena_dashboard_process = ProcessIsolatedDashboard(tool_names=[]) - port = serena_dashboard_process.start() - webbrowser.open(f"http://localhost:{port}/dashboard/index.html") - log.info(f"Dashboard started on port {port}") - except Exception as e: - log.error(f"Failed to start dashboard: {e}") - serena_dashboard_process = None - - serena_agent_process.start() - # Update tools now that the process agent is running - update_tools() - - # Start monitoring tasks - tasks = [] - monitor_task = asyncio.create_task(monitor_worker_process()) - tasks.append(monitor_task) - - # Start dashboard shutdown monitor if dashboard is running - if serena_dashboard_process is not None: - dashboard_monitor_task = asyncio.create_task(monitor_dashboard_shutdown()) - tasks.append(dashboard_monitor_task) - - # Set up signal handlers in main process - def signal_handler(signum: int, frame: Any) -> None: - log.info(f"Received signal {signum} in main process") - shutdown_event.set() - # Give a short delay for graceful shutdown, then force exit - import os - import threading - - def force_exit() -> None: - import time - - time.sleep(2.0) # Wait 2 seconds for graceful shutdown - log.warning("Forcing exit after timeout") - os._exit(1) - - threading.Thread(target=force_exit, daemon=True).start() - - import signal - - signal.signal(signal.SIGTERM, signal_handler) - signal.signal(signal.SIGINT, signal_handler) - - try: - # Start the server and wait for shutdown - yield - except (KeyboardInterrupt, SystemExit): - log.info("Received shutdown signal") - shutdown_event.set() - + yield + except (KeyboardInterrupt, SystemExit): + log.info("Received shutdown signal") + request_global_shutdown() except Exception as e: log.error(f"Error in server lifespan: {e}") - shutdown_event.set() + request_global_shutdown() finally: - log.info("Starting server shutdown cleanup") + # Cancel monitor task + monitor_task.cancel() - if "tasks" in locals(): - for task in tasks: - task.cancel() - with contextlib.suppress(TimeoutError, asyncio.CancelledError): - await asyncio.wait_for(task, timeout=1.0) + with contextlib.suppress(asyncio.CancelledError): + await monitor_task - # Stop dashboard first + serena_agent_process.stop() if serena_dashboard_process is not None: - try: - serena_dashboard_process.stop() - except Exception as e: - log.error(f"Error stopping dashboard: {e}") - - # Stop process agent - try: - serena_agent_process.stop() - except Exception as e: - log.error(f"Error stopping process agent: {e}") - - log.info("Server shutdown cleanup completed") + serena_dashboard_process.stop() + request_global_shutdown() + log.info("Shutting down all processes") + signal.signal(signal.SIGINT, sigint_singal) + signal.signal(signal.SIGTERM, sigterm_signal) mcp_settings = Settings(lifespan=server_lifespan, host=host, port=port) mcp = FastMCP(**mcp_settings.model_dump()) - - update_tools() - return mcp, serena_agent_process @@ -494,9 +401,6 @@ def start_mcp_server( f"Used path: {project_file}" ) - log.info( - f"Starting process-isolated serena agent in MCP server with config:\n{agent.serena_config}." - f"\n Log level: {agent.serena_config.log_level}" - ) + log.info(f"Starting MCP server with config:\n{agent.serena_config}.\n Log level: {agent.serena_config.log_level}") mcp_server.run(transport=transport) diff --git a/src/serena/process_isolated_agent.py b/src/serena/process_isolated_agent.py index 7b16767..768e45f 100644 --- a/src/serena/process_isolated_agent.py +++ b/src/serena/process_isolated_agent.py @@ -1,27 +1,20 @@ -""" -Process-isolated SerenaAgent to prevent asyncio event loop contamination between MCP server and language server. - -This module provides: -1. ProcessIsolatedSerenaAgent - A wrapper that runs SerenaAgent in a separate process -2. SerenaAgentWorker - The worker process that hosts the actual SerenaAgent -3. JSON-RPC based IPC for communication between processes -""" - import asyncio import contextlib import logging import multiprocessing -import signal +import os import traceback +import webbrowser from collections.abc import Callable from enum import StrEnum from logging.handlers import QueueHandler from multiprocessing.connection import Connection +from multiprocessing.sharedctypes import Synchronized +from multiprocessing.synchronize import Event as EventClass from typing import Any, Literal, Self import uvicorn from mcp.server.fastmcp.utilities.func_metadata import FuncMetadata -from sensai.util.logging import LOG_DEFAULT_FORMAT from serena.agent import SerenaAgent, SerenaConfig, SerenaConfigBase, Tool, ToolInterface, ToolRegistry from serena.config import SerenaAgentContext, SerenaAgentMode @@ -29,296 +22,182 @@ from serena.dashboard import MemoryLogHandler, SerenaDashboardAPI log = logging.getLogger(__name__) -_global_log_q: "multiprocessing.Queue[str]" = multiprocessing.Queue() +# Global synchronization primitives +_global_log_queue: multiprocessing.Queue = multiprocessing.Queue() +_dashboard_ready_event = multiprocessing.Event() +_dashboard_port_value = multiprocessing.Value("i", 0) +global_shutdown_event = multiprocessing.Event() -# 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() +def request_global_shutdown() -> None: + """Signal the global shutdown event.""" + global_shutdown_event.set() + log.info("Global shutdown event set") -_shutdown_registry = _GlobalShutdownRegistry() +def _dashboard_worker( + tool_names: list[str], + log_q: "multiprocessing.Queue[Any]", + dashboard_ready_event: EventClass, + port_value: "Synchronized[int]", + shutdown_evt: EventClass, +) -> None: + """Entry point for the dashboard process.""" + # Route all logging to the shared queue + root = logging.getLogger() + root.handlers.clear() + root.setLevel(logging.DEBUG) + root.addHandler(QueueHandler(log_q)) + + async def _process_logs(api: SerenaDashboardAPI) -> None: + while not shutdown_evt.is_set(): + while not log_q.empty(): + record = log_q.get_nowait() + if record is None: + break + api.memory_log_handler.emit(record) + # Small delay to avoid busy waiting + await asyncio.sleep(0.1) + + async def _monitor_shutdown() -> None: + # Poll the multiprocessing Event in async context + # Check every 100ms until shutdown is requested + loop = asyncio.get_event_loop() + while True: + # Check in executor to avoid blocking + result = await loop.run_in_executor(None, shutdown_evt.is_set) + if result: + break + await asyncio.sleep(0.1) + + async def _async_main() -> None: + api = SerenaDashboardAPI( + memory_log_handler=MemoryLogHandler(), + tool_names=tool_names, + shutdown_callback=shutdown_evt.set, + ) + # Pick a free port and signal readiness + port = api._find_first_free_port(0x5EDA) + port_value.value = port + + # Start server first, then log processing + config = uvicorn.Config( + app=api._app, + host="0.0.0.0", + port=port, + workers=1, + log_config=None, + log_level="critical", + ) + server = uvicorn.Server(config) + + # Serve and await shutdown signal + server_task = asyncio.create_task(server.serve()) + shutdown_task = asyncio.create_task(_monitor_shutdown()) + logging_loop_task = asyncio.create_task(_process_logs(api)) + dashboard_ready_event.set() + + done, pending = await asyncio.wait( + [server_task, shutdown_task, logging_loop_task], + return_when=asyncio.FIRST_COMPLETED, + ) + if shutdown_task in done: + server.should_exit = True + await server_task + + # Cancel remaining tasks + for task in pending: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + try: + asyncio.run(_async_main()) + except BaseException: + logging.exception("Dashboard worker crashed") + finally: + logging.info("Dashboard worker exiting") -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() +def _shutdown_process(proc: multiprocessing.Process, timeout: float = 1.0) -> None: + """Helper to shutdown a process gracefully.""" + if proc.is_alive(): + proc.terminate() + proc.join(timeout=timeout) + if proc.is_alive(): + log.error("Process did not terminate gracefully, forcing kill") + proc.kill() + proc.join(timeout=timeout) + if proc.is_alive(): + log.error("Process did not respond to kill") class ProcessIsolatedDashboard: - """Process-isolated dashboard to prevent asyncio contamination with MCP server and worker.""" + """Dashboard running in a separate process to avoid asyncio contamination.""" def __init__(self, tool_names: list[str]): self.tool_names = tool_names self.process: multiprocessing.Process | None = None - self.shutdown_queue: multiprocessing.Queue[bool] | None = None - self.port_queue: multiprocessing.Queue[int] | None = None - self._log_handler: ProcessDashboardLogHandler | None = None - def start(self) -> int: - """Start the dashboard process and return the port number.""" + def start(self, timeout: float = 10.0) -> None: + """Start the dashboard process.""" if self.process is not None: - raise RuntimeError("Dashboard process already started") + raise RuntimeError("Dashboard already started") - log.info("Starting process-isolated dashboard") - - # Create communication queues - self.shutdown_queue = multiprocessing.Queue() - self.port_queue = multiprocessing.Queue() - - # Create and start dashboard process - dashboard_worker = DashboardWorker(shutdown_queue=self.shutdown_queue, port_queue=self.port_queue, tool_names=self.tool_names) - self.process = multiprocessing.Process(target=dashboard_worker.run) + self.process = multiprocessing.Process( + target=_dashboard_worker, + args=(self.tool_names, _global_log_queue, _dashboard_ready_event, _dashboard_port_value, global_shutdown_event), + daemon=True, + ) self.process.start() - # Set up log handler to send logs to dashboard process - logging.Logger.root.addHandler(ProcessDashboardLogHandler()) + if not _dashboard_ready_event.wait(timeout): + self.process.terminate() + self.process.join(timeout=1.0) + self.process = None + raise RuntimeError("Dashboard failed to start within timeout") - # Wait for dashboard to start and return port - try: - port = self.port_queue.get(timeout=10) # Wait up to 10 seconds for startup - log.debug(f"Dashboard started on port {port}") - return port - except Exception as e: - self.stop() - raise RuntimeError(f"Failed to start dashboard: {e}") from e + port = _dashboard_port_value.value + log.info(f"Dashboard started on port {port}") + webbrowser.open(f"http://localhost:{port}/dashboard/index.html") - def stop(self) -> None: - """Stop the dashboard process.""" + def stop(self, timeout: float = 1.0) -> None: + """Signal shutdown and wait for the dashboard process to exit.""" if self.process is None: return - log.info("Stopping process-isolated dashboard") - - try: - # Signal shutdown - if self.shutdown_queue is not None: - import contextlib - - with contextlib.suppress(Exception): - self.shutdown_queue.put_nowait(True) - - # Wait for process to terminate - if self.process.is_alive(): - self.process.join(timeout=5.0) - - if self.process.is_alive(): - log.warning("Dashboard process did not terminate gracefully, forcing termination") - self.process.terminate() - self.process.join(timeout=3.0) - - if self.process.is_alive(): - log.error("Dashboard process could not be terminated, killing it") - self.process.kill() - self.process.join() - - except Exception as e: - log.error(f"Error stopping dashboard process: {e}") - - finally: - self.process = None - self.shutdown_queue = None - self.port_queue = None - - log.info("Process-isolated dashboard stopped") - - def update_tool_names(self, tool_names: list[str]) -> None: - """Update tool names (for future enhancement).""" - self.tool_names = tool_names - # TODO: Could send update to dashboard process if needed - - -class ProcessDashboardLogHandler(logging.Handler): - """Log handler that sends log messages to dashboard process via queue.""" - - def __init__(self, level: int = logging.NOTSET): - super().__init__(level=level) - self.setFormatter(logging.Formatter(LOG_DEFAULT_FORMAT)) - - def emit(self, record: logging.LogRecord) -> None: - msg = self.format(record) - # Non-blocking put to avoid deadlocks - _global_log_q.put_nowait(msg) - - -class DashboardWorker: - """Worker process that hosts the dashboard with its own asyncio loop.""" - - def __init__( - self, - shutdown_queue: "multiprocessing.Queue[bool]", - port_queue: "multiprocessing.Queue[int]", - tool_names: list[str], - ): - self.shutdown_queue = shutdown_queue - self.port_queue = port_queue - self.tool_names = tool_names - - def run(self) -> None: - """Main dashboard worker loop - runs in separate process.""" - try: - log.info("Dashboard worker process started") - asyncio.run(self._async_main()) - except Exception as e: - log.error(f"Fatal error in dashboard worker process: {e}") - finally: - log.info("Dashboard worker process stopped") - - async def _async_main(self) -> None: - """Main async loop for dashboard process.""" - log_processor_task = None - try: - # Set up dashboard API with shutdown callback - dashboard_api = SerenaDashboardAPI( - memory_log_handler=MemoryLogHandler(), tool_names=self.tool_names, shutdown_callback=self._handle_shutdown - ) - - # Start log processor task - log_processor_task = asyncio.create_task(self._process_log_queue(dashboard_api)) - - # Find free port and signal to parent - port = dashboard_api._find_first_free_port(0x5EDA) - self.port_queue.put(port) - - config = uvicorn.Config(app=dashboard_api._app, host="0.0.0.0", port=port, workers=1, log_config=None, log_level="critical") - server = uvicorn.Server(config) - - # Start server and shutdown monitor concurrently - shutdown_task = asyncio.create_task(self._monitor_shutdown()) - server_task = asyncio.create_task(server.serve()) - - # Wait for either server to complete or shutdown signal - done, pending = await asyncio.wait([server_task, shutdown_task], return_when=asyncio.FIRST_COMPLETED) - # Cancel pending tasks - for task in pending: - task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await task - - except Exception as e: - log.error(f"Error in dashboard worker: {e}") - finally: - # Clean up - if log_processor_task and not log_processor_task.done(): - log_processor_task.cancel() - - with contextlib.suppress(asyncio.CancelledError): - await log_processor_task - - @staticmethod - async def _process_log_queue(dashboard_api: SerenaDashboardAPI) -> None: - """Process log messages from the queue and forward to memory handler.""" - while True: - try: - # Check for log messages - while not _global_log_q.empty(): - try: - log_msg = _global_log_q.get_nowait() - record = logging.LogRecord( - name="forwarded", level=logging.INFO, pathname="", lineno=0, msg=log_msg, args=(), exc_info=None - ) - dashboard_api.memory_log_handler.emit(record) - except Exception as e: - log.error(f"Error processing log message: {e}") - await asyncio.sleep(1) - - # Small delay to avoid busy waiting - await asyncio.sleep(0.1) - except asyncio.CancelledError: - break - except Exception as e: - log.error(f"Error processing log queue: {e}") - await asyncio.sleep(1) - - async def _monitor_shutdown(self) -> None: - """Monitor for shutdown signals.""" - while True: - try: - # Check for shutdown signal - if not self.shutdown_queue.empty(): - self.shutdown_queue.get_nowait() - log.info("Dashboard shutdown signal received") - break - - await asyncio.sleep(0.5) - - except asyncio.CancelledError: - break - except Exception as e: - log.error(f"Error monitoring shutdown: {e}") - await asyncio.sleep(1) - - def _handle_shutdown(self) -> None: - """Handle shutdown request from dashboard UI.""" - log.info("Dashboard UI shutdown requested") - # Signal shutdown to the main process through the parent dashboard - if self.shutdown_queue is not None: - try: - self.shutdown_queue.put_nowait(True) - log.info("Sent shutdown signal to main process via queue") - except Exception as e: - log.error(f"Failed to send shutdown signal: {e}") - - # Also try the global shutdown as fallback - try: - call_global_shutdown() - except Exception as e: - log.error(f"Global shutdown failed: {e}") - - # Force exit this dashboard process - import os - import threading - - def delayed_exit() -> None: - import time - - time.sleep(0.5) # Give time for shutdown signal to be sent - log.info("Dashboard process forcing exit") - os._exit(0) - - threading.Thread(target=delayed_exit, daemon=True).start() + log.info("Stopping dashboard process") + request_global_shutdown() + _shutdown_process(self.process, timeout=timeout) + self.process = None class SerenaAgentWorker: """Worker process that hosts the actual SerenaAgent.""" + class RequestMethod(StrEnum): + INITIALIZE = "initialize" + TOOL_CALL = "tool_call" + GET_ACTIVE_TOOL_NAMES = "get_active_tool_names" + IS_LANGUAGE_SERVER_RUNNING = "is_language_server_running" + RESET_LANGUAGE_SERVER = "reset_language_server" + GET_EXPOSED_TOOL_NAMES = "get_exposed_tool_names" + SHUTDOWN = "shutdown" + def __init__(self, conn: Connection): self.conn = conn self.agent: SerenaAgent | None = None - self._shutdown_requested = False - def run(self, q: "multiprocessing.Queue[str]") -> None: + def run(self, log_queue: "multiprocessing.Queue[str]") -> None: """Main worker loop - runs in separate process.""" - qh = QueueHandler(q) + qh = QueueHandler(log_queue) root = logging.getLogger() root.setLevel(logging.DEBUG) root.addHandler(qh) + log.info("SerenaAgent worker process started") try: - # Set up signal handler for clean shutdown - signal.signal(signal.SIGTERM, self._signal_handler) - signal.signal(signal.SIGINT, self._signal_handler) - - log.info("SerenaAgent worker process started") - - while not self._shutdown_requested: + while not global_shutdown_event.is_set(): try: # Use polling to avoid blocking indefinitely if self.conn.poll(timeout=0.5): # Poll every 500ms @@ -354,17 +233,12 @@ class SerenaAgentWorker: except Exception as e: log.error(f"Error in worker main loop: {e}") break - except Exception as e: log.error(f"Fatal error in worker process: {e}") finally: self._cleanup() log.info("SerenaAgent worker process stopped") - - def _signal_handler(self, signum: int, frame: Any) -> None: - """Handle shutdown signals gracefully.""" - log.info(f"Received signal {signum}, initiating shutdown") - self._shutdown_requested = True + os._exit(0) # Exit without raising any further exceptions def _handle_request(self, request: dict[str, Any]) -> dict[str, Any]: """Handle a single request.""" @@ -386,10 +260,9 @@ class SerenaAgentWorker: case self.RequestMethod.GET_EXPOSED_TOOL_NAMES: return self._get_exposed_tool_names() case self.RequestMethod.SHUTDOWN: - return self._shutdown() + return self.shutdown() case _: return {"error": f"Unknown method: {method}"} - except Exception as e: return {"error": str(e), "traceback": traceback.format_exc()} @@ -427,17 +300,6 @@ class SerenaAgentWorker: trace_lsp_communication=trace_lsp_communication, 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()} @@ -515,18 +377,15 @@ 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.""" + def shutdown(self) -> dict[str, Any]: try: - log.info("Shutdown requested from dashboard") - # Clean up resources before exiting + log.info("Shutting down SerenaAgent worker process on request") 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) + request_global_shutdown() + # Return successful response before exiting + return {"result": "Shutdown initiated"} except Exception as e: + log.error(f"Error during shutdown: {e}") return {"error": str(e), "traceback": traceback.format_exc()} def _cleanup(self) -> None: @@ -539,17 +398,6 @@ class SerenaAgentWorker: log.error(f"Error stopping language server: {e}") self.agent = None - class RequestMethod(StrEnum): - """Enumeration of available request methods.""" - - INITIALIZE = "initialize" - TOOL_CALL = "tool_call" - GET_ACTIVE_TOOL_NAMES = "get_active_tool_names" - 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: """Process-isolated wrapper for SerenaAgent that prevents asyncio contamination.""" @@ -603,7 +451,7 @@ class ProcessIsolatedSerenaAgent: # Create and start worker process, passing along the dashboard's queue if available worker = SerenaAgentWorker(child_conn) - self.process = multiprocessing.Process(target=worker.run, args=[_global_log_q]) + self.process = multiprocessing.Process(target=worker.run, args=[_global_log_queue]) self.process.start() # Initialize the agent in the worker process @@ -634,31 +482,13 @@ class ProcessIsolatedSerenaAgent: """Stop the worker process.""" if self.process is None: return - - log.info("Stopping process-isolated SerenaAgent") - + log.info("Stopping SerenaAgent process") try: # Close connection to signal worker to shutdown if self.conn is not None: self.conn.close() - - # Wait for process to terminate with shorter timeout - if self.process.is_alive(): - self.process.join(timeout=3.0) # Reduced from 10s to 3s - - if self.process.is_alive(): - log.warning("Worker process did not terminate gracefully, sending SIGTERM") - self.process.terminate() - self.process.join(timeout=2.0) # Reduced from 5s to 2s - - if self.process.is_alive(): - log.error("Worker process did not respond to SIGTERM, sending SIGKILL") - self.process.kill() - self.process.join(timeout=1.0) # Add timeout for kill as well - - if self.process.is_alive(): - log.error("Worker process could not be killed - forcing cleanup") - + _shutdown_process(self.process, timeout=2.0) + self.process = None except KeyboardInterrupt: log.warning("Keyboard interrupt during shutdown - forcing termination") if self.process and self.process.is_alive(): @@ -666,12 +496,11 @@ class ProcessIsolatedSerenaAgent: self.process.join(timeout=1.0) except Exception as e: log.error(f"Error stopping worker process: {e}") - finally: self.process = None self.conn = None - log.info("Process-isolated SerenaAgent stopped") + log.info("SerenaAgent stopped") def _make_request(self, method: SerenaAgentWorker.RequestMethod, params: dict[str, Any] | None = None) -> dict[str, Any]: """Make a request to the worker process.""" @@ -693,8 +522,7 @@ class ProcessIsolatedSerenaAgent: timeout = self.serena_config.tool_timeout if self.conn.poll(timeout): try: - response = self.conn.recv() - return response + return self.conn.recv() except (EOFError, BrokenPipeError) as e: raise RuntimeError("Failed to receive response: worker process may have crashed") from e else: