feat(mcp): add plugin MCP-tools hook and promote shared connection helpers (TG-1124)

Add PluginSpec.get_mcp_tools() (default empty) and register plugin-provided tools in
build_mcp_server through the same safe_tool wrapper as core tools. Promote the shared
connection render/validation helpers (effective_mode, raise_validation_error,
render_connection_body, authentication_label) to mcp/tools/common.py so the read/test
tools and any plugin tools share them. Tag the connection and table-group tool modules
with _DOC_GROUP=MANAGE so they group under 'Manage TestGen configuration'.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: rboni-dk

testgen/mcp/server.py

@@ -378,6 +378,17 @@ def safe_prompt(fn):
     safe_prompt(profiling_overview)
     safe_prompt(hygiene_triage)
 
+    # Register plugin-provided tools through the same safe_tool wrapper as the core tools above.
+    from testgen.utils.plugins import discover
+
+    for plugin in discover():
+        try:
+            spec = plugin.load()
+            for tool in spec.get_mcp_tools():
+                safe_tool(tool)
+        except Exception:
+            LOG.warning("Plugin %s failed to load; skipping its MCP tools", plugin.package)
+
     return mcp
 
 

testgen/mcp/tools/common.py

@@ -42,6 +42,7 @@
 from testgen.common.models.test_suite import TestSuite
 from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError
 from testgen.mcp.permissions import get_project_permissions
