diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index a3b7a88a..4e2938ce 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -1652,6 +1652,16 @@ class ToolExecutor: f"{resolved.connector_name or 'This service'} needs to be connected", connection_id=resolved.connection_id, ) + if tool_data.get("name") == "mcp_tool": + from docsgpt.connectors import catalog + + # A connection's secret only goes to the server it was stored for. + stored_for = catalog.base_url(resolved.row.get("server_url")) + 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, + ) audit_delegation( resolved, invoker=self.user, diff --git a/docsgpt/api/user/tools/mcp.py b/docsgpt/api/user/tools/mcp.py index bfe6fb09..f2efaa5d 100644 --- a/docsgpt/api/user/tools/mcp.py +++ b/docsgpt/api/user/tools/mcp.py @@ -115,6 +115,26 @@ def _mcp_connection(user, config, auth_type, auth_credentials, display_name): return str(row["id"]) if row else None +def _previous_connection(existing_doc, config) -> str | None: + """The saved tool's connection, kept only while the server is unchanged. + + An edit that cannot resolve a connection of its own (the default + encryption key blocks storing a new secret) must not carry the old + one over to a different server: its key would be sent there. + """ + from docsgpt.connectors import catalog + + connection_id = (existing_doc or {}).get("connection_id") + base = catalog.base_url(config.get("server_url")) + if not connection_id or not base: + return None + with db_readonly() as conn: + row = ConnectorSessionsRepository(conn).get(str(connection_id)) + if row is None or catalog.base_url(row.get("server_url")) != base: + return None + return str(connection_id) + + def _mcp_policy_error(config: dict): """A 403 when an admin turned this MCP server's connector off, else None. @@ -363,7 +383,7 @@ class MCPServerSave(Resource): display_name = data["displayName"] connection_id = _mcp_connection( user, storage_config, auth_type, auth_credentials, display_name, - ) or (str(existing_doc["connection_id"]) if existing_doc and existing_doc.get("connection_id") else None) + ) or _previous_connection(existing_doc, storage_config) if connection_id and auth_type != "oauth": # The secret lives on the connection only. storage_config.pop("encrypted_credentials", None) diff --git a/tests/connectors/test_runtime.py b/tests/connectors/test_runtime.py index 2649ef7b..6d0a3d16 100644 --- a/tests/connectors/test_runtime.py +++ b/tests/connectors/test_runtime.py @@ -301,3 +301,32 @@ def test_redaction_covers_connection_secrets(secret_key): from docsgpt.storage.db.redaction import REDACTED, redact_secrets assert redact_secrets({secret_key: {"x": "y"}})[secret_key] == REDACTED + + +class TestMcpServerMismatch: + def test_connection_for_another_server_is_not_applied(self, pg_conn): + """A key stored for one MCP server is never sent to a tool now pointing at another.""" + from docsgpt.connectors import service + + cid = _connection(pg_conn, provider="custom_mcp", server_url="https://old.example.com", + secrets={"credentials": {"bearer_token": "old-server-secret"}}) + tool = {**_tool(cid, name="mcp_tool"), "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_save_keeps_previous_connection_only_for_the_same_server(self, pg_conn): + from docsgpt.api.user.tools.mcp import _previous_connection + + @contextmanager + def _yield(): + yield pg_conn + + cid = _connection(pg_conn, provider="custom_mcp", server_url="https://old.example.com") + existing = {"connection_id": cid} + with patch("docsgpt.api.user.tools.mcp.db_readonly", _yield): + assert _previous_connection(existing, {"server_url": "https://old.example.com/mcp"}) == cid + assert _previous_connection(existing, {"server_url": "https://new.example.com/mcp"}) is None + assert _previous_connection(None, {"server_url": "https://old.example.com/mcp"}) is None