More flexible handling of enabled tools:

* In addition to exclusion, support notion of "optional tools" that
    are disabled by default and enabled only upon request
  * Unify the handling of inclusions/exclusions across
      - SerenaConfig
      - ProjectConfig
      - SerenaAgentContext
      - SerenaAgentMode
  * Make ToolRegistry a singleton and simplify all related operations

pyproject.toml

@@ -225,6 +225,7 @@ ignore = [
   "B904",     # forces use of raise from other_exception
   "RUF012",   # forbids mutable attributes as ClassVar
   "SIM117",   # forbids nested with statements
+  "C400",     # wants to unnecessarily force use of list comprehension
 ]
 unfixable = ["F841", "F601", "F602", "B018"]
 extend-fixable = ["F401", "B905", "W291"]

scripts/print_tool_overview.py

@@ -1,4 +1,4 @@
 from serena.agent import ToolRegistry
 
 if __name__ == "__main__":
-    ToolRegistry.print_tool_overview()
+    ToolRegistry().print_tool_overview()

src/serena/agent.py

@@ -23,7 +23,7 @@ from sensai.util.logging import LogTime
 
 from serena import serena_version
 from serena.config.context_mode import SerenaAgentContext, SerenaAgentMode
-from serena.config.serena_config import Project, SerenaConfig, get_serena_managed_dir
+from serena.config.serena_config import Project, SerenaConfig, ToolSet, get_serena_managed_dir
 from serena.constants import (
     SERENA_LOG_FORMAT,
 )
@@ -234,13 +234,13 @@ class SerenaAgent:
         self._context = context
 
         # instantiate all tool classes
-        self._all_tools: dict[type[Tool], Tool] = {tool_class: tool_class(self) for tool_class in ToolRegistry.get_all_tool_classes()}
+        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 context
-        # (which is fixed for the session)
-        excluded_tool_classes = set(self._context.get_excluded_tool_classes())
-        self._exposed_tools = {tc: t for tc, t in self._all_tools.items() if tc not in excluded_tool_classes}
+        # 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:
@@ -408,42 +408,20 @@ class SerenaAgent:
     def _update_active_tools(self) -> None:
         """
         Update the active tools based on context, modes, and project configuration.
-        All tool exclusions are merged together.
         """
-        excluded_tool_classes: set[type[Tool]] = set()
-        # modes
-        for mode in self._modes:
-            mode_excluded_tool_classes = mode.get_excluded_tool_classes()
-            if len(mode_excluded_tool_classes) > 0:
-                log.info(
-                    f"Mode {mode.name} excluded {len(mode_excluded_tool_classes)} tools: {', '.join([tool.get_name_from_cls() for tool in mode_excluded_tool_classes])}"
-                )
-                excluded_tool_classes.update(mode_excluded_tool_classes)
-        # context
-        context_excluded_tool_classes = self._context.get_excluded_tool_classes()
-        if len(context_excluded_tool_classes) > 0:
-            log.info(
-                f"Context {self._context.name} excluded {len(context_excluded_tool_classes)} tools: {', '.join([tool.get_name_from_cls() for tool in context_excluded_tool_classes])}"
-            )
-            excluded_tool_classes.update(context_excluded_tool_classes)
-        # project config
+        tool_set = ToolSet().apply(self.serena_config, self._context, *self._modes)
         if self._active_project is not None:
-            project_excluded_tool_classes = self._active_project.project_config.get_excluded_tool_classes()
-            if len(project_excluded_tool_classes) > 0:
-                log.info(
-                    f"Project {self._active_project.project_name} excluded {len(project_excluded_tool_classes)} tools: {', '.join([tool.get_name_from_cls() for tool in project_excluded_tool_classes])}"
-                )
-                excluded_tool_classes.update(project_excluded_tool_classes)
+            tool_set.apply(self._active_project.project_config)
             if self._active_project.project_config.read_only:
-                for tool_class in self._all_tools:
-                    if tool_class.can_edit():
-                        excluded_tool_classes.add(tool_class)
+                tool_set.exclude_editing_tools()
 
         self._active_tools = {
-            tool_class: tool_instance for tool_class, tool_instance in self._all_tools.items() if tool_class not in excluded_tool_classes
+            tool_class: tool_instance
+            for tool_class, tool_instance in self._all_tools.items()
+            if tool_set.includes_name(tool_instance.get_name())
         }
 
