mirror of
https://github.com/tiennm99/serena.git
synced 2026-10-14 18:19:14 +00:00
Refactoring of MCP creation, introducing SerenaMCPFactory
This commit is contained in:
1 parent
ac172dcab6
commit
75354a30ca
1 file changed
+160
-149
+160
-149
@@ -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)
|
||||
Reference in new issue
Block a user