From cbf4933c04959f423f1d037199c514d50d3cd775 Mon Sep 17 00:00:00 2001 From: Michael Panchenko Date: Sun, 15 Jun 2025 20:47:36 +0200 Subject: [PATCH] Forwarding logs from agent process, WIP --- src/serena/dashboard.py | 4 + src/serena/mcp.py | 61 +++++++-------- src/serena/process_isolated_agent.py | 106 +++++++++++---------------- 3 files changed, 72 insertions(+), 99 deletions(-) diff --git a/src/serena/dashboard.py b/src/serena/dashboard.py index 549c003c..63bf79ab 100644 --- a/src/serena/dashboard.py +++ b/src/serena/dashboard.py @@ -78,6 +78,10 @@ class SerenaDashboardAPI: self._app = FastAPI(title="Serena Dashboard") self._setup_routes() + @property + def memory_log_handler(self) -> MemoryLogHandler: + return self._memory_log_handler + def _setup_routes(self) -> None: self._app.mount("/dashboard", StaticFiles(directory=SERENA_DASHBOARD_DIR), name="dashboard") diff --git a/src/serena/mcp.py b/src/serena/mcp.py index 6c91d0bd..ff4d4c8e 100644 --- a/src/serena/mcp.py +++ b/src/serena/mcp.py @@ -2,13 +2,14 @@ The Serena Model Context Protocol (MCP) Server """ +import contextlib import sys from collections.abc import AsyncIterator, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass from logging import Formatter, Logger, StreamHandler from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal import click # Add click import import docstring_parser @@ -21,10 +22,7 @@ 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 - -# Import for type annotations -if TYPE_CHECKING: - from serena.process_isolated_agent import ProcessIsolatedSerenaAgent +from serena.process_isolated_agent import ProcessIsolatedDashboard, ProcessIsolatedSerenaAgent log = logging.getLogger(__name__) LOG_FORMAT = "%(levelname)-5s %(asctime)-15s %(name)s:%(funcName)s:%(lineno)d - %(message)s" @@ -141,10 +139,7 @@ def create_mcp_server_and_agent( tool_timeout=tool_timeout, ) - # Create process-isolated agent with the configuration - from serena.process_isolated_agent import ProcessIsolatedDashboard, ProcessIsolatedSerenaAgent - - process_agent = ProcessIsolatedSerenaAgent(serena_config=serena_config) + serena_agent_process = ProcessIsolatedSerenaAgent(serena_config=serena_config) except Exception as e: show_fatal_exception_safe(e) @@ -152,10 +147,10 @@ def create_mcp_server_and_agent( def update_tools() -> None: """Update the tools in the MCP server - adapted for process isolation.""" - nonlocal mcp, process_agent + nonlocal mcp, serena_agent_process # Check if process agent is running - if process_agent.process is None or not process_agent.process.is_alive(): + 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 @@ -165,7 +160,7 @@ def create_mcp_server_and_agent( # 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 = process_agent.get_exposed_tool_names() + 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 @@ -175,27 +170,27 @@ def create_mcp_server_and_agent( mcp._tool_manager._tools = {} for tool_name in tool_names: - process_isolated_tool = ProcessIsolatedTool(process_agent=process_agent, tool_name=tool_name) + process_isolated_tool = ProcessIsolatedTool(process_agent=serena_agent_process, tool_name=tool_name) mcp_tool = make_tool(process_isolated_tool) mcp._tool_manager._tools[tool_name] = mcp_tool @asynccontextmanager async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[None]: """Manage server startup and shutdown lifecycle.""" - nonlocal process_agent + nonlocal serena_agent_process import asyncio import webbrowser mark_used(mcp_server) - dashboard: ProcessIsolatedDashboard | None = None + serena_dashboard_process: ProcessIsolatedDashboard | None = None shutdown_event = asyncio.Event() async def monitor_worker_process() -> None: """Monitor the worker process and shutdown server if it exits.""" while not shutdown_event.is_set(): try: - if process_agent.process is None or not process_agent.process.is_alive(): + 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 @@ -208,14 +203,14 @@ def create_mcp_server_and_agent( async def monitor_dashboard_shutdown() -> None: """Monitor the dashboard shutdown queue.""" - if dashboard is None or dashboard.shutdown_queue is None: + if serena_dashboard_process is None or serena_dashboard_process.shutdown_queue is None: return while not shutdown_event.is_set(): try: # Check if dashboard sent shutdown signal - if not dashboard.shutdown_queue.empty(): - dashboard.shutdown_queue.get_nowait() + 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 @@ -247,27 +242,26 @@ def create_mcp_server_and_agent( os._exit(0) try: - process_agent.start() # Set up global shutdown function for dashboard from serena.process_isolated_agent import set_global_shutdown_func set_global_shutdown_func(shutdown_server) - # Update tools now that the process agent is running - update_tools() - # Start process-isolated dashboard if enabled if serena_config.web_dashboard: try: - tool_names = process_agent.get_exposed_tool_names() - dashboard = ProcessIsolatedDashboard(tool_names) - port = dashboard.start() + 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}") - dashboard = None + serena_dashboard_process = None + + serena_agent_process.start() + # Update tools now that the process agent is running + update_tools() # Start monitoring tasks tasks = [] @@ -275,7 +269,7 @@ def create_mcp_server_and_agent( tasks.append(monitor_task) # Start dashboard shutdown monitor if dashboard is running - if dashboard is not None: + if serena_dashboard_process is not None: dashboard_monitor_task = asyncio.create_task(monitor_dashboard_shutdown()) tasks.append(dashboard_monitor_task) @@ -314,9 +308,6 @@ def create_mcp_server_and_agent( finally: log.info("Starting server shutdown cleanup") - # Cancel monitoring tasks - import contextlib - if "tasks" in locals(): for task in tasks: task.cancel() @@ -324,15 +315,15 @@ def create_mcp_server_and_agent( await asyncio.wait_for(task, timeout=1.0) # Stop dashboard first - if dashboard is not None: + if serena_dashboard_process is not None: try: - dashboard.stop() + serena_dashboard_process.stop() except Exception as e: log.error(f"Error stopping dashboard: {e}") # Stop process agent try: - process_agent.stop() + serena_agent_process.stop() except Exception as e: log.error(f"Error stopping process agent: {e}") @@ -343,7 +334,7 @@ def create_mcp_server_and_agent( update_tools() - return mcp, process_agent + return mcp, serena_agent_process class ProjectType(click.ParamType): diff --git a/src/serena/process_isolated_agent.py b/src/serena/process_isolated_agent.py index a9b59a7b..7b16767b 100644 --- a/src/serena/process_isolated_agent.py +++ b/src/serena/process_isolated_agent.py @@ -8,16 +8,20 @@ This module provides: """ import asyncio +import contextlib import logging import multiprocessing import signal import traceback from collections.abc import Callable from enum import StrEnum +from logging.handlers import QueueHandler from multiprocessing.connection import Connection 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 @@ -25,6 +29,8 @@ from serena.dashboard import MemoryLogHandler, SerenaDashboardAPI log = logging.getLogger(__name__) +_global_log_q: "multiprocessing.Queue[str]" = multiprocessing.Queue() + # Global shutdown function that can be called by the dashboard class _GlobalShutdownRegistry: @@ -62,7 +68,6 @@ class ProcessIsolatedDashboard: def __init__(self, tool_names: list[str]): self.tool_names = tool_names self.process: multiprocessing.Process | None = None - self.log_queue: multiprocessing.Queue[str] | None = None self.shutdown_queue: multiprocessing.Queue[bool] | None = None self.port_queue: multiprocessing.Queue[int] | None = None self._log_handler: ProcessDashboardLogHandler | None = None @@ -75,25 +80,21 @@ class ProcessIsolatedDashboard: log.info("Starting process-isolated dashboard") # Create communication queues - self.log_queue = multiprocessing.Queue() self.shutdown_queue = multiprocessing.Queue() self.port_queue = multiprocessing.Queue() # Create and start dashboard process - dashboard_worker = DashboardWorker( - log_queue=self.log_queue, shutdown_queue=self.shutdown_queue, port_queue=self.port_queue, tool_names=self.tool_names - ) + 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.start() # Set up log handler to send logs to dashboard process - self._log_handler = ProcessDashboardLogHandler(self.log_queue) - logging.getLogger().addHandler(self._log_handler) + logging.Logger.root.addHandler(ProcessDashboardLogHandler()) # Wait for dashboard to start and return port try: port = self.port_queue.get(timeout=10) # Wait up to 10 seconds for startup - log.info(f"Dashboard started on port {port}") + log.debug(f"Dashboard started on port {port}") return port except Exception as e: self.stop() @@ -107,11 +108,6 @@ class ProcessIsolatedDashboard: log.info("Stopping process-isolated dashboard") try: - # Remove log handler - if self._log_handler is not None: - logging.getLogger().removeHandler(self._log_handler) - self._log_handler = None - # Signal shutdown if self.shutdown_queue is not None: import contextlib @@ -138,7 +134,6 @@ class ProcessIsolatedDashboard: finally: self.process = None - self.log_queue = None self.shutdown_queue = None self.port_queue = None @@ -153,19 +148,14 @@ class ProcessIsolatedDashboard: class ProcessDashboardLogHandler(logging.Handler): """Log handler that sends log messages to dashboard process via queue.""" - def __init__(self, log_queue: "multiprocessing.Queue[str]", level: int = logging.NOTSET): + def __init__(self, level: int = logging.NOTSET): super().__init__(level=level) - self.setFormatter(logging.Formatter("%(levelname)s:%(name)s:%(message)s")) - self.log_queue = log_queue + self.setFormatter(logging.Formatter(LOG_DEFAULT_FORMAT)) def emit(self, record: logging.LogRecord) -> None: - try: - msg = self.format(record) - # Non-blocking put to avoid deadlocks - self.log_queue.put_nowait(msg) - except Exception: - # Silently ignore queue errors to avoid recursive logging - pass + msg = self.format(record) + # Non-blocking put to avoid deadlocks + _global_log_q.put_nowait(msg) class DashboardWorker: @@ -173,18 +163,13 @@ class DashboardWorker: def __init__( self, - log_queue: "multiprocessing.Queue[str]", shutdown_queue: "multiprocessing.Queue[bool]", port_queue: "multiprocessing.Queue[int]", tool_names: list[str], ): - self.log_queue = log_queue self.shutdown_queue = shutdown_queue self.port_queue = port_queue self.tool_names = tool_names - self.memory_log_handler: MemoryLogHandler | None = None - self.dashboard_api: SerenaDashboardAPI | None = None - self.log_processor_task: asyncio.Task | None = None def run(self) -> None: """Main dashboard worker loop - runs in separate process.""" @@ -198,28 +183,21 @@ class DashboardWorker: async def _async_main(self) -> None: """Main async loop for dashboard process.""" + log_processor_task = None try: - # Set up memory log handler - self.memory_log_handler = MemoryLogHandler() - # Set up dashboard API with shutdown callback - self.dashboard_api = SerenaDashboardAPI( - memory_log_handler=self.memory_log_handler, tool_names=self.tool_names, shutdown_callback=self._handle_shutdown + dashboard_api = SerenaDashboardAPI( + memory_log_handler=MemoryLogHandler(), tool_names=self.tool_names, shutdown_callback=self._handle_shutdown ) # Start log processor task - self.log_processor_task = asyncio.create_task(self._process_log_queue()) + log_processor_task = asyncio.create_task(self._process_log_queue(dashboard_api)) # Find free port and signal to parent - port = self.dashboard_api._find_first_free_port(0x5EDA) + port = dashboard_api._find_first_free_port(0x5EDA) self.port_queue.put(port) - # Start dashboard server - import uvicorn - - config = uvicorn.Config( - app=self.dashboard_api._app, host="0.0.0.0", port=port, workers=1, log_config=None, log_level="critical" - ) + 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 @@ -228,12 +206,9 @@ class DashboardWorker: # 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() - import contextlib - with contextlib.suppress(asyncio.CancelledError): await task @@ -241,33 +216,31 @@ class DashboardWorker: log.error(f"Error in dashboard worker: {e}") finally: # Clean up - if self.log_processor_task and not self.log_processor_task.done(): - self.log_processor_task.cancel() - import contextlib + if log_processor_task and not log_processor_task.done(): + log_processor_task.cancel() with contextlib.suppress(asyncio.CancelledError): - await self.log_processor_task + await log_processor_task - async def _process_log_queue(self) -> None: + @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 self.log_queue.empty(): + while not _global_log_q.empty(): try: - log_msg = self.log_queue.get_nowait() - if self.memory_log_handler: - # Create a dummy log record and emit it - record = logging.LogRecord( - name="forwarded", level=logging.INFO, pathname="", lineno=0, msg=log_msg, args=(), exc_info=None - ) - self.memory_log_handler.emit(record) - except Exception: - break + 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: @@ -331,8 +304,13 @@ class SerenaAgentWorker: self.agent: SerenaAgent | None = None self._shutdown_requested = False - def run(self) -> None: + def run(self, q: "multiprocessing.Queue[str]") -> None: """Main worker loop - runs in separate process.""" + qh = QueueHandler(q) + root = logging.getLogger() + root.setLevel(logging.DEBUG) + root.addHandler(qh) + try: # Set up signal handler for clean shutdown signal.signal(signal.SIGTERM, self._signal_handler) @@ -623,9 +601,9 @@ class ProcessIsolatedSerenaAgent: parent_conn, child_conn = multiprocessing.Pipe() self.conn = parent_conn - # Create and start worker process + # Create and start worker process, passing along the dashboard's queue if available worker = SerenaAgentWorker(child_conn) - self.process = multiprocessing.Process(target=worker.run) + self.process = multiprocessing.Process(target=worker.run, args=[_global_log_q]) self.process.start() # Initialize the agent in the worker process