-        log.info(f"Active tools after all exclusions ({len(self._active_tools)}): {', '.join(self.get_active_tool_names())}")
+        log.info(f"Active tools ({len(self._active_tools)}): {', '.join(self.get_active_tool_names())}")
 
     def issue_task(self, task: Callable[[], Any], name: str | None = None) -> Future:
         """
@@ -648,7 +626,7 @@ class SerenaAgent:
         return self._all_tools[tool_class]  # type: ignore
 
     def print_tool_overview(self) -> None:
-        ToolRegistry.print_tool_overview(self._active_tools.values())
+        ToolRegistry().print_tool_overview(self._active_tools.values())
 
     def mark_file_modified(self, relativ_path: str) -> None:
         assert self.lines_read is not None

src/serena/config/context_mode.py

@@ -4,7 +4,7 @@ Context and Mode configuration loader
 
 import os
 from copy import copy
-from dataclasses import asdict, dataclass, field
+from dataclasses import asdict, dataclass
 from enum import Enum
 from pathlib import Path
 from typing import TYPE_CHECKING, Self
@@ -13,16 +13,17 @@ import yaml
 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
 
 if TYPE_CHECKING:
-    from serena.agent import Tool
+    pass
 
 log = logging.getLogger(__name__)
 
 
-@dataclass
-class SerenaAgentMode:
+@dataclass(kw_only=True)
+class SerenaAgentMode(ToolInclusionDefinition):
     """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.
@@ -31,7 +32,6 @@ class SerenaAgentMode:
     name: str
     prompt: str
     description: str = ""
-    excluded_tools: set[str] = field(default_factory=set)
 
     def to_json_dict(self) -> dict[str, str | list[str]]:
         result = asdict(self)
@@ -50,12 +50,6 @@ class SerenaAgentMode:
         if self.excluded_tools:
             print(" excluded tools:\n  " + ", ".join(sorted(self.excluded_tools)))
 
-    def get_excluded_tool_classes(self) -> list[type["Tool"]]:
-        """Get the list of tool classes that are excluded from the mode."""
-        from serena.agent import ToolRegistry
-
-        return [ToolRegistry.get_tool_class_by_name(tool_name) for tool_name in self.excluded_tools]
-
     @classmethod
     def from_yaml(cls, yaml_path: str | Path) -> Self:
         """Load a mode from a YAML file."""
@@ -93,8 +87,8 @@ class SerenaAgentMode:
             return cls.from_yaml(name_or_path)
 
 
-@dataclass
-class SerenaAgentContext(ToStringMixin):
+@dataclass(kw_only=True)
+class SerenaAgentContext(ToolInclusionDefinition, ToStringMixin):
     """Represents a context where the agent is operating (an IDE, a chat, etc.), typically read off a YAML file.
     An agent can only be in a single context at a time.
     The contexts cannot be changed after the agent is running.
@@ -103,7 +97,6 @@ class SerenaAgentContext(ToStringMixin):
     name: str
     prompt: str
     description: str = ""
-    excluded_tools: set[str] = field(default_factory=set)
 
     def _tostring_includes(self) -> list[str]:
         return ["name"]
@@ -119,12 +112,6 @@ class SerenaAgentContext(ToStringMixin):
         data["excluded_tools"] = set(data["excluded_tools"])
         return cls(**data)
 
-    def get_excluded_tool_classes(self) -> list[type["Tool"]]:
-        """Get the list of tool classes that are excluded from the context."""
-        from serena.agent import ToolRegistry
-
-        return [ToolRegistry.get_tool_class_by_name(tool_name) for tool_name in self.excluded_tools]
-
     @classmethod
     def from_yaml(cls, yaml_path: str | Path) -> Self:
         """Load a context from a YAML file."""

src/serena/config/serena_config.py

@@ -4,6 +4,7 @@ The Serena Model Context Protocol (MCP) Server
 
 import os
 import shutil
+from collections.abc import Iterable
 from copy import deepcopy
 from dataclasses import dataclass, field
 from functools import cached_property
@@ -22,7 +23,6 @@ from serena.constants import (
     SELENA_CONFIG_TEMPLATE_FILE,
     SERENA_MANAGED_DIR_NAME,
 )
-from serena.tools import Tool, ToolRegistry
 from serena.util.general import load_yaml, save_yaml
 from serena.util.inspection import determine_programming_language_composition
 from solidlsp.ls_config import Language
@@ -35,6 +35,53 @@ T = TypeVar("T")
 DEFAULT_TOOL_TIMEOUT: float = 240
 
 
+class ToolSet:
+    def __init__(self) -> None:
+        from serena.tools import ToolRegistry
+
+        self._tool_names = set(ToolRegistry().get_tool_names_default_enabled())
+
+    def apply(self, *tool_inclusion_definitions: "ToolInclusionDefinition") -> Self:
+        from serena.tools import ToolRegistry
+
+        registry = ToolRegistry()
+        for definition in tool_inclusion_definitions:
+            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)
+            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
+
+    def exclude_editing_tools(self) -> Self:
+        from serena.tools import ToolRegistry
+
+        registry = ToolRegistry()
+        for tool_name in list(self._tool_names):
+            if registry.get_tool_class_by_name(tool_name).can_edit():
+                self._tool_names.remove(tool_name)
+        return self
+
+    def get_tool_names(self) -> set[str]:
+        """
+        Returns the names of the tools that are currently included in the tool set.
+        """
+        return self._tool_names
+
+    def includes_name(self, tool_name: str) -> bool:
+        return tool_name in self._tool_names
+
+
+@dataclass
+class ToolInclusionDefinition:
+    excluded_tools: Iterable[str] = ()
+    included_optional_tools: Iterable[str] = ()
+
+
 class SerenaConfigError(Exception):
     pass
 
@@ -56,12 +103,11 @@ def is_running_in_docker() -> bool:
         return False
 
 
-@dataclass
-class ProjectConfig(ToStringMixin):
+@dataclass(kw_only=True)
+class ProjectConfig(ToolInclusionDefinition, ToStringMixin):
     project_name: str
     language: Language
     ignored_paths: list[str] = field(default_factory=list)
-    excluded_tools: set[str] = field(default_factory=set)
     read_only: bool = False
     ignore_all_files_in_gitignore: bool = True
     initial_prompt: str = ""
@@ -123,7 +169,8 @@ class ProjectConfig(ToStringMixin):
             project_name=project_name,
             language=language,
             ignored_paths=data.get("ignored_paths", []),
-            excluded_tools=set(data.get("excluded_tools", [])),
+            excluded_tools=data.get("excluded_tools", []),
+            included_optional_tools=data.get("included_optional_tools", []),
             read_only=data.get("read_only", False),
             ignore_all_files_in_gitignore=data.get("ignore_all_files_in_gitignore", True),
             initial_prompt=data.get("initial_prompt", ""),
@@ -148,9 +195,6 @@ class ProjectConfig(ToStringMixin):
             yaml_data["project_name"] = project_root.name
         return cls._from_dict(yaml_data)
 
-    def get_excluded_tool_classes(self) -> set[type["Tool"]]:
-        return set(ToolRegistry.get_tool_class_by_name(tool_name) for tool_name in self.excluded_tools)
-
 
 @dataclass
 class Project:
@@ -178,7 +222,7 @@ class Project:
 
 
 @dataclass(kw_only=True)
-class SerenaConfig:
+class SerenaConfig(ToolInclusionDefinition):
     """
     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.
