From f49802a040ae129194fd538c0afa9c56d957e2fe Mon Sep 17 00:00:00 2001 From: Dominik Jain Date: Sun, 6 Jul 2025 01:04:26 +0200 Subject: [PATCH] Allow to enable JetBrains tools via config entry 'jetbrains' which enables an internal SerenaAgentMode * Added notion of internal SerenaAgentModes, which cannot be enabled dynamically * Introduce 'jetbrains' SerenaAgentMode * Add 'jetbrains' option to SerenaConfig, which enables the respective mode at startup, appropriately configuring the set of exposed tools * Refactoring: Made ToolSet immutable; methods return new instances --- src/serena/agent.py | 29 +++++----- src/serena/config/context_mode.py | 15 ++++- src/serena/config/serena_config.py | 58 +++++++++++++++---- src/serena/constants.py | 1 + .../config/internal_modes/jetbrains.yml | 11 ++++ .../resources/serena_config.template.yml | 6 ++ 6 files changed, 94 insertions(+), 26 deletions(-) create mode 100644 src/serena/resources/config/internal_modes/jetbrains.yml diff --git a/src/serena/agent.py b/src/serena/agent.py index bc5e33a..f570ac3 100644 --- a/src/serena/agent.py +++ b/src/serena/agent.py @@ -237,11 +237,6 @@ class SerenaAgent: self._all_tools: dict[type[Tool], Tool] = {tool_class: tool_class(self) for tool_class in ToolRegistry().get_all_tool_classes()} tool_names = [tool.get_name_from_cls() for tool in self._all_tools.values()] - # determine the set exposed tools (which e.g. the MCP shall see), limited by the Serena config - # as well as the context (which is fixed for the session) - tool_set = ToolSet().apply(self.serena_config, self._context) - self._exposed_tools = {tc: t for tc, t in self._all_tools.items() if tool_set.includes_name(t.get_name())} - # If GUI log window is enabled, set the tool names for highlighting if self._gui_log_handler is not None: self._gui_log_handler.log_viewer.set_tool_names(tool_names) @@ -261,6 +256,16 @@ class SerenaAgent: log.info(f"Starting Serena server (version={serena_version()}, process id={os.getpid()}, parent process id={os.getppid()})") log.info("Configuration file: %s", self.serena_config.config_file_path) log.info("Available projects: {}".format(", ".join(self.serena_config.project_names))) + log.info(f"Loaded tools ({len(self._all_tools)}): {', '.join([tool.get_name_from_cls() for tool in self._all_tools.values()])}") + + # determine the base toolset defining the set of exposed tools (which e.g. the MCP shall see), + # limited by the Serena config, the context (which is fixed for the session) and JetBrains mode + tool_inclusion_definitions = [self.serena_config, self._context] + if self.serena_config.jetbrains: + tool_inclusion_definitions.append(SerenaAgentMode.from_name_internal("jetbrains")) + self._base_tool_set = ToolSet.default().apply(*tool_inclusion_definitions) + self._exposed_tools = {tc: t for tc, t in self._all_tools.items() if self._base_tool_set.includes_name(t.get_name())} + log.info(f"Number of exposed tools: {len(self._exposed_tools)}") # create executor for starting the language server and running tools in another thread # This executor is used to achieve linear task execution, so it is important to use a single-threaded executor. @@ -287,10 +292,6 @@ class SerenaAgent: modes = SerenaAgentMode.load_default_modes() self._modes = modes - # log tool information - log.info(f"Loaded tools ({len(self._all_tools)}): {', '.join([tool.get_name_from_cls() for tool in self._all_tools.values()])}") - log.info(f"Number of exposed tools given {self._context}: {len(self._exposed_tools)}") - self._active_tools: dict[type[Tool], Tool] = {} self._update_active_tools() @@ -407,13 +408,15 @@ class SerenaAgent: def _update_active_tools(self) -> None: """ - Update the active tools based on context, modes, and project configuration. + Update the active tools based on enabled modes and the active project. + The base tool set already takes the Serena configuration and the context into account + (as well as any internal modes that are not handled dynamically, such as JetBrains mode). """ - tool_set = ToolSet().apply(self.serena_config, self._context, *self._modes) + tool_set = self._base_tool_set.apply(*self._modes) if self._active_project is not None: - tool_set.apply(self._active_project.project_config) + tool_set = tool_set.apply(self._active_project.project_config) if self._active_project.project_config.read_only: - tool_set.exclude_editing_tools() + tool_set = tool_set.without_editing_tools() self._active_tools = { tool_class: tool_instance diff --git a/src/serena/config/context_mode.py b/src/serena/config/context_mode.py index 0ee3c26..dfd690a 100644 --- a/src/serena/config/context_mode.py +++ b/src/serena/config/context_mode.py @@ -14,7 +14,7 @@ from sensai.util import logging from sensai.util.string import ToStringMixin from serena.config.serena_config import ToolInclusionDefinition -from serena.constants import CONTEXT_YAMLS_DIR, DEFAULT_CONTEXT, DEFAULT_MODES, MODE_YAMLS_DIR +from serena.constants import CONTEXT_YAMLS_DIR, DEFAULT_CONTEXT, DEFAULT_MODES, INTERNAL_MODE_YAMLS_DIR, MODE_YAMLS_DIR if TYPE_CHECKING: pass @@ -23,7 +23,7 @@ log = logging.getLogger(__name__) @dataclass(kw_only=True) -class SerenaAgentMode(ToolInclusionDefinition): +class SerenaAgentMode(ToolInclusionDefinition, ToStringMixin): """Represents a mode of operation for the agent, typically read off a YAML file. An agent can be in multiple modes simultaneously as long as they are not mutually exclusive. The modes can be adjusted after the agent is running, for example for switching from planning to editing. @@ -33,6 +33,9 @@ class SerenaAgentMode(ToolInclusionDefinition): prompt: str description: str = "" + def _tostring_includes(self) -> list[str]: + return ["name"] + def to_json_dict(self) -> dict[str, str | list[str]]: result = asdict(self) result["excluded_tools"] = list(result["excluded_tools"]) @@ -69,6 +72,14 @@ class SerenaAgentMode(ToolInclusionDefinition): ) return cls.from_yaml(yaml_path) + @classmethod + def from_name_internal(cls, name: str) -> Self: + """Loads an internal Serena mode""" + yaml_path = os.path.join(INTERNAL_MODE_YAMLS_DIR, f"{name}.yml") + if not os.path.exists(yaml_path): + raise FileNotFoundError(f"Internal mode '{name}' not found in {INTERNAL_MODE_YAMLS_DIR}") + return cls.from_yaml(yaml_path) + @classmethod def list_registered_mode_names(cls) -> list[str]: """Names of all registered modes (from the corresponding YAML files in the serena repo).""" diff --git a/src/serena/config/serena_config.py b/src/serena/config/serena_config.py index 306393c..1f710bb 100644 --- a/src/serena/config/serena_config.py +++ b/src/serena/config/serena_config.py @@ -36,35 +36,60 @@ DEFAULT_TOOL_TIMEOUT: float = 240 class ToolSet: - def __init__(self) -> None: + def __init__(self, tool_names: set[str]) -> None: + self._tool_names = tool_names + + @classmethod + def default(cls) -> "ToolSet": + """ + :return: the default tool set, which contains all tools that are enabled by default + """ from serena.tools import ToolRegistry - self._tool_names = set(ToolRegistry().get_tool_names_default_enabled()) + return cls(set(ToolRegistry().get_tool_names_default_enabled())) - def apply(self, *tool_inclusion_definitions: "ToolInclusionDefinition") -> Self: + def apply(self, *tool_inclusion_definitions: "ToolInclusionDefinition") -> "ToolSet": + """ + :param tool_inclusion_definitions: the definitions to apply + :return: a new tool set with the definitions applied + """ from serena.tools import ToolRegistry registry = ToolRegistry() + tool_names = set(self._tool_names) for definition in tool_inclusion_definitions: + included_tools = [] + excluded_tools = [] for included_tool in definition.included_optional_tools: if not registry.is_valid_tool_name(included_tool): raise ValueError(f"Invalid tool name '{included_tool}' provided for inclusion") - self._tool_names.add(included_tool) + if included_tool not in tool_names: + tool_names.add(included_tool) + included_tools.append(included_tool) for excluded_tool in definition.excluded_tools: if not registry.is_valid_tool_name(excluded_tool): raise ValueError(f"Invalid tool name '{excluded_tool}' provided for exclusion") if excluded_tool in self._tool_names: - self._tool_names.remove(excluded_tool) - return self + tool_names.remove(excluded_tool) + excluded_tools.append(excluded_tool) + if included_tools: + log.info(f"{definition} included {len(included_tools)} tools: {', '.join(included_tools)}") + if excluded_tools: + log.info(f"{definition} excluded {len(excluded_tools)} tools: {', '.join(excluded_tools)}") + return ToolSet(tool_names) - def exclude_editing_tools(self) -> Self: + def without_editing_tools(self) -> "ToolSet": + """ + :return: a new tool set that excludes all tools that can edit + """ from serena.tools import ToolRegistry registry = ToolRegistry() - for tool_name in list(self._tool_names): + tool_names = set(self._tool_names) + for tool_name in self._tool_names: if registry.get_tool_class_by_name(tool_name).can_edit(): - self._tool_names.remove(tool_name) - return self + tool_names.remove(tool_name) + return ToolSet(tool_names) def get_tool_names(self) -> set[str]: """ @@ -115,6 +140,9 @@ class ProjectConfig(ToolInclusionDefinition, ToStringMixin): SERENA_DEFAULT_PROJECT_FILE = "project.yml" + def _tostring_includes(self) -> list[str]: + return ["project_name"] + @classmethod def autogenerate(cls, project_root: str | Path, project_name: str | None = None, save_to_disk: bool = True) -> Self: """ @@ -222,7 +250,7 @@ class Project: @dataclass(kw_only=True) -class SerenaConfig(ToolInclusionDefinition): +class SerenaConfig(ToolInclusionDefinition, ToStringMixin): """ Holds the Serena agent configuration, which is typically loaded from a YAML configuration file (when instantiated via :method:`from_config_file`), which is updated when projects are added or removed. @@ -242,10 +270,17 @@ class SerenaConfig(ToolInclusionDefinition): the path to the configuration file to which updates of the configuration shall be saved; if None, the configuration is not saved to disk """ + jetbrains: bool = False + """ + whether to apply JetBrains mode + """ CONFIG_FILE = "serena_config.yml" CONFIG_FILE_DOCKER = "serena_config.docker.yml" # Docker-specific config file; auto-generated if missing, mounted via docker-compose for user customization + def _tostring_includes(self) -> list[str]: + return ["config_file_path"] + @classmethod def _generate_config_file(cls, config_file_path: str) -> None: """ @@ -333,6 +368,7 @@ class SerenaConfig(ToolInclusionDefinition): instance.trace_lsp_communication = loaded_commented_yaml.get("trace_lsp_communication", False) instance.excluded_tools = loaded_commented_yaml.get("excluded_tools", []) instance.included_optional_tools = loaded_commented_yaml.get("included_optional_tools", []) + instance.jetbrains = loaded_commented_yaml.get("jetbrains", False) # re-save the configuration file if any migrations were performed if num_project_migrations > 0: diff --git a/src/serena/constants.py b/src/serena/constants.py index 447364b..668c2f9 100644 --- a/src/serena/constants.py +++ b/src/serena/constants.py @@ -7,6 +7,7 @@ REPO_ROOT = str(_repo_root_path) PROMPT_TEMPLATES_DIR = str(_serena_pkg_path / "resources" / "config" / "prompt_templates") CONTEXT_YAMLS_DIR = str(_serena_pkg_path / "resources" / "config" / "contexts") MODE_YAMLS_DIR = str(_serena_pkg_path / "resources" / "config" / "modes") +INTERNAL_MODE_YAMLS_DIR = str(_serena_pkg_path / "resources" / "config" / "internal_modes") SERENA_DASHBOARD_DIR = str(_serena_pkg_path / "resources" / "dashboard") SERENA_ICON_DIR = str(_serena_pkg_path / "resources" / "icons") diff --git a/src/serena/resources/config/internal_modes/jetbrains.yml b/src/serena/resources/config/internal_modes/jetbrains.yml new file mode 100644 index 0000000..fdeb41d --- /dev/null +++ b/src/serena/resources/config/internal_modes/jetbrains.yml @@ -0,0 +1,11 @@ +description: JetBrains tools replace language server-based tools +prompt: | + You have access to the very powerful JetBrains tools for symbolic operations: + * `jet_brains_find_symbol` replaces `find_symbol` + * `jet_brains_find_symbol_references` replaces `find_referencing_symbols` +excluded_tools: + - find_symbol + - find_referencing_symbols +included_optional_tools: + - jet_brains_find_symbol + - jet_brains_find_symbol_references \ No newline at end of file diff --git a/src/serena/resources/serena_config.template.yml b/src/serena/resources/serena_config.template.yml index e0a6a09..fcffbc4 100644 --- a/src/serena/resources/serena_config.template.yml +++ b/src/serena/resources/serena_config.template.yml @@ -42,6 +42,12 @@ excluded_tools: [] included_optional_tools: [] # list of optional tools (which are disabled by default) to be included +jetbrains: False +# whether to enable JetBrains mode and use tools based on the Serena JetBrains IDE plugin +# instead of language server-based tools +# NOTE: The plugin is yet unreleased. This is for Serena developers only. + + # MANAGED BY SERENA, KEEP AT THE BOTTOM OF THE YAML AND DON'T EDIT WITHOUT NEED # The list of registered projects. # To add a project, within a chat, simply ask Serena to "activate the project /path/to/project" or,