From babd9754654f6d02ddfb00dc1c5519215c3481f1 Mon Sep 17 00:00:00 2001 From: "alan.cl" <1165243776@qq.com> Date: Tue, 26 May 2026 12:06:12 +0800 Subject: [PATCH] feat(connector): inject only user-selected connectors per request - Read ext_info.connector_ids (list) at start of _react_agent_stream; fall back to legacy ext_info.connector_id (string) for older clients - Filter pack injection to selected ids via get_connector_tools() - Empty/missing selection no longer leaks all active connectors into prompt - Warn (not error) on stale ids so users can recover gracefully --- .../openapi/api_v1/agentic_data_api.py | 92 +++++++++++-- .../tests/test_connector_selection.py | 123 ++++++++++++++++++ 2 files changed, 202 insertions(+), 13 deletions(-) create mode 100644 packages/dbgpt-app/src/dbgpt_app/tests/test_connector_selection.py diff --git a/packages/dbgpt-app/src/dbgpt_app/openapi/api_v1/agentic_data_api.py b/packages/dbgpt-app/src/dbgpt_app/openapi/api_v1/agentic_data_api.py index 90ce27941..444a1967c 100644 --- a/packages/dbgpt-app/src/dbgpt_app/openapi/api_v1/agentic_data_api.py +++ b/packages/dbgpt-app/src/dbgpt_app/openapi/api_v1/agentic_data_api.py @@ -8,7 +8,7 @@ import tempfile import uuid import zipfile from pathlib import Path -from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional +from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Tuple from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, Query, Request, UploadFile @@ -35,6 +35,7 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: from dbgpt.agent.core.memory.gpts import GptsMemory + from dbgpt.agent.resource.connector.manager import ConnectorManager REACT_AGENT_MEMORY_CACHE: Dict[str, "GptsMemory"] = {} @@ -68,6 +69,54 @@ def _extract_auto_data_markers(text: str) -> tuple[str, Dict[str, str]]: return cleaned, extracted +def _parse_connector_ids(ext_info: Optional[Dict[str, Any]]) -> List[str]: + """Extract connector IDs from ``ext_info``. + + Supports two shapes: + - ``ext_info.connector_ids`` — a list of UUID strings (preferred). + - ``ext_info.connector_id`` — a single UUID string (legacy back-compat). + + Returns an empty list when no valid IDs are found. + """ + if not ext_info or not isinstance(ext_info, dict): + return [] + raw_ids = ext_info.get("connector_ids") + if isinstance(raw_ids, list): + return [cid for cid in raw_ids if isinstance(cid, str) and cid] + legacy = ext_info.get("connector_id") + if isinstance(legacy, str) and legacy: + return [legacy] + return [] + + +def _select_connector_packs( + connector_ids: List[str], + connector_manager: Optional["ConnectorManager"], +) -> Tuple[List["MCPToolPack"], List[str]]: + """Return tool packs for the requested *connector_ids*. + + Args: + connector_ids: IDs the user selected in the frontend. + connector_manager: A :class:`ConnectorManager` instance (or ``None``). + + Returns: + A tuple of ``(packs, missing_ids)`` where *packs* are the resolved + :class:`MCPToolPack` objects and *missing_ids* lists IDs that could + not be resolved (e.g. deleted mid-conversation). + """ + packs: List["MCPToolPack"] = [] + missing_ids: List[str] = [] + if connector_manager is None or not connector_ids: + return packs, missing_ids + for cid in connector_ids: + pack = connector_manager.get_connector_tools(cid) + if pack is None: + missing_ids.append(cid) + continue + packs.append(pack) + return packs, missing_ids + + async def _execute_skill_script_impl( skill_name: str, script_name: str, args: dict ) -> str: @@ -856,6 +905,9 @@ async def _react_agent_stream( ) database_name = dialogue.ext_info.get("database_name") + # Connector selection (Task C): only inject user-selected connectors. + connector_ids: List[str] = _parse_connector_ids(dialogue.ext_info) + def infer_phase(action: str) -> str: """根据 action 名称推断所属阶段，返回自由文本描述""" if not action: @@ -2652,7 +2704,7 @@ print(json.dumps(summary, ensure_ascii=False)) is_skill_mode = pre_matched_skill is not None _skill_name = pre_matched_skill.metadata.name if pre_matched_skill else "skill" - # Inject connector tools (safe: no-op if manager not registered or no active connectors) + # Inject connector tools — only the ones the user explicitly selected. connector_tool_extras: List[Any] = [] try: from dbgpt.agent.resource.connector.manager import ( @@ -2662,13 +2714,21 @@ print(json.dumps(summary, ensure_ascii=False)) _connector_manager = CFG.SYSTEM_APP.get_component( "connector_manager", _ConnectorManager, default_component=None ) - if _connector_manager is not None: - _connector_packs = _connector_manager.get_all_tools() - if _connector_packs: - connector_tool_extras.extend(_connector_packs) + if _connector_manager is not None and connector_ids: + connector_tool_extras, _missing = _select_connector_packs( + connector_ids, _connector_manager + ) + for _mid in _missing: + logger.warning( + "_react_agent_stream: connector_id %s not active, skipping", + _mid, + ) + if connector_tool_extras: logger.info( - "_react_agent_stream: injected %d connector tool pack(s)", - len(_connector_packs), + "_react_agent_stream: injected %d connector tool pack(s) " + "(selected: %d)", + len(connector_tool_extras), + len(connector_ids), ) except Exception: pass # graceful degradation — connector tools are optional @@ -2933,14 +2993,20 @@ Action Input: The JSON format of tool parameters _cm = CFG.SYSTEM_APP.get_component( "connector_manager", _ConnectorManager, default_component=None ) - if _cm is not None: + if _cm is not None and connector_ids: _active = _cm.list_active() - if _active: + # Only describe connectors the user explicitly selected. + # Iterate connector_ids (not _active) so prompt order matches + # user selection order and stays consistent with _select_connector_packs. + _active_map = {c["connector_id"]: c for c in _active if isinstance(c, dict)} + _selected = [ + _active_map[cid] for cid in connector_ids if cid in _active_map + ] + if _selected: _connector_lines = [] - for _c in _active: + for _c in _selected: _tool_names = ", ".join( - t.name - for t in (_c.get("tools", []) if isinstance(_c, dict) else []) + t.get("name", "unknown") for t in _c.get("tools", []) ) _connector_lines.append( f"- {_c.get('name', 'unknown')} ({_c.get('connector_type', 'unknown')}): " diff --git a/packages/dbgpt-app/src/dbgpt_app/tests/test_connector_selection.py b/packages/dbgpt-app/src/dbgpt_app/tests/test_connector_selection.py new file mode 100644 index 000000000..e9e1f47a9 --- /dev/null +++ b/packages/dbgpt-app/src/dbgpt_app/tests/test_connector_selection.py @@ -0,0 +1,123 @@ +"""Unit tests for connector selection helpers in agentic_data_api. + +Tests ``_parse_connector_ids`` and ``_select_connector_packs`` which are +extracted helpers that ``_react_agent_stream`` delegates to. +""" + +from typing import Any, Dict +from unittest.mock import MagicMock + +from dbgpt_app.openapi.api_v1.agentic_data_api import ( + _parse_connector_ids, + _select_connector_packs, +) + +# --------------------------------------------------------------------------- +# _parse_connector_ids +# --------------------------------------------------------------------------- + + +class TestParseConnectorIds: + """Verify ext_info parsing for connector_ids / connector_id.""" + + def test_none_ext_info(self): + assert _parse_connector_ids(None) == [] + + def test_empty_dict(self): + assert _parse_connector_ids({}) == [] + + def test_not_a_dict(self): + assert _parse_connector_ids("not-a-dict") == [] + + def test_connector_ids_list(self): + ext = {"connector_ids": ["id1", "id2"]} + assert _parse_connector_ids(ext) == ["id1", "id2"] + + def test_connector_ids_filters_non_strings(self): + ext = {"connector_ids": ["id1", 123, "", None, "id2"]} + assert _parse_connector_ids(ext) == ["id1", "id2"] + + def test_connector_ids_empty_list(self): + ext = {"connector_ids": []} + assert _parse_connector_ids(ext) == [] + + def test_legacy_connector_id_string(self): + ext = {"connector_id": "legacy-id"} + assert _parse_connector_ids(ext) == ["legacy-id"] + + def test_legacy_connector_id_empty_string(self): + ext = {"connector_id": ""} + assert _parse_connector_ids(ext) == [] + + def test_legacy_connector_id_not_string(self): + ext = {"connector_id": 42} + assert _parse_connector_ids(ext) == [] + + def test_connector_ids_takes_precedence_over_legacy(self): + ext = {"connector_ids": ["new-id"], "connector_id": "old-id"} + assert _parse_connector_ids(ext) == ["new-id"] + + def test_connector_ids_empty_falls_through_to_legacy(self): + """Empty list means user explicitly chose nothing -> no fallback.""" + ext = {"connector_ids": [], "connector_id": "old-id"} + # connector_ids is present (empty list), so we return empty -- no + # fallback to legacy. This matches the code: isinstance([], list) + # is True, so the list comprehension runs and returns []. + assert _parse_connector_ids(ext) == [] + + +# --------------------------------------------------------------------------- +# _select_connector_packs +# --------------------------------------------------------------------------- + + +def _make_manager(packs: Dict[str, Any]) -> MagicMock: + """Create a mock ConnectorManager with the given id->pack mapping.""" + mgr = MagicMock() + mgr.get_connector_tools.side_effect = lambda cid: packs.get(cid) + return mgr + + +class TestSelectConnectorPacks: + """Verify filtered connector pack selection.""" + + def test_empty_ids_returns_empty(self): + mgr = _make_manager({"id1": "pack1"}) + packs, missing = _select_connector_packs([], mgr) + assert packs == [] + assert missing == [] + + def test_none_manager_returns_empty(self): + packs, missing = _select_connector_packs(["id1"], None) + assert packs == [] + assert missing == [] + + def test_both_active(self): + mgr = _make_manager({"id1": "pack1", "id2": "pack2"}) + packs, missing = _select_connector_packs(["id1", "id2"], mgr) + assert packs == ["pack1", "pack2"] + assert missing == [] + + def test_one_missing(self): + mgr = _make_manager({"id1": "pack1"}) + packs, missing = _select_connector_packs(["id1", "missing"], mgr) + assert packs == ["pack1"] + assert missing == ["missing"] + + def test_all_missing(self): + mgr = _make_manager({}) + packs, missing = _select_connector_packs(["a", "b"], mgr) + assert packs == [] + assert missing == ["a", "b"] + + def test_single_id(self): + mgr = _make_manager({"solo": "solo-pack"}) + packs, missing = _select_connector_packs(["solo"], mgr) + assert packs == ["solo-pack"] + assert missing == [] + + def test_preserves_order(self): + mgr = _make_manager({"c": "pc", "a": "pa", "b": "pb"}) + packs, missing = _select_connector_packs(["a", "b", "c"], mgr) + assert packs == ["pa", "pb", "pc"] + assert missing == []