diff --git a/src/multilspy/language_server.py b/src/multilspy/language_server.py index 41fc43c9..3ae97807 100644 --- a/src/multilspy/language_server.py +++ b/src/multilspy/language_server.py @@ -1827,7 +1827,13 @@ class SyncLanguageServer: self._server_context = self.language_server.start_server() asyncio.run_coroutine_threadsafe(self._server_context.__aenter__(), loop=self.loop).result() return self - + + def is_running(self) -> bool: + """ + Check if the language server is running. + """ + return self.loop is not None and self.loop_thread is not None and self.loop_thread.is_alive() + def stop(self) -> None: """ Shuts down the language server process and cleans up resources. diff --git a/src/serena/agent.py b/src/serena/agent.py index b98efaa6..c428c19a 100644 --- a/src/serena/agent.py +++ b/src/serena/agent.py @@ -10,8 +10,7 @@ import sys import traceback from abc import ABC from collections import defaultdict -from collections.abc import Callable, Generator, Iterable, Iterator -from contextlib import contextmanager +from collections.abc import Callable, Generator, Iterable from logging import Logger from pathlib import Path from typing import Any, TypeVar, cast @@ -93,9 +92,9 @@ class SerenaAgent: ) # create and start the language server instance - config = MultilspyConfig(code_language=self.language) - logger = MultilspyLogger() - self.language_server = SyncLanguageServer.create(config, logger, self.project_root) + self._config = MultilspyConfig(code_language=self.language) + self._ls_logger = MultilspyLogger() + self.language_server = SyncLanguageServer.create(self._config, self._ls_logger, self.project_root) self.prompt_factory = PromptFactory() self.symbol_manager = SymbolManager(self.language_server, self) @@ -123,6 +122,14 @@ class SerenaAgent: if self._start_language_server: log.info("Starting the language server ...") self.language_server.start() + if not self.language_server.is_running(): + raise RuntimeError(f"Failed to start the language server for {self.language} and project {self.project_root}") + + def reset_language_server(self) -> None: + self.language_server.stop() + self.language_server.start() + if not self.language_server.is_running(): + raise RuntimeError(f"Failed to start (reset) the language server for {self.language} and project {self.project_root}") def get_tool(self, tool_class: type[TTool]) -> TTool: return self._all_tools[tool_class] # type: ignore @@ -151,16 +158,6 @@ class SerenaAgent: self._gui_log_handler.stop_viewer() Logger.root.removeHandler(self._gui_log_handler) - @contextmanager - def language_server_lifecycle_context(self) -> Iterator[None]: - """ - Context manager for the language server's lifecycle - """ - if self._start_language_server: - raise Exception("This context manager can only be used if the instance is created with start_language_server=True") - with self.language_server.start_server(): - yield - class MemoriesManager: def __init__(self, memory_dir: str): @@ -195,12 +192,30 @@ class MemoriesManager: class Component(ABC): def __init__(self, agent: "SerenaAgent"): self.agent = agent - self.language_server = agent.language_server - self.project_root = agent.project_root - self.project_config = agent.project_config - self.prompt_factory = agent.prompt_factory - self.memories_manager = agent.memories_manager - self.symbol_manager = agent.symbol_manager + + @property + def language_server(self) -> SyncLanguageServer: + return self.agent.language_server + + @property + def project_root(self) -> str: + return self.agent.project_root + + @property + def project_config(self) -> dict[str, Any]: + return self.agent.project_config + + @property + def prompt_factory(self) -> PromptFactory: + return self.agent.prompt_factory + + @property + def memories_manager(self) -> MemoriesManager: + return self.agent.memories_manager + + @property + def symbol_manager(self) -> SymbolManager: + return self.agent.symbol_manager _DEFAULT_MAX_ANSWER_LENGTH = int(2e5) @@ -261,6 +276,9 @@ class Tool(Component): """ Applies the tool with the given arguments """ + if not self.language_server.is_running(): + self.agent.reset_language_server() + apply_fn = getattr(self, "apply") if apply_fn is None: raise ValueError(f"apply not defined in {self}") diff --git a/src/serena/mcp.py b/src/serena/mcp.py index 7d815f14..871aa96a 100644 --- a/src/serena/mcp.py +++ b/src/serena/mcp.py @@ -13,7 +13,6 @@ from mcp.server.fastmcp.server import FastMCP, Settings from mcp.server.fastmcp.tools.base import Tool as MCPTool from mcp.server.fastmcp.utilities.func_metadata import func_metadata from sensai.util import logging -from sensai.util.helper import mark_used from serena.agent import SerenaAgent, Tool @@ -43,8 +42,6 @@ def make_tool( tool: Tool, ) -> MCPTool: """Create a Tool from a function.""" - from mcp.server.fastmcp import Context - func_name = tool.get_name() apply_fn = getattr(tool, "apply") @@ -57,8 +54,7 @@ def make_tool( func_arg_metadata = func_metadata(apply_fn) parameters = func_arg_metadata.arg_model.model_json_schema() - def execute_fn(ctx: Context, **kwargs) -> str: # type: ignore - mark_used(ctx) + def execute_fn(**kwargs) -> str: # type: ignore return tool.apply_ex(log_call=True, catch_exceptions=True, **kwargs) return MCPTool( @@ -68,10 +64,23 @@ def make_tool( parameters=parameters, fn_metadata=func_arg_metadata, is_async=is_async, - context_kwarg="ctx", + context_kwarg=None, ) +_SERENA_AGENT_REGISTRY: dict[str, SerenaAgent] = {} +"""Map of project file paths to their corresponding SerenaAgent instances.""" + + +def _get_create_serena_agent(project_file_path: str) -> SerenaAgent: + if project_file_path not in _SERENA_AGENT_REGISTRY: + log.info(f"Creating new SerenaAgent for project file: {project_file_path}") + _SERENA_AGENT_REGISTRY[project_file_path] = SerenaAgent(project_file_path, start_language_server=True) + else: + log.info(f"Reusing existing SerenaAgent for project file: {project_file_path}") + return _SERENA_AGENT_REGISTRY[project_file_path] + + def create_mcp_server() -> FastMCP: argv = sys.argv[1:] if len(argv) != 1: @@ -79,19 +88,22 @@ def create_mcp_server() -> FastMCP: sys.exit(1) project_file_path = argv[0] - agent = SerenaAgent(project_file_path) + + agent = _get_create_serena_agent(project_file_path) @asynccontextmanager - async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[SerenaMCPRequestContext]: + async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[None]: """Manage server startup and shutdown lifecycle.""" - with agent.language_server_lifecycle_context(): - yield SerenaMCPRequestContext(agent=agent) + try: + yield + finally: + for agent in _SERENA_AGENT_REGISTRY.values(): + agent.language_server.stop() mcp_settings = Settings(lifespan=server_lifespan) mcp = FastMCP(**mcp_settings.model_dump()) for tool in agent.tools.values(): mcp._tool_manager._tools[tool.get_name()] = make_tool(tool) - return mcp