refactor: improve session management in ToolManager

api/core/tools/tool_manager.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast
 
 import sqlalchemy as sa
 from pydantic import TypeAdapter
+from sqlalchemy.orm import Session
 from yarl import URL
 
 import contexts
@@ -617,8 +618,9 @@ class ToolManager:
                 WHERE tenant_id = :tenant_id
                 ORDER BY tenant_id, provider, is_default DESC, created_at DESC
                 """
-        ids = [row.id for row in db.session.execute(sa.text(sql), {"tenant_id": tenant_id}).all()]
+        with Session(db.engine).no_autoflush as session:
-        return db.session.query(BuiltinToolProvider).where(BuiltinToolProvider.id.in_(ids)).all()
+            ids = [row.id for row in session.execute(sa.text(sql), {"tenant_id": tenant_id}).all()]
+            return session.query(BuiltinToolProvider).where(BuiltinToolProvider.id.in_(ids)).all()
 
     @classmethod
     def list_providers_from_api(

api/services/tools/builtin_tools_manage_service.py

@@ -546,54 +546,53 @@ class BuiltinToolManageService:
         # get all builtin providers
         provider_controllers = ToolManager.list_builtin_providers(tenant_id)
 
-        with db.session.no_autoflush:
+        # get all user added providers
-            # get all user added providers
+        db_providers: list[BuiltinToolProvider] = ToolManager.list_default_builtin_providers(tenant_id)
-            db_providers: list[BuiltinToolProvider] = ToolManager.list_default_builtin_providers(tenant_id)
 
-            # rewrite db_providers
+        # rewrite db_providers
-            for db_provider in db_providers:
+        for db_provider in db_providers:
-                db_provider.provider = str(ToolProviderID(db_provider.provider))
+            db_provider.provider = str(ToolProviderID(db_provider.provider))
 
-            # find provider
+        # find provider
-            def find_provider(provider):
+        def find_provider(provider):
-                return next(filter(lambda db_provider: db_provider.provider == provider, db_providers), None)
+            return next(filter(lambda db_provider: db_provider.provider == provider, db_providers), None)
 
-            result: list[ToolProviderApiEntity] = []
+        result: list[ToolProviderApiEntity] = []
 
-            for provider_controller in provider_controllers:
+        for provider_controller in provider_controllers:
-                try:
+            try:
-                    # handle include, exclude
+                # handle include, exclude
-                    if is_filtered(
+                if is_filtered(
-                        include_set=dify_config.POSITION_TOOL_INCLUDES_SET,  # type: ignore
+                    include_set=dify_config.POSITION_TOOL_INCLUDES_SET,  # type: ignore
-                        exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET,  # type: ignore
+                    exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET,  # type: ignore
-                        data=provider_controller,
+                    data=provider_controller,
-                        name_func=lambda x: x.identity.name,
+                    name_func=lambda x: x.identity.name,
-                    ):
+                ):
-                        continue
+                    continue
 
-                    # convert provider controller to user provider
+                # convert provider controller to user provider
-                    user_builtin_provider = ToolTransformService.builtin_provider_to_user_provider(
+                user_builtin_provider = ToolTransformService.builtin_provider_to_user_provider(
-                        provider_controller=provider_controller,
+                    provider_controller=provider_controller,
-                        db_provider=find_provider(provider_controller.entity.identity.name),
+                    db_provider=find_provider(provider_controller.entity.identity.name),
-                        decrypt_credentials=True,
+                    decrypt_credentials=True,
+                )
+
+                # add icon
+                ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_builtin_provider)
+
+                tools = provider_controller.get_tools()
+                for tool in tools or []:
+                    user_builtin_provider.tools.append(
+                        ToolTransformService.convert_tool_entity_to_api_entity(
+                            tenant_id=tenant_id,
+                            tool=tool,
+                            labels=ToolLabelManager.get_tool_labels(provider_controller),
+                        )
                     )
 
-                    # add icon
+                result.append(user_builtin_provider)
-                    ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_builtin_provider)
+            except Exception as e:
-
+                raise e
-                    tools = provider_controller.get_tools()
-                    for tool in tools or []:
-                        user_builtin_provider.tools.append(
-                            ToolTransformService.convert_tool_entity_to_api_entity(
-                                tenant_id=tenant_id,
-                                tool=tool,
-                                labels=ToolLabelManager.get_tool_labels(provider_controller),
-                            )
-                        )
-
-                    result.append(user_builtin_provider)
-                except Exception as e:
-                    raise e
 
         return BuiltinToolProviderSort.sort(result)
 
@@ -604,7 +603,7 @@ class BuiltinToolManageService:
         1.if the default provider exists, return the default provider
         2.if the default provider does not exist, return the oldest provider
         """
-        with Session(db.engine) as session:
+        with Session(db.engine).no_autoflush as session:
             try:
                 full_provider_name = provider_name
                 provider_id_entity = ToolProviderID(provider_name)