Merge pull request #1106 from oraios/fix-exposed-tools

Refactor active mode and tool computation, fixing modes not being able to activate optional tools

src/serena/agent.py

@@ -20,6 +20,7 @@ from serena.config.context_mode import SerenaAgentContext, SerenaAgentMode
 from serena.config.serena_config import (
     LanguageBackend,
     ModeSelectionDefinition,
+    NamedToolInclusionDefinition,
     RegisteredProject,
     SerenaConfig,
     ToolInclusionDefinition,
@@ -89,6 +90,9 @@ class ToolSet:
     def __init__(self, tool_names: set[str]) -> None:
         self._tool_names = tool_names
 
+    def __len__(self) -> int:
+        return len(self._tool_names)
+
     @classmethod
     def default(cls) -> "ToolSet":
         """
@@ -172,6 +176,9 @@ class ToolSet:
     def includes_name(self, tool_name: str) -> bool:
         return tool_name in self._tool_names
 
+    def to_available_tools(self, all_tools: dict[type[Tool], Tool]) -> AvailableTools:
+        return AvailableTools([t for t in all_tools.values() if self.includes_name(t.get_name())])
+
 
 class ActiveModes:
     def __init__(self) -> None:
@@ -321,30 +328,6 @@ class SerenaAgent:
         else:
             log.info(f"Using language backend from global configuration: {self._language_backend.name}")
 
-        # determine the base toolset defining the set of exposed tools (which e.g. the MCP shall see),
-        # determined by the
-        #   * dashboard availability/opening on launch
-        #   * Serena config
-        #   * the context (which is fixed for the session)
-        #   * single-project mode reductions (if applicable)
-        #   * JetBrains mode
-        tool_inclusion_definitions: list[ToolInclusionDefinition] = []
-        if (
-            self.serena_config.web_dashboard
-            and not self.serena_config.web_dashboard_open_on_launch
-            and not self.serena_config.gui_log_window
-        ):
-            tool_inclusion_definitions.append(ToolInclusionDefinition(included_optional_tools=[OpenDashboardTool.get_name_from_cls()]))
-        tool_inclusion_definitions.append(self.serena_config)
-        tool_inclusion_definitions.append(self._context)
-        if self._context.single_project:
-            tool_inclusion_definitions.extend(self._single_project_context_tool_inclusion_definitions(registered_project_to_activate))
-        if self._language_backend == LanguageBackend.JETBRAINS:
-            tool_inclusion_definitions.append(SerenaAgentMode.from_name_internal("jetbrains"))
-        self._base_tool_set = ToolSet.default().apply(*tool_inclusion_definitions)
-        self._exposed_tools = AvailableTools([t for t in self._all_tools.values() 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
         self._task_executor = TaskExecutor("SerenaAgentTaskExecutor")
@@ -353,19 +336,29 @@ class SerenaAgent:
         self.prompt_factory = SerenaPromptFactory()
         self._project_activation_callback = project_activation_callback
 
-        # activate a project configuration (if provided or if there is only a single project available)
+        # activate the given project (if any), also updating the active modes
+        # Note: We cannot update the active tools yet, because the base toolset has not been computed yet
+        #       (and its computation depends on the active project)
+        self._active_modes: ActiveModes
+        self._mode_overrides = modes
         if project is not None:
             try:
-                self.activate_project_from_path_or_name(project, update_modes_and_tools=False)
+                self.activate_project_from_path_or_name(project, update_active_modes=True, update_active_tools=False)
             except Exception as e:
                 log.error(f"Error activating project '{project}' at startup: {e}", exc_info=e)
+        else:
+            self._update_active_modes()
 
-        # update active modes and active tools (considering the active project, if any)
-        # declared attributes are set in the call to _update_active_modes_and_tools()
-        self._mode_overrides = modes
-        self._active_modes: ActiveModes
+        # determine the base toolset defining the set of exposed tools (which e.g. the MCP shall see),
+        self._base_toolset = self._create_base_toolset(
+            self.serena_config, self._language_backend, self._context, self._active_modes, self._active_project
+        )
+        self._exposed_tools = self._base_toolset.to_available_tools(self._all_tools)
+        log.info(f"Number of exposed tools: {len(self._exposed_tools)}")
+
+        # update the active tools (considering the active project, if any)
         self._active_tools: AvailableTools
-        self._update_active_modes_and_tools()
+        self._update_active_tools()
 
         # start the dashboard (web frontend), registering its log handler
         # should be the last thing to happen in the initialization since the dashboard
@@ -386,6 +379,73 @@ class SerenaAgent:
             if self._gui_log_viewer is not None:
                 self._gui_log_viewer.set_dashboard_url(dashboard_url)
 
+    @classmethod
+    def _create_base_toolset(
+        cls,
+        serena_config: SerenaConfig,
+        language_backend: LanguageBackend,
+        context: SerenaAgentContext,
+        modes: ActiveModes,
+        project: Project | None,
+    ) -> ToolSet:
+        """
+        Determines the base toolset defining the set of exposed tools (which e.g. the MCP shall see).
+        It depends on ...
+           * dashboard availability/opening on launch
+           * Serena config
+           * the context (which is fixed for the session)
+           * the optional tools enabled by initial modes
+           * single-project mode reductions (if applicable)
+           * JetBrains mode
+        """
+        # determine whether to include the OpenDashboardTool based on the Serena configuration
+        tool_inclusion_definitions: list[ToolInclusionDefinition] = []
+        if serena_config.web_dashboard and not serena_config.web_dashboard_open_on_launch and not serena_config.gui_log_window:
+            tool_inclusion_definitions.append(
+                NamedToolInclusionDefinition(name="OpenDashboard", included_optional_tools=[OpenDashboardTool.get_name_from_cls()])
+            )
+
+        # consider Serena configuration and the active context
+        tool_inclusion_definitions.append(serena_config)
+        tool_inclusion_definitions.append(context)
+
+        # consider modes
+        # Since modes can be dynamically turned on and off, we don't include their definitions directly,
+        # but for the initially active modes, we make sure that the tools they enable are included.
+        for mode in modes.get_modes():
+            tool_inclusion_definitions.append(
+                NamedToolInclusionDefinition(
+                    name=f"InitialModeInclusions[{mode.name}]", included_optional_tools=mode.included_optional_tools
+                )
+            )
+
+        # When in a single-project context, the agent is assumed to work on a single project, and we thus
+        # want to apply that project's tool exclusions/inclusions from the get-go, limiting the set
+        # of tools that will be exposed to the client.
+        # Furthermore, we disable tools that are only relevant for project activation.
+        # So if the project exists, we apply all the aforementioned exclusions.
+        if context.single_project and project is not None:
+            log.info(
+                "Applying tool inclusion/exclusion definitions for single-project context based on project '%s'",
+                project.project_name,
+            )
+            tool_inclusion_definitions.append(
+                NamedToolInclusionDefinition(
+                    name="SingleProjectExclusions",
+                    excluded_tools=[ActivateProjectTool.get_name_from_cls(), GetCurrentConfigTool.get_name_from_cls()],
+                )
+            )
+            tool_inclusion_definitions.append(project.project_config)
+
+        # enabled the internal 'jetbrains' mode for the JetBrains backend
+        if language_backend == LanguageBackend.JETBRAINS:
+            tool_inclusion_definitions.append(SerenaAgentMode.from_name_internal("jetbrains"))
+
+        # compute the resulting tool set
+        base_toolset = ToolSet.default().apply(*tool_inclusion_definitions)
+        log.info(f"Number of exposed tools: {len(base_toolset)}")
+        return base_toolset
+
     def get_language_backend(self) -> LanguageBackend:
         return self._language_backend
 
@@ -437,36 +497,6 @@ class SerenaAgent:
                 os.environ["COMSPEC"] = ""  # force use of default shell
                 log.info("Adjusting COMSPEC environment variable to use the default shell instead of '%s'", comspec)
 
-    def _single_project_context_tool_inclusion_definitions(
-        self, registered_project: RegisteredProject | None
-    ) -> list[ToolInclusionDefinition]:
-        """
-        When in a single-project context, the agent is assumed to work on a single project, and we thus
-        want to apply that project's tool exclusions/inclusions from the get-go, limiting the set
-        of tools that will be exposed to the client.
-        Furthermore, we disable tools that are only relevant for project activation.
-        So if the project exists, we apply all the aforementioned exclusions.
-
-        NOTE: This uses an existing registered project (instead of project auto-generation),
-              because the result must be returned instantaneously
-              (project generation could take too much time), so as not to delay MCP server startup
-              and provide responses to the client immediately.
-
-        :param registered_project: the registered project or None
-        :return: the list of tool inclusion definitions to apply for a single-project context, based on the given registered project (if any)
-        """
-        tool_inclusion_definitions = []
-        if registered_project is not None:
-            log.info(
-                "Applying tool inclusion/exclusion definitions for single-project context based on project '%s'",
-                registered_project.project_name,
-            )
-            tool_inclusion_definitions.append(
-                ToolInclusionDefinition(excluded_tools=[ActivateProjectTool.get_name_from_cls(), GetCurrentConfigTool.get_name_from_cls()])
-            )
-            tool_inclusion_definitions.append(registered_project.project_config)
-        return tool_inclusion_definitions
-
     def record_tool_usage(self, input_kwargs: dict, tool_result: str | dict, tool: Tool) -> None:
         """
         Record the usage of a tool with the given input and output strings if tool usage statistics recording is enabled.
@@ -548,7 +578,8 @@ class SerenaAgent:
         :param mode_names: List of mode names or paths to use
         """
         self._mode_overrides = ModeSelectionDefinition(default_modes=mode_names)
-        self._update_active_modes_and_tools()
+        self._update_active_modes()
+        self._update_active_tools()
 
         log.info(f"Set modes to {[mode.name for mode in self.get_active_modes()]}")
 
@@ -580,13 +611,11 @@ class SerenaAgent:
         log.info("System prompt:\n%s", system_prompt)
         return system_prompt
 
-    def _update_active_modes_and_tools(self) -> None:
+    def _update_active_modes(self) -> None:
         """
-        Updates the active modes and, subsequently, the active tools based on the 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).
+        Updates the active modes based on the Serena configuration, the active project configuration (if any),
+        and mode overrides (if any).
         """
-        # determine active modes from serena config, active project, and mode overrides
         self._active_modes = ActiveModes()
         self._active_modes.apply(self.serena_config)
         if self._active_project:
@@ -594,8 +623,14 @@ class SerenaAgent:
         if self._mode_overrides:
             self._active_modes.apply(self._mode_overrides)
 
+    def _update_active_tools(self) -> None:
+        """
+        Updates the active tools based on the active modes and the active project.
+        The base tool set already takes the Serena configuration and the context into account
+        (as well as many other aspects, such as JetBrains mode).
+        """
         # apply modes
-        tool_set = self._base_tool_set.apply(*self._active_modes.get_modes())
+        tool_set = self._base_toolset.apply(*self._active_modes.get_modes())
 
         # apply active project configuration (if any)
         if self._active_project is not None:
@@ -603,11 +638,18 @@ class SerenaAgent:
             if self._active_project.project_config.read_only:
                 tool_set = tool_set.without_editing_tools()
 
-        tools = [t for t in self._all_tools.values() if tool_set.includes_name(t.get_name())]
-        self._active_tools = AvailableTools(tools)
-
+        self._active_tools = tool_set.to_available_tools(self._all_tools)
         log.info(f"Active tools ({len(self._active_tools)}): {', '.join(self._active_tools.tool_names)}")
 
+        # check if a tool was activated that is not in the exposed tool set and issue a warning if so
+        active_tools_not_exposed = set(self._active_tools.tool_names) - set(self._exposed_tools.tool_names)
+        if active_tools_not_exposed:
+            log.warning(
+                "The following active tools are not in the exposed tool set and thus won't be available to clients:\n"
+                f"{active_tools_not_exposed}\n"
+                "Consider adjusting your configuration to include these tools if you want to use them."
+            )
+
     def issue_task(
         self, task: Callable[[], T], name: str | None = None, logged: bool = True, timeout: float | None = None
     ) -> TaskExecutor.Task[T]:
@@ -642,7 +684,7 @@ class SerenaAgent:
         """
         return self._language_backend == LanguageBackend.LSP
 
