fix: database lock timeout by separating external MCP calls from transactions (#22821)

api/services/tools/mcp_tools_manage_service.py

@@ -112,19 +112,27 @@ class MCPToolManageService:
     @classmethod
     def list_mcp_tool_from_remote_server(cls, tenant_id: str, provider_id: str) -> ToolProviderApiEntity:
         mcp_provider = cls.get_mcp_provider_by_provider_id(provider_id, tenant_id)
+        server_url = mcp_provider.decrypted_server_url
+        authed = mcp_provider.authed
+
         try:
-            with MCPClient(
+            with MCPClient(server_url, provider_id, tenant_id, authed=authed, for_list=True) as mcp_client:
-                mcp_provider.decrypted_server_url, provider_id, tenant_id, authed=mcp_provider.authed, for_list=True
-            ) as mcp_client:
                 tools = mcp_client.list_tools()
         except MCPAuthError:
             raise ValueError("Please auth the tool first")
         except MCPError as e:
             raise ValueError(f"Failed to connect to MCP server: {e}")
-        mcp_provider.tools = json.dumps([tool.model_dump() for tool in tools])
+
-        mcp_provider.authed = True
+        try:
-        mcp_provider.updated_at = datetime.now()
+            mcp_provider = cls.get_mcp_provider_by_provider_id(provider_id, tenant_id)
-        db.session.commit()
+            mcp_provider.tools = json.dumps([tool.model_dump() for tool in tools])
+            mcp_provider.authed = True
+            mcp_provider.updated_at = datetime.now()
+            db.session.commit()
+        except Exception:
+            db.session.rollback()
+            raise
+
         user = mcp_provider.load_user()
         return ToolProviderApiEntity(
             id=mcp_provider.id,
@@ -160,22 +168,35 @@ class MCPToolManageService:
         server_identifier: str,
     ):
         mcp_provider = cls.get_mcp_provider_by_provider_id(provider_id, tenant_id)
-        mcp_provider.updated_at = datetime.now()
+
-        mcp_provider.name = name
+        reconnect_result = None
-        mcp_provider.icon = (
+        encrypted_server_url = None
-            json.dumps({"content": icon, "background": icon_background}) if icon_type == "emoji" else icon
+        server_url_hash = None
-        )
-        mcp_provider.server_identifier = server_identifier
 
         if UNCHANGED_SERVER_URL_PLACEHOLDER not in server_url:
             encrypted_server_url = encrypter.encrypt_token(tenant_id, server_url)
-            mcp_provider.server_url = encrypted_server_url
             server_url_hash = hashlib.sha256(server_url.encode()).hexdigest()
 
             if server_url_hash != mcp_provider.server_url_hash:
-                cls._re_connect_mcp_provider(mcp_provider, provider_id, tenant_id)
+                reconnect_result = cls._re_connect_mcp_provider(server_url, provider_id, tenant_id)
-                mcp_provider.server_url_hash = server_url_hash
+
         try:
+            mcp_provider.updated_at = datetime.now()
+            mcp_provider.name = name
+            mcp_provider.icon = (
+                json.dumps({"content": icon, "background": icon_background}) if icon_type == "emoji" else icon
+            )
+            mcp_provider.server_identifier = server_identifier
+
+            if encrypted_server_url is not None and server_url_hash is not None:
+                mcp_provider.server_url = encrypted_server_url
+                mcp_provider.server_url_hash = server_url_hash
+
+                if reconnect_result:
+                    mcp_provider.authed = reconnect_result["authed"]
+                    mcp_provider.tools = reconnect_result["tools"]
+                    mcp_provider.encrypted_credentials = reconnect_result["encrypted_credentials"]
+
             db.session.commit()
         except IntegrityError as e:
             db.session.rollback()
@@ -187,6 +208,9 @@ class MCPToolManageService:
             if "unique_mcp_provider_server_identifier" in error_msg:
                 raise ValueError(f"MCP tool {server_identifier} already exists")
             raise
+        except Exception:
+            db.session.rollback()
+            raise
 
     @classmethod
     def update_mcp_provider_credentials(
@@ -207,23 +231,22 @@ class MCPToolManageService:
         db.session.commit()
 
     @classmethod
-    def _re_connect_mcp_provider(cls, mcp_provider: MCPToolProvider, provider_id: str, tenant_id: str):
+    def _re_connect_mcp_provider(cls, server_url: str, provider_id: str, tenant_id: str):
-        """re-connect mcp provider"""
         try:
             with MCPClient(
-                mcp_provider.decrypted_server_url,
+                server_url,
                 provider_id,
                 tenant_id,
                 authed=False,
                 for_list=True,
             ) as mcp_client:
                 tools = mcp_client.list_tools()
-                mcp_provider.authed = True
+                return {
-                mcp_provider.tools = json.dumps([tool.model_dump() for tool in tools])
+                    "authed": True,
+                    "tools": json.dumps([tool.model_dump() for tool in tools]),
+                    "encrypted_credentials": "{}",
+                }
         except MCPAuthError:
-            mcp_provider.authed = False
+            return {"authed": False, "tools": "[]", "encrypted_credentials": "{}"}
-            mcp_provider.tools = "[]"
         except MCPError as e:
             raise ValueError(f"Failed to re-connect MCP server: {e}") from e
-        # reset credentials
-        mcp_provider.encrypted_credentials = "{}"