diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index 2d8aac8c..add4df1f 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -1879,32 +1879,50 @@ class ToolExecutor: def _apply_connection(self, tool_data: Dict, tool_id: str, tool_config: Dict, resolved) -> None: """Merge a connection's credentials into ``tool_config``. + A connection only serves the tools its connector provides (a Telegram + bot token never reaches an ntfy tool), and an MCP connection's secret + only goes to the server it was stored for. + + Args: + tool_data: The ``user_tools`` row being run. + tool_id: The tool's key in this run. + tool_config: The config the tool is loaded with, updated in place. + resolved: The connection :func:`resolve_connection` picked. + Raises: - ConnectionUnavailable: The connection needs signing in again. + ConnectionUnavailable: The connection needs signing in again, or + does not belong to this tool's connector or server. """ - from docsgpt.connectors import service + from docsgpt.connectors import catalog, service from docsgpt.connectors.resolve import audit_delegation + unavailable = service.ConnectionUnavailable( + f"{resolved.connector_name or 'This service'} needs to be connected", + connection_id=resolved.connection_id, + ) if not resolved.available or resolved.row is None: - raise service.ConnectionUnavailable( - f"{resolved.connector_name or 'This service'} needs to be connected", - connection_id=resolved.connection_id, + raise unavailable + tool_name = tool_data.get("name") + definition = catalog.get_definition(resolved.connector_key) + if definition is None or tool_name not in definition.tool_templates: + logger.warning( + "tool %s (%s) points at a %s connection", tool_data.get("id") or tool_id, tool_name, + resolved.connector_key, ) - if tool_data.get("name") == "mcp_tool": - from docsgpt.connectors import catalog - + raise unavailable + if tool_name == "mcp_tool": # A connection's secret only goes to the server it was stored for: - # a custom server's own URL, or a built-in connector's MCP server - # (GitHub's token only ever goes to GitHub's). + # a custom server's own URL (legacy rows name it in ``provider``), + # or a preset's or built-in connector's MCP server (GitHub's token + # only ever goes to GitHub's). No server known, nothing is sent. stored_for = catalog.base_url(resolved.row.get("server_url")) - definition = catalog.get_definition(resolved.connector_key) - if not stored_for and definition is not None and definition.publisher == "built_in": + provider = str(resolved.row.get("provider") or "") + if not stored_for and provider.startswith("mcp:"): + stored_for = catalog.base_url(provider[len("mcp:"):]) + if not stored_for: stored_for = definition.mcp_base_url or "" - if stored_for and stored_for != catalog.base_url(tool_config.get("server_url")): - raise service.ConnectionUnavailable( - f"{resolved.connector_name or 'This service'} needs to be connected", - connection_id=resolved.connection_id, - ) + if not stored_for or stored_for != catalog.base_url(tool_config.get("server_url")): + raise unavailable if service.builtin_mcp_config(definition) is not None: # Only the connector's own endpoints: GitHub's write one while # the tool opted in and an admin allows it, else read-only. diff --git a/docsgpt/connectors/catalog.py b/docsgpt/connectors/catalog.py index d294d409..1e4e66d9 100644 --- a/docsgpt/connectors/catalog.py +++ b/docsgpt/connectors/catalog.py @@ -479,7 +479,9 @@ def connector_key_for_row(row: dict) -> Optional[str]: """ key = row.get("connector_key") provider = row.get("provider") or "" - if not key: + if not key or key.startswith("mcp:"): + # ``mcp:`` is a provider value, never a catalog key. + key = None if provider.startswith("mcp:"): key = "custom_mcp" elif get_definition(provider): diff --git a/tests/connectors/test_runtime.py b/tests/connectors/test_runtime.py index a39aaebf..7488323a 100644 --- a/tests/connectors/test_runtime.py +++ b/tests/connectors/test_runtime.py @@ -533,6 +533,41 @@ class TestMcpServerMismatch: _executor()._get_or_load_tool(tool, "t1", "search") manager.return_value.load_tool.assert_not_called() + @pytest.mark.parametrize("tool_name, provider, server_url", [ + # A service's key without a server of its own is not an MCP server's. + ("mcp_tool", "telegram", None), + ("mcp_tool", "ntfy", None), + # A custom server connection that names no server has nowhere to go. + ("mcp_tool", "custom_mcp", None), + # A tool only runs on a connection of the connector that provides it. + ("ntfy", "telegram", None), + ("telegram", "ntfy", None), + ("telegram", "custom_mcp", "https://new.example.com"), + ]) + def test_connection_of_another_connector_is_not_applied(self, pg_conn, tool_name, provider, server_url): + from docsgpt.connectors import service + + cid = _connection(pg_conn, provider=provider, server_url=server_url, + secrets={"credentials": {"token": "bot-token", "bearer_token": "bot-token"}}) + tool = {**_tool(cid, name=tool_name), "config": {"server_url": "https://new.example.com/mcp", + "auth_type": "bearer"}} + with _service_db(pg_conn), patch("docsgpt.agents.tool_executor.ToolManager") as manager: + with pytest.raises(service.ConnectionUnavailable): + _executor()._get_or_load_tool(tool, "t1", "search") + manager.return_value.load_tool.assert_not_called() + + def test_legacy_mcp_connection_is_still_applied_to_its_server(self, pg_conn): + """Rows from before connector keys are named from their ``mcp:`` provider.""" + cid = _connection(pg_conn, provider="mcp:https://m.example.com", auth_kind="mcp_oauth", + server_url=None, secrets={"tokens": {"access_token": "x"}}) + pg_conn.execute(text("UPDATE connector_sessions SET connector_key = NULL WHERE id = CAST(:i AS uuid)"), + {"i": cid}) + tool = {**_tool(cid, name="mcp_tool"), "config": {"server_url": "https://m.example.com/mcp", + "auth_type": "oauth"}} + with _service_db(pg_conn), patch("docsgpt.agents.tool_executor.ToolManager") as manager: + _executor()._get_or_load_tool(tool, "t1", "search") + assert manager.return_value.load_tool.call_args.kwargs["tool_config"]["connection_id"] == cid + def test_save_keeps_previous_connection_only_for_the_same_server(self, pg_conn): from docsgpt.api.user.tools.mcp import _previous_connection