@@ -287,6 +331,8 @@ class SerenaConfig:
         instance.web_dashboard_open_on_launch = loaded_commented_yaml.get("web_dashboard_open_on_launch", True)
         instance.tool_timeout = loaded_commented_yaml.get("tool_timeout", DEFAULT_TOOL_TIMEOUT)
         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", [])
 
         # re-save the configuration file if any migrations were performed
         if num_project_migrations > 0:

src/serena/resources/serena_config.template.yml

@@ -36,6 +36,12 @@ trace_lsp_communication: False
 tool_timeout: 240
 # timeout, in seconds, after which tool executions are terminated
 
+excluded_tools: []
+# list of tools to be globally excluded
+
+included_optional_tools: []
+# list of optional tools (which are disabled by default) to be included
+
 # 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,

src/serena/tools/jetbrains_tools.py

@@ -1,10 +1,10 @@
 import json
 
-from serena.tools import TOOL_DEFAULT_MAX_ANSWER_LENGTH, Tool
+from serena.tools import TOOL_DEFAULT_MAX_ANSWER_LENGTH, Tool, ToolMarkerOptional
 from serena.tools.jetbrains_plugin_client import JetBrainsPluginClient
 
 
-class JetBrainsFindSymbolTool(Tool):
+class JetBrainsFindSymbolTool(Tool, ToolMarkerOptional):
     """
     Performs a global (or local) search for symbols with/containing a given name/substring (optionally filtered by type).
     """
@@ -70,7 +70,7 @@ class JetBrainsFindSymbolTool(Tool):
         return self._limit_length(result, max_answer_chars)
 
 
-class JetBrainsFindSymbolReferencesTool(Tool):
+class JetBrainsFindSymbolReferencesTool(Tool, ToolMarkerOptional):
     """
     Finds symbols that reference the given symbol
     """

src/serena/tools/tools_base.py

@@ -2,8 +2,8 @@ import inspect
 import os
 import traceback
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Generator, Iterable
-from copy import copy
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Self, TypeVar
 
@@ -13,6 +13,7 @@ from sensai.util.string import dict_string
 
 from serena.prompt_factory import PromptFactory
 from serena.symbol import SymbolManager
+from serena.util.class_decorators import singleton
 from serena.util.inspection import iter_subclasses
 from solidlsp import SolidLanguageServer
 
@@ -72,6 +73,12 @@ class ToolMarkerDoesNotRequireActiveProject:
     pass
 
 