+from testgen.mcp.tools.markdown import MdDoc
 
 # User-facing label for ``Disposition.INACTIVE`` is "Muted" — accept that label on input.
 _DISPOSITION_USER_TO_DB: dict[str, Disposition] = {
@@ -1374,3 +1375,54 @@ def validate_connection_fields(connection: Connection) -> list[str]:
         errors.append(f"`max_query_chars` must be between {_MAX_QUERY_CHARS_MIN} and {_MAX_QUERY_CHARS_MAX}.")
 
     return errors
+
+
+def effective_mode(connection: Connection, connection_mode: str | None) -> str | None:
+    """Mode label to apply: the explicit override, else the connection's current mode."""
+    if connection_mode is not None:
+        return connection_mode
+    inferred = infer_mode(connection)
+    return str(inferred) if inferred is not None else None
+
+
+def raise_validation_error(errors: list[str], header: str) -> None:
+    bullets = "\n".join(f"- {err}" for err in errors)
+    raise MCPUserError(f"{header}\n\n{bullets}")
+
+
+def render_connection_body(doc: MdDoc, connection: Connection) -> None:
+    """Render every non-secret connection field below the heading.
+
+    Encrypted columns are filtered out via ``ConnField.secret``.
+    """
+    doc.field("ID", connection.connection_id, code=True)
+    doc.field("Project", connection.project_code, code=True)
+    doc.field("Type", format_flavor_label(connection.sql_flavor_code))
+
+    # Each populated, non-secret field under its flavor-specific label
+    # (e.g. "Catalog" for Databricks, "Login URL" for Salesforce).
+    for fld in connection_display_fields(connection):
+        if fld.secret:
+            continue
+        value = getattr(connection, fld.column, None)
+        if value in (None, ""):
+            continue
+        doc.field(fld.label, value, code=fld.column != "project_port")
+
+    doc.field("Authentication", authentication_label(connection))
+    if connection.max_threads is not None:
+        doc.field("Max Threads", connection.max_threads)
+    if connection.max_query_chars is not None:
+        doc.field("Max Expression Length", connection.max_query_chars)
+
+
+def authentication_label(connection: Connection) -> str:
+    """The connection's auth method: the active connection mode for multi-mode
+    flavors, else the implicit method (service account key, else password).
+    """
+    mode = infer_mode(connection)
+    if mode is not None:
+        return str(mode)
+    if connection.service_account_key:
+        return "Service Account Key"
+    return "Password"

testgen/mcp/tools/connections.py

@@ -19,20 +19,24 @@
 from testgen.mcp.permissions import get_project_permissions, mcp_permission
 from testgen.mcp.tools.common import (
     SQL_FLAVOR_CODE_TO_LABEL,
+    DocGroup,
     apply_connection_params,
-    connection_display_fields,
+    effective_mode,
     format_flavor_label,
     format_page_footer,
     format_page_info,
-    infer_mode,
     parse_sql_flavor,
+    raise_validation_error,
+    render_connection_body,
     resolve_connection,
     validate_connection_fields,
     validate_limit,
     validate_page,
 )
 from testgen.mcp.tools.markdown import MdDoc
 
+_DOC_GROUP = DocGroup.MANAGE
+
 
 @with_database_session
 @mcp_permission("view")
@@ -99,7 +103,7 @@ def get_connection(connection_id: int) -> str:
     connection = resolve_connection(connection_id)
     doc = MdDoc()
     doc.heading(1, f"Connection `{connection.connection_name}`")
-    _render_connection_body(doc, connection)
+    render_connection_body(doc, connection)
     return doc.render()
 
 
@@ -151,7 +155,7 @@ def test_connection(
         connection = Connection(sql_flavor=family, sql_flavor_code=code)
 
     if connection_params is not None or connection_mode is not None:
-        mode = connection_mode if inline else _effective_mode(connection, connection_mode)
+        mode = connection_mode if inline else effective_mode(connection, connection_mode)
         apply_connection_params(connection, connection.sql_flavor_code, mode, connection_params or {})
 
     normalize_auth_fields(connection)
@@ -160,7 +164,7 @@ def test_connection(
 
     errors = validate_connection_fields(connection)
     if errors:
-        _raise_validation_error(errors, "Cannot test connection. Required fields missing or invalid.")
+        raise_validation_error(errors, "Cannot test connection. Required fields missing or invalid.")
 
     status = test_connection_status(connection)
 
@@ -181,59 +185,3 @@ def test_connection(
     if status.details:
         doc.code_block(status.details)
     return doc.render()
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-
-def _effective_mode(connection: Connection, connection_mode: str | None) -> str | None:
-    """Mode label to apply: the explicit override, else the connection's current mode."""
-    if connection_mode is not None:
-        return connection_mode
-    inferred = infer_mode(connection)
-    return str(inferred) if inferred is not None else None
-
-
-def _raise_validation_error(errors: list[str], header: str) -> None:
-    bullets = "\n".join(f"- {err}" for err in errors)
-    raise MCPUserError(f"{header}\n\n{bullets}")
-
-
-def _render_connection_body(doc: MdDoc, connection: Connection) -> None:
-    """Render every non-secret connection field below the heading.
-
-    Encrypted columns are filtered out via ``ConnField.secret``.
-    """
-    doc.field("ID", connection.connection_id, code=True)
-    doc.field("Project", connection.project_code, code=True)
-    doc.field("Type", format_flavor_label(connection.sql_flavor_code))
-
-    # Each populated, non-secret field under its flavor-specific label
-    # (e.g. "Catalog" for Databricks, "Login URL" for Salesforce).
-    for fld in connection_display_fields(connection):
-        if fld.secret:
-            continue
-        value = getattr(connection, fld.column, None)
-        if value in (None, ""):
-            continue
-        doc.field(fld.label, value, code=fld.column != "project_port")
-
-    doc.field("Authentication", _authentication_label(connection))
-    if connection.max_threads is not None:
-        doc.field("Max Threads", connection.max_threads)
-    if connection.max_query_chars is not None:
-        doc.field("Max Expression Length", connection.max_query_chars)
-
-
-def _authentication_label(connection: Connection) -> str:
-    """The connection's auth method: the active connection mode for multi-mode
-    flavors, else the implicit method (service account key, else password).
-    """
-    mode = infer_mode(connection)
-    if mode is not None:
-        return str(mode)
-    if connection.service_account_key:
-        return "Service Account Key"
-    return "Password"

testgen/mcp/tools/table_groups.py

@@ -21,6 +21,7 @@
 from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError
 from testgen.mcp.permissions import get_project_permissions, mcp_permission
 from testgen.mcp.tools.common import (
+    DocGroup,
     format_flavor_label,
     format_page_footer,
     format_page_info,
@@ -32,6 +33,8 @@
 from testgen.mcp.tools.markdown import MdDoc
 from testgen.utils import friendly_score
 
+_DOC_GROUP = DocGroup.MANAGE
+
 _DUPLICATE_NAME_MESSAGE = "A Table Group with the same name already exists."
 _PII_FLAG_DENIED_MESSAGE = (
     "Changing PII detection requires permission to view PII. "

testgen/utils/plugins.py

@@ -3,7 +3,7 @@
 import importlib
 import importlib.metadata
 import inspect
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from typing import ClassVar, get_args
 
 from testgen.ui.assets import get_asset_path
@@ -61,6 +61,16 @@ def configure_ui(cls) -> None:
         is actually running. Called by ``Plugin.load_streamlit()``, never by ``Plugin.load()``.
         """
 
+    @classmethod
+    def get_mcp_tools(cls) -> list[Callable]:
+        """Return MCP tool callables to register into the MCP server.
+
+        Override in plugins to register additional MCP tools. Implementations import their
+        tool modules inside the method (like ``configure_ui``) so plugin import stays light
+        for processes that never build the MCP server.
+        """
+        return []
+
 
 class PluginHook:
     """Singleton holding resolved plugin values, pre-loaded with defaults."""