diff --git a/docsgpt/api/user/tools/mcp.py b/docsgpt/api/user/tools/mcp.py index f2efaa5d..1842c5b8 100644 --- a/docsgpt/api/user/tools/mcp.py +++ b/docsgpt/api/user/tools/mcp.py @@ -115,12 +115,13 @@ 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: +def _previous_connection(existing_doc, config, user) -> 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. + one over to a different server (its key would be sent there), to a + connection the caller does not own, or to one that signs in another way. """ from docsgpt.connectors import catalog @@ -129,9 +130,12 @@ def _previous_connection(existing_doc, config) -> str | None: if not connection_id or not base: return None with db_readonly() as conn: - row = ConnectorSessionsRepository(conn).get(str(connection_id)) + row = ConnectorSessionsRepository(conn).get_for_user(str(connection_id), user) if row is None or catalog.base_url(row.get("server_url")) != base: return None + wanted = {"oauth": "mcp_oauth", "none": "none"}.get(config.get("auth_type") or "none", "api_key") + if (row.get("auth_kind") or "") != wanted: + return None return str(connection_id) @@ -383,7 +387,7 @@ class MCPServerSave(Resource): display_name = data["displayName"] connection_id = _mcp_connection( user, storage_config, auth_type, auth_credentials, display_name, - ) or _previous_connection(existing_doc, storage_config) + ) or _previous_connection(existing_doc, storage_config, user) if connection_id and auth_type != "oauth": # The secret lives on the connection only. storage_config.pop("encrypted_credentials", None) diff --git a/docsgpt/api/user/tools/routes.py b/docsgpt/api/user/tools/routes.py index 91058a47..8d849f4a 100644 --- a/docsgpt/api/user/tools/routes.py +++ b/docsgpt/api/user/tools/routes.py @@ -610,19 +610,11 @@ class UpdateTool(Resource): else {} ) existing_config = tool_doc.get("config", {}) or {} - has_existing_secrets = "encrypted_credentials" in existing_config - if tool_doc.get("connection_id"): - has_existing_secrets = True - new_config = _update_connection_secrets( - conn, user, tool_doc, data["config"], config_requirements, - ) - if new_config is None: - return make_response( - jsonify({"success": False, "message": "Only the owner can change the credentials"}), - 403, - ) - data = {**data, "config": new_config} - + has_existing_secrets = ( + "encrypted_credentials" in existing_config or bool(tool_doc.get("connection_id")) + ) + # Validate before touching the connection: a rejected + # edit must not have rotated its credentials. if config_requirements: validation_errors = _validate_config( data["config"], config_requirements, @@ -637,6 +629,16 @@ class UpdateTool(Resource): }), 400, ) + if tool_doc.get("connection_id"): + new_config = _update_connection_secrets( + conn, user, tool_doc, data["config"], config_requirements, + ) + if new_config is None: + return make_response( + jsonify({"success": False, "message": "Only the owner can change the credentials"}), + 403, + ) + data = {**data, "config": new_config} update_data["config"] = _merge_secrets_on_update( data["config"], existing_config, config_requirements, user diff --git a/tests/connectors/test_runtime.py b/tests/connectors/test_runtime.py index 6d0a3d16..e1b9d0e2 100644 --- a/tests/connectors/test_runtime.py +++ b/tests/connectors/test_runtime.py @@ -326,7 +326,11 @@ class TestMcpServerMismatch: cid = _connection(pg_conn, provider="custom_mcp", server_url="https://old.example.com") existing = {"connection_id": cid} + same = {"server_url": "https://old.example.com/mcp", "auth_type": "bearer"} 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 + assert _previous_connection(existing, same, "alice") == cid + assert _previous_connection(existing, {**same, "server_url": "https://new.example.com/mcp"}, "alice") is None + assert _previous_connection(None, same, "alice") is None + # Someone else's connection, or one that signs in another way, is not kept. + assert _previous_connection(existing, same, "bob") is None + assert _previous_connection(existing, {**same, "auth_type": "oauth"}, "alice") is None diff --git a/tests/connectors/test_tool_routes.py b/tests/connectors/test_tool_routes.py index 8ac0785e..e93ae31a 100644 --- a/tests/connectors/test_tool_routes.py +++ b/tests/connectors/test_tool_routes.py @@ -183,3 +183,14 @@ class TestUpdateConnectedTool: assert resp.status_code == 200 row = _connection_row(pg_conn, created["connection_id"]) assert service.read_secrets(row)["credentials"]["token"] == "123456:SECRETTOKEN" + + def test_rejected_edit_leaves_the_connection_untouched(self, app, pg_conn): + from docsgpt.api.user.tools.routes import UpdateTool + + created = self._create(app, pg_conn) + with _db(pg_conn), patch("docsgpt.api.user.tools.routes._validate_config", + return_value={"timeout": "Timeout must be between 1 and 300"}): + resp = _call(app, UpdateTool, {"id": created["id"], "config": {"token": "999999:ROTATED"}}) + assert resp.status_code == 400 + row = _connection_row(pg_conn, created["connection_id"]) + assert service.read_secrets(row)["credentials"]["token"] == "123456:SECRETTOKEN"