+class ToolMarkerOptional:
+    """
+    Marker class for optional tools that are disabled by default.
+    """
+
+
 class ToolInterface(ABC):
     """Protocol defining the complete interface that make_tool() expects from a tool."""
 
@@ -311,58 +318,44 @@ class EditedFileContext:
             # If they do not, we may have to add a call to notify it.
 
 
+@dataclass(kw_only=True)
+class RegisteredTool:
+    tool_class: type[Tool]
+    is_optional: bool
+    tool_name: str
+
+
+@singleton
 class ToolRegistry:
-    _tool_dict: dict[str, type[Tool]] | None = None
-    """maps tool name to the corresponding tool class"""
-
-    @staticmethod
-    def _iter_tool_classes() -> Generator[type[Tool], None, None]:
-        """
-        Iterate over Tool subclasses.
-        """
+    def __init__(self) -> None:
+        self._tool_dict: dict[str, RegisteredTool] = {}
         for cls in iter_subclasses(Tool):
-            if cls.__module__.startswith("serena.tools"):
-                yield cls
+            if not cls.__module__.startswith("serena.tools"):
+                continue
+            is_optional = issubclass(cls, ToolMarkerOptional)
+            name = cls.get_name_from_cls()
+            if name in self._tool_dict:
+                raise ValueError(f"Duplicate tool name found: {name}. Tool classes must have unique names.")
+            self._tool_dict[name] = RegisteredTool(tool_class=cls, is_optional=is_optional, tool_name=name)
 
-    @classmethod
-    def _get_tool_dict(cls) -> dict[str, type[Tool]]:
-        if cls._tool_dict is None:
-            cls._tool_dict = {}
-            for tool_class in cls._iter_tool_classes():
-                name = tool_class.get_name_from_cls()
-                if name in cls._tool_dict:
-                    raise ValueError(f"Duplicate tool name found: {name}. Tool classes must have unique names.")
-                cls._tool_dict[name] = tool_class
-        return cls._tool_dict
+    def get_tool_class_by_name(self, tool_name: str) -> type[Tool]:
+        return self._tool_dict[tool_name].tool_class
 
-    @classmethod
-    def get_tool_class_by_name(cls, tool_name: str) -> type[Tool]:
-        try:
-            return cls._get_tool_dict()[tool_name]
-        except KeyError as e:
-            available_tools = "\n".join(ToolRegistry.get_tool_names())
-            raise ValueError(f"Tool with name {tool_name} not found. Available tools:\n{available_tools}") from e
+    def get_all_tool_classes(self) -> list[type[Tool]]:
+        return list(t.tool_class for t in self._tool_dict.values())
 
-    @classmethod
-    def get_all_tool_classes(cls) -> list[type[Tool]]:
-        return list(cls._get_tool_dict().values())
+    def get_tool_names_default_enabled(self) -> list[str]:
+        """
+        :return: the list of tool names that are enabled by default (i.e. non-optional tools).
+        """
+        return [t.tool_name for t in self._tool_dict.values() if not t.is_optional]
 
-    @classmethod
-    def get_tool_names(cls) -> list[str]:
-        return list(cls._get_tool_dict().keys())
-
-    @classmethod
-    def tool_dict(cls) -> dict[str, type[Tool]]:
-        """Maps tool name to the corresponding tool class"""
-        return copy(cls._get_tool_dict())
-
-    @classmethod
-    def print_tool_overview(cls, tools: Iterable[type[Tool] | Tool] | None = None) -> None:
+    def print_tool_overview(self, tools: Iterable[type[Tool] | Tool] | None = None) -> None:
         """
         Print a summary of the tools. If no tools are passed, a summary of all tools is printed.
         """
         if tools is None:
-            tools = cls._get_tool_dict().values()
+            tools = [tool.tool_class for tool in self._tool_dict.values() if not tool.is_optional]
 
         tool_dict: dict[str, type[Tool] | Tool] = {}
         for tool_class in tools:
@@ -370,3 +363,6 @@ class ToolRegistry:
         for tool_name in sorted(tool_dict.keys()):
             tool_class = tool_dict[tool_name]
             print(f" * `{tool_name}`: {tool_class.get_tool_description().strip()}")
+
+    def is_valid_tool_name(self, tool_name: str) -> bool:
+        return tool_name in self._tool_dict

test/serena/test_mcp.py

@@ -271,7 +271,7 @@ def is_test_mock_class(tool_class: type) -> bool:
     )
 
 
-@pytest.mark.parametrize("tool_class", ToolRegistry.get_all_tool_classes())
+@pytest.mark.parametrize("tool_class", ToolRegistry().get_all_tool_classes())
 def test_make_tool_all_tools(tool_class) -> None:
     """Test that make_tool works for all tools in the codebase."""