Refactoring of MCP creation, introducing SerenaMCPFactory

This commit is contained in:
Dominik Jain committed 2025-06-21 20:18:11 +02:00
1 parent ac172dcab6
commit 75354a30ca
1 file changed
+160 -149
+160 -149
View File
@@ -23,12 +23,12 @@ from mcp.server.fastmcp.server import FastMCP, Settings
from mcp.server.fastmcp.tools.base import Tool as MCPTool
from pydantic_settings import SettingsConfigDict
from sensai.util import logging
from sensai.util.helper import mark_used
from serena.agent import (
ActivateProjectTool,
Project,
SerenaAgent,
SerenaConfig,
ToolInterface,
ToolRegistry,
create_serena_config,
@@ -66,152 +66,169 @@ class SerenaMCPRequestContext:
agent: SerenaAgent
def make_tool(
tool: ToolInterface,
) -> MCPTool:
func_name = tool.get_name()
func_doc = tool.get_apply_docstring() or ""
func_arg_metadata = tool.get_apply_fn_metadata()
is_async = False
parameters = func_arg_metadata.arg_model.model_json_schema()
class SerenaMCPFactory:
def __init__(self, context: str = DEFAULT_CONTEXT, project: str | None = None):
"""
:param context: The context name or path to context file
:param project: Either an absolute path to the project directory or a name of an already registered project.
If the project passed here hasn't been registered yet, it will be registered automatically and can be activated by its name
afterward.
"""
self.context = SerenaAgentContext.load(context)
docstring = docstring_parser.parse(func_doc)
# load project (if any)
project_instance: Project | None = None
if project is not None:
# Fail early if the project cannot be loaded
log.info(f"Will activate project {project} at mcp server startup")
try:
project_instance = Project.load(project)
except Exception as e:
log.error(f"Failed to load or generate project config for {project}: {e}")
raise
else:
log.warning("No project passed at startup, you will have to activate a project yourself (by asking the agent to do so).")
self.project_instance = project_instance
self.project = project
# Mount the tool description as a combination of the docstring description and
# the return value description, if it exists.
if docstring.description:
func_doc = f"{docstring.description.strip().strip('.')}."
else:
func_doc = ""
if docstring.returns and (docstring_returns_descr := docstring.returns.description):
# Only add a space before "Returns" if func_doc is not empty
prefix = " " if func_doc else ""
func_doc = f"{func_doc}{prefix}Returns {docstring_returns_descr.strip().strip('.')}."
self.active_tool_names = self._determine_active_tool_names(self.context, self.project_instance)
# Parse the parameter descriptions from the docstring and add pass its description
# to the parameters schema.
docstring_params = {param.arg_name: param for param in docstring.params}
parameters_properties: dict[str, dict[str, Any]] = parameters["properties"]
for parameter, properties in parameters_properties.items():
if (param_doc := docstring_params.get(parameter)) and param_doc.description:
param_desc = f"{param_doc.description.strip().strip('.') + '.'}"
properties["description"] = param_desc[0].upper() + param_desc[1:]
self.serena_agent_process: ProcessIsolatedSerenaAgent | None = None
self.serena_dashboard_process: ProcessIsolatedDashboard | None = None
def execute_fn(**kwargs) -> str: # type: ignore
return tool.apply_ex(log_call=True, catch_exceptions=True, **kwargs)
@staticmethod
def _determine_active_tool_names(context: SerenaAgentContext, project: Project) -> set[str]:
"""
Determine the names of tools that should be included in this session based on the context.
"""
tools_excluded_in_this_session = context.get_excluded_tool_classes()
return MCPTool(
fn=execute_fn,
name=func_name,
description=func_doc,
parameters=parameters,
fn_metadata=func_arg_metadata,
is_async=is_async,
context_kwarg=None,
annotations=None,
)
def create_mcp_server_and_agent(
project: str | None,
host: str = "0.0.0.0",
port: int = 8000,
context: str = DEFAULT_CONTEXT,
modes: Sequence[str] = DEFAULT_MODES,
enable_web_dashboard: bool | None = None,
enable_gui_log_window: bool | None = None,
log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] | None = None,
trace_lsp_communication: bool | None = None,
tool_timeout: float | None = None,
) -> tuple[FastMCP, ProcessIsolatedSerenaAgent]:
"""
Create an MCP server with process-isolated SerenaAgent to prevent asyncio contamination.
:param project: "Either an absolute path to the project directory or a name of an already registered project. "
"If the project passed here hasn't been registered yet, it will be registered automatically and can be activated by its name "
"afterwards.
:param host: The host to bind to
:param port: The port to bind to
:param context: The context name or path to context file
:param modes: List of mode names or paths to mode files
:param enable_web_dashboard: Whether to enable the web dashboard. If not specified, will take the value from the serena configuration.
:param enable_gui_log_window: Whether to enable the GUI log window. It currently does not work on macOS, and setting this to True will be ignored then.
If not specified, will take the value from the serena configuration.
:param log_level: Log level. If not specified, will take the value from the serena configuration.
:param trace_lsp_communication: Whether to trace the communication between Serena and the language servers.
This is useful for debugging language server issues.
:param tool_timeout: Timeout in seconds for tool execution. If not specified, will take the value from the serena configuration.
"""
mcp: FastMCP | None = None
context_instance = SerenaAgentContext.load(context)
modes_instances = [SerenaAgentMode.load(mode) for mode in modes]
is_ide_assistant = context_instance.name == RegisteredContext.IDE_ASSISTANT.value
project_instance: Project | None = None
if project is not None:
# Fail early if the project cannot be loaded
log.info(f"Will activate project {project} at mcp server startup")
try:
project_instance = Project.load(project)
except Exception as e:
log.error(f"Failed to load or generate project config for {project}: {e}")
raise
else:
log.warning("No project passed at startup, you will have to activate a project yourself (by asking the agent to do so).")
tools_excluded_in_this_session = context_instance.get_excluded_tool_classes()
if is_ide_assistant and project_instance is not None:
tools_excluded_in_this_session.extend(project_instance.project_config.get_excluded_tool_classes())
# if a project has been loaded, it will be activated at startup and in ide-assistant context,
# we assume that no other project will be activated in this session.
# Therefore, we exclude the activate project tool
tools_excluded_in_this_session.append(ActivateProjectTool)
tool_names_excluded_in_this_session = {tool.get_name_from_cls() for tool in tools_excluded_in_this_session}
is_ide_assistant = context.name == RegisteredContext.IDE_ASSISTANT.value
if is_ide_assistant and project is not None:
tools_excluded_in_this_session.extend(project.project_config.get_excluded_tool_classes())
tools_excluded_in_this_session.append(ActivateProjectTool)
all_tool_names = set(ToolRegistry.get_tool_names())
tool_names_included_in_this_session = all_tool_names - tool_names_excluded_in_this_session
tool_names_excluded_in_this_session = {tool.get_name_from_cls() for tool in tools_excluded_in_this_session}
try:
# Start agent and dashboard processes (the latter only if enabled)
serena_dashboard_process = None
serena_config = create_serena_config(
enable_web_dashboard=enable_web_dashboard,
enable_gui_log_window=enable_gui_log_window,
log_level=log_level,
trace_lsp_communication=trace_lsp_communication,
tool_timeout=tool_timeout,
)
if serena_config.web_dashboard:
serena_dashboard_process = ProcessIsolatedDashboard(tool_names=sorted(tool_names_included_in_this_session))
serena_agent_process = ProcessIsolatedSerenaAgent(
project=project, serena_config=serena_config, modes=modes_instances, context=context_instance
all_tool_names = set(ToolRegistry.get_tool_names())
tool_names_included_in_this_session = all_tool_names - tool_names_excluded_in_this_session
return tool_names_included_in_this_session
def _make_tool(self, tool: ToolInterface) -> MCPTool:
func_name = tool.get_name()
func_doc = tool.get_apply_docstring() or ""
func_arg_metadata = tool.get_apply_fn_metadata()
is_async = False
parameters = func_arg_metadata.arg_model.model_json_schema()
docstring = docstring_parser.parse(func_doc)
# Mount the tool description as a combination of the docstring description and
# the return value description, if it exists.
if docstring.description:
func_doc = f"{docstring.description.strip().strip('.')}."
else:
func_doc = ""
if docstring.returns and (docstring_returns_descr := docstring.returns.description):
# Only add a space before "Returns" if func_doc is not empty
prefix = " " if func_doc else ""
func_doc = f"{func_doc}{prefix}Returns {docstring_returns_descr.strip().strip('.')}."
# Parse the parameter descriptions from the docstring and add pass its description
# to the parameter schema.
docstring_params = {param.arg_name: param for param in docstring.params}
parameters_properties: dict[str, dict[str, Any]] = parameters["properties"]
for parameter, properties in parameters_properties.items():
if (param_doc := docstring_params.get(parameter)) and param_doc.description:
param_desc = f"{param_doc.description.strip().strip('.') + '.'}"
properties["description"] = param_desc[0].upper() + param_desc[1:]
def execute_fn(**kwargs) -> str: # type: ignore
return tool.apply_ex(log_call=True, catch_exceptions=True, **kwargs)
return MCPTool(
fn=execute_fn,
name=func_name,
description=func_doc,
parameters=parameters,
fn_metadata=func_arg_metadata,
is_async=is_async,
context_kwarg=None,
annotations=None,
)
except Exception as e:
show_fatal_exception_safe(e)
raise
def update_tools() -> None:
"""Update the tools in the MCP server - adapted for process isolation."""
nonlocal mcp, serena_agent_process
# Get tool names from agent
# Tools may change as a result of project or mode 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 potentially active tools at startup, unfortunately.
def _set_mcp_tools(self, mcp: FastMCP) -> None:
"""Update the tools in the MCP server"""
if mcp is not None:
mcp._tool_manager._tools = {}
for tool_name in tool_names_included_in_this_session:
process_isolated_tool = ProcessIsolatedTool(process_agent=serena_agent_process, tool_name=tool_name)
mcp_tool = make_tool(process_isolated_tool)
for tool_name in self.active_tool_names:
process_isolated_tool = ProcessIsolatedTool(process_agent=self.serena_agent_process, tool_name=tool_name)
mcp_tool = self._make_tool(process_isolated_tool)
mcp._tool_manager._tools[tool_name] = mcp_tool
def _instantiate_agent(self, serena_config: SerenaConfig, modes: list[SerenaAgentMode]) -> None:
if serena_config.web_dashboard:
self.serena_dashboard_process = ProcessIsolatedDashboard(tool_names=sorted(self.active_tool_names))
self.serena_agent_process = ProcessIsolatedSerenaAgent(
project=self.project, serena_config=serena_config, modes=modes, context=self.context
)
def create_mcp_server(
self,
host: str = "0.0.0.0",
port: int = 8000,
modes: Sequence[str] = DEFAULT_MODES,
enable_web_dashboard: bool | None = None,
enable_gui_log_window: bool | None = None,
log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] | None = None,
trace_lsp_communication: bool | None = None,
tool_timeout: float | None = None,
) -> FastMCP:
"""
Create an MCP server with process-isolated SerenaAgent to prevent asyncio contamination.
:param host: The host to bind to
:param port: The port to bind to
:param modes: List of mode names or paths to mode files
:param enable_web_dashboard: Whether to enable the web dashboard. If not specified, will take the value from the serena configuration.
:param enable_gui_log_window: Whether to enable the GUI log window. It currently does not work on macOS, and setting this to True will be ignored then.
If not specified, will take the value from the serena configuration.
:param log_level: Log level. If not specified, will take the value from the serena configuration.
:param trace_lsp_communication: Whether to trace the communication between Serena and the language servers.
This is useful for debugging language server issues.
:param tool_timeout: Timeout in seconds for tool execution. If not specified, will take the value from the serena configuration.
"""
try:
# Start agent and dashboard processes (the latter only if enabled)
serena_config = create_serena_config(
enable_web_dashboard=enable_web_dashboard,
enable_gui_log_window=enable_gui_log_window,
log_level=log_level,
trace_lsp_communication=trace_lsp_communication,
tool_timeout=tool_timeout,
)
modes_instances = [SerenaAgentMode.load(mode) for mode in modes]
self._instantiate_agent(serena_config, modes_instances)
except Exception as e:
show_fatal_exception_safe(e)
raise
# Override model_config to disable the use of `.env` files for reading settings, because user projects are likely to contain
# `.env` files (e.g. containing LOG_LEVEL) that are not supposed to override the MCP settings;
# retain only FASTMCP_ prefix for already set environment variables.
Settings.model_config = SettingsConfigDict(env_prefix="FASTMCP_")
mcp_settings = Settings(lifespan=self.server_lifespan, host=host, port=port)
mcp = FastMCP(**mcp_settings.model_dump())
return mcp
@asynccontextmanager
async def server_lifespan(mcp_server: FastMCP) -> AsyncIterator[None]:
async def server_lifespan(self, mcp_server: FastMCP) -> AsyncIterator[None]:
"""Manage server startup and shutdown lifecycle."""
mark_used(mcp_server)
def signal_handler(signum: int, frame: Any) -> None:
log.info(f"Received signal {signum} in main process")
@@ -220,6 +237,8 @@ def create_mcp_server_and_agent(
def force_exit() -> None:
time.sleep(2.0) # Wait 2 seconds for graceful shutdown
log.warning("Forcing exit after timeout")
# noinspection PyProtectedMember
# noinspection PyUnresolvedReferences
os._exit(1)
threading.Thread(target=force_exit, daemon=True).start()
@@ -228,12 +247,13 @@ def create_mcp_server_and_agent(
sigint_singal = signal.signal(signal.SIGINT, signal_handler)
sigterm_signal = signal.signal(signal.SIGTERM, signal_handler)
if serena_dashboard_process is not None:
if self.serena_dashboard_process is not None:
log.info("Starting dashboard process")
serena_dashboard_process.start()
self.serena_dashboard_process.start()
log.info("Starting serena agent process")
serena_agent_process.start()
update_tools()
self.serena_agent_process.start()
self._set_mcp_tools(mcp_server)
async def monitor_global_shutdown() -> None:
"""Monitor the global shutdown event and trigger local shutdown."""
@@ -264,23 +284,14 @@ def create_mcp_server_and_agent(
with contextlib.suppress(asyncio.CancelledError):
await monitor_task
serena_agent_process.stop()
if serena_dashboard_process is not None:
serena_dashboard_process.stop()
self.serena_agent_process.stop()
if self.serena_dashboard_process is not None:
self.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)
# Override model_config to disable the use of `.env` files for reading settings, because user projects are likely to contain
# `.env` files (e.g. containing LOG_LEVEL) that are not supposed to override the MCP settings;
# retain only FASTMCP_ prefix for already set environment variables.
Settings.model_config = SettingsConfigDict(env_prefix="FASTMCP_")
mcp_settings = Settings(lifespan=server_lifespan, host=host, port=port)
mcp = FastMCP(**mcp_settings.model_dump())
return mcp, serena_agent_process
class ProjectType(click.ParamType):
name = "[PROJECT_NAME|PROJECT_PATH]"
@@ -417,12 +428,12 @@ def start_mcp_server(
# This is for backward compatibility with the old CLI, should be removed in the future!
project_file = project_file_arg if project_file_arg is not None else project
mcp_factory = SerenaMCPFactory(context=context, project=project_file)
# Use process isolation by default to prevent asyncio event loop contamination
mcp_server, agent = create_mcp_server_and_agent(
project=project_file,
mcp_server = mcp_factory.create_mcp_server(
host=host,
port=port,
context=context,
modes=modes,
enable_web_dashboard=enable_web_dashboard,
enable_gui_log_window=enable_gui_log_window,
@@ -439,6 +450,6 @@ def start_mcp_server(
f"Used path: {project_file}"
)
log.info(f"Starting MCP server with config:\n{agent.serena_config}.\n Log level: {agent.serena_config.log_level}")
log.info("Starting MCP server ...")
mcp_server.run(transport=transport)