-    def _activate_project(self, project: Project, update_modes_and_tools: bool = True) -> None:
+    def _activate_project(self, project: Project, update_active_modes: bool = True, update_active_tools: bool = True) -> None:
         log.info(f"Activating {project.project_name} at {project.project_root}")
 
         # Check if the project requires a different language backend than the one initialized at startup
@@ -657,8 +699,11 @@ class SerenaAgent:
 
         self._active_project = project
 
-        if update_modes_and_tools:
-            self._update_active_modes_and_tools()
+        if update_active_modes:
+            self._update_active_modes()
+
+        if update_active_tools:
+            self._update_active_tools()
 
         def init_language_server_manager() -> None:
             # start the language server
@@ -672,14 +717,14 @@ class SerenaAgent:
         if self._project_activation_callback is not None:
             self._project_activation_callback()
 
-    def activate_project_from_path_or_name(self, project_root_or_name: str, update_modes_and_tools: bool = True) -> Project:
+    def activate_project_from_path_or_name(
+        self, project_root_or_name: str, update_active_modes: bool = True, update_active_tools: bool = True
+    ) -> Project:
         """
         Activate a project from a path or a name.
         If the project was already registered, it will just be activated.
         If the argument is a path at which no Serena project previously existed, the project will be created beforehand.
         Raises ProjectNotFoundError if the project could neither be found nor created.
-
-        :return: the project instance
         """
         project_instance: Project | None = self.serena_config.get_project(project_root_or_name)
         if project_instance is not None:
@@ -694,7 +739,7 @@ class SerenaAgent:
                 f"Existing project names: {self.serena_config.project_names}"
             )
 
-        self._activate_project(project_instance, update_modes_and_tools=update_modes_and_tools)
+        self._activate_project(project_instance, update_active_modes=update_active_modes, update_active_tools=update_active_tools)
 
         return project_instance
 

src/serena/config/serena_config.py

@@ -128,6 +128,14 @@ class ToolInclusionDefinition:
         return num_fixed > 0
 
 
+@dataclass
+class NamedToolInclusionDefinition(ToolInclusionDefinition):
+    name: str | None = None
+
+    def __str__(self) -> str:
+        return f"ToolInclusionDefinition[{self.name}]"
+
+
 @dataclass
 class ModeSelectionDefinition:
     base_modes: Sequence[str] | None = None