diff --git a/packages/dbgpt-core/src/dbgpt/agent/resource/connector/manager.py b/packages/dbgpt-core/src/dbgpt/agent/resource/connector/manager.py index 8aa7f1481..c981d9384 100644 --- a/packages/dbgpt-core/src/dbgpt/agent/resource/connector/manager.py +++ b/packages/dbgpt-core/src/dbgpt/agent/resource/connector/manager.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import re import secrets from enum import Enum from typing import TYPE_CHECKING, Any, Dict, List, Optional @@ -34,6 +35,20 @@ class ConnectorStatus(str, Enum): disconnected = "disconnected" +def _slugify(text: str) -> str: + """Convert *text* to a lowercase ASCII slug. + + Rules: lowercase, replace whitespace and underscores with ``-``, + strip characters outside ``[a-z0-9-]``, collapse consecutive dashes. + """ + s = text.lower() + s = re.sub(r"[\s_]+", "-", s) + s = re.sub(r"[^a-z0-9-]", "", s) + s = re.sub(r"-{2,}", "-", s) + s = s.strip("-") + return s + + class ConnectorManager(BaseComponent): """Core manager that aggregates catalog, credential store and confirmation system. @@ -79,6 +94,77 @@ class ConnectorManager(BaseComponent): self._salts: Dict[str, str] = {} self._connector_types: Dict[str, str] = {} + # ------------------------------------------------------------------ + # Tool prefix helpers + # ------------------------------------------------------------------ + + def _has_multiple_instances_of_type(self, connector_type: str) -> bool: + """True if at least one other connector of the same type is already active.""" + return any(t == connector_type for cid, t in self._connector_types.items()) + + def compute_tool_prefix( + self, + connector_type: str, + display_name: str, + connector_id: str, + ) -> str: + """Return the prefix to prepend to every tool name in this connector's pack. + + Prefix rules: + - Built-in, single instance: ``{connector_type}`` + - Built-in, multiple instances: ``{connector_type}-{slug(display_name)}`` + - Custom (type starts with ``custom_``): ``{slug(display_name)}`` + - Fallback when slug is empty: first 6 chars of *connector_id* + """ + is_custom = connector_type.startswith("custom_") + slug = _slugify(display_name) + + if is_custom: + return slug if slug else connector_id[:6] + + if self._has_multiple_instances_of_type(connector_type): + suffix = slug if slug else connector_id[:6] + return f"{connector_type}-{suffix}" + + return connector_type + + def _apply_tool_prefix( + self, pack: "MCPToolPack", prefix: str, display_name: str + ) -> None: + """Rename every tool in *pack* to add the prefix and update tool_server_map.""" + new_server_map: Dict[str, str] = {} + default_server = ( + pack._mcp_servers + if isinstance(pack._mcp_servers, str) + else (pack._mcp_servers[0] if pack._mcp_servers else "") + ) + new_resources: Dict[str, Any] = {} + + for tool in list(pack.sub_resources): + if not isinstance(tool, BaseTool): + continue + old_name = tool.name + prefixed = f"{prefix}_{old_name}" + # Skip if already prefixed (idempotency) + if old_name.startswith(f"{prefix}_"): + new_server_map[old_name] = pack.tool_server_map.get( + old_name, default_server + ) + new_resources[old_name] = tool + continue + new_name = prefixed + new_server_map[new_name] = pack.tool_server_map.get( + old_name, default_server + ) + # Rename tool internals + tool._name = new_name + if not tool._description.endswith(f"(via {display_name})"): + tool._description = f"{tool._description} (via {display_name})" + new_resources[new_name] = tool + + pack.tool_server_map = new_server_map + pack._resources = new_resources + def init_app(self, system_app: SystemApp) -> None: """Bind the manager to *system_app*. @@ -171,8 +257,11 @@ class ConnectorManager(BaseComponent): try: await pack.preload_resource() + prefix = self.compute_tool_prefix(connector_type, pack_name, connector_id) + self._apply_tool_prefix(pack, prefix, pack_name) self._active_packs[connector_id] = pack self._statuses[connector_id] = ConnectorStatus.active + self._connector_types[connector_id] = connector_type tool_count = sum( 1 for tool in pack.sub_resources if isinstance(tool, BaseTool) ) @@ -193,7 +282,6 @@ class ConnectorManager(BaseComponent): exc_info=True, ) - self._connector_types[connector_id] = connector_type return connector_id async def remove_connector(self, connector_id: str) -> None: @@ -213,6 +301,7 @@ class ConnectorManager(BaseComponent): self._statuses[connector_id] = ConnectorStatus.disconnected self._active_packs.pop(connector_id, None) self._salts.pop(connector_id, None) + self._connector_types.pop(connector_id, None) # TODO(T6): cancel APScheduler jobs associated with connector_id logger.info("Connector id=%s removed", connector_id) diff --git a/packages/dbgpt-core/src/dbgpt/agent/resource/connector/tests/test_manager.py b/packages/dbgpt-core/src/dbgpt/agent/resource/connector/tests/test_manager.py index a950ad4fb..1b17c5259 100644 --- a/packages/dbgpt-core/src/dbgpt/agent/resource/connector/tests/test_manager.py +++ b/packages/dbgpt-core/src/dbgpt/agent/resource/connector/tests/test_manager.py @@ -171,3 +171,119 @@ def test_get_catalog_after_load_returns_entries(tmp_path): assert len(catalog.list()) == 1 assert catalog.list()[0].type == "test_conn" assert catalog.list()[0].display_name == "Test Connector" + + +# --------------------------------------------------------------------------- +# Task B – tool auto-prefix tests +# --------------------------------------------------------------------------- + + +def test_compute_tool_prefix_builtin_single(): + """Built-in type with no prior instances => prefix is just the type.""" + manager = ConnectorManager() + prefix = manager.compute_tool_prefix("yuque", "My Yuque", "abc123") + assert prefix == "yuque" + + +def test_compute_tool_prefix_builtin_multiple(): + """Built-in type with an existing instance => prefix includes slug of display_name.""" + manager = ConnectorManager() + manager._connector_types["other-id"] = "yuque" + prefix = manager.compute_tool_prefix("yuque", "Team Yuque", "new-id") + assert prefix == "yuque-team-yuque" + + +def test_compute_tool_prefix_custom(): + """Custom connector type => prefix is slug of display_name only.""" + manager = ConnectorManager() + prefix = manager.compute_tool_prefix("custom_mcp", "My Internal KB", "xyz789") + assert prefix == "my-internal-kb" + + +def test_compute_tool_prefix_slug_fallback(): + """When slug is empty, fall back to first 6 chars of connector_id.""" + manager = ConnectorManager() + prefix = manager.compute_tool_prefix("custom_mcp", "!!!", "abc1234567890") + assert prefix == "abc123" + + +def test_apply_tool_prefix_renames_and_updates_map(): + """_apply_tool_prefix renames tools, updates tool_server_map, and annotates descriptions.""" + manager = ConnectorManager() + + pack = MCPToolPack(mcp_servers="http://example.com/sse", name="Yuque") + pack.add_command( + command_label="Create a doc", + command_name="create_doc", + args={}, + function=lambda: None, + ) + pack.add_command( + command_label="Delete a doc", + command_name="delete_doc", + args={}, + function=lambda: None, + ) + pack.tool_server_map = { + "create_doc": "http://example.com/sse", + "delete_doc": "http://example.com/sse", + } + + manager._apply_tool_prefix(pack, "yuque", "Yuque") + + tools = pack.sub_resources + assert tools[0].name == "yuque_create_doc" + assert tools[1].name == "yuque_delete_doc" + assert pack.tool_server_map == { + "yuque_create_doc": "http://example.com/sse", + "yuque_delete_doc": "http://example.com/sse", + } + assert tools[0].description.endswith("(via Yuque)") + assert tools[1].description.endswith("(via Yuque)") + + +def test_apply_tool_prefix_idempotent(): + """Calling _apply_tool_prefix twice with the same prefix doesn't double-prefix.""" + manager = ConnectorManager() + + pack = MCPToolPack(mcp_servers="http://example.com/sse", name="Yuque") + pack.add_command( + command_label="Create a doc", + command_name="create_doc", + args={}, + function=lambda: None, + ) + pack.tool_server_map = { + "create_doc": "http://example.com/sse", + } + + manager._apply_tool_prefix(pack, "yuque", "Yuque") + manager._apply_tool_prefix(pack, "yuque", "Yuque") + + tools = pack.sub_resources + assert tools[0].name == "yuque_create_doc" + assert pack.tool_server_map == { + "yuque_create_doc": "http://example.com/sse", + } + assert tools[0].description.count("(via Yuque)") == 1 + + +def test_remove_connector_clears_connector_types(): + """After remove_connector, _connector_types should not retain the entry. + + Otherwise compute_tool_prefix would treat re-created connectors as multi-instance. + """ + manager = ConnectorManager() + cid = "ghost-id" + manager._connector_types[cid] = "yuque" + manager._active_packs[cid] = MCPToolPack( + mcp_servers="http://example.com/sse", name="Ghost" + ) + manager._statuses[cid] = ConnectorStatus.active + + import asyncio + + asyncio.run(manager.remove_connector(cid)) + + assert cid not in manager._connector_types + assert not manager._has_multiple_instances_of_type("yuque")