test: use real plugin runtime objects

Co-authored-by: Byron Wang <byron@dify.ai>

.github/workflows/main-ci.yml

@@ -63,9 +63,11 @@ jobs:
               - 'api/**'
               - 'scripts/ast_grep_guard.py'
               - 'scripts/check_no_new_getattr.py'
+              - 'scripts/check_no_spec_mock.py'
               - 'scripts/check_no_new_controller_sqlalchemy.py'
               - 'scripts/lint_controller_sqlalchemy.py'
               - 'scripts/ast_grep_rules/no_new_getattr.yml'
+              - 'scripts/ast_grep_rules/no_spec_mock.yml'
               - 'scripts/ast_grep_rules/no_new_controller_sqlalchemy.yml'
               - 'scripts/ast_grep_rules/no_new_session_mock.yml'
               - '.github/workflows/style.yml'
@@ -162,10 +164,12 @@ jobs:
               - 'api/**'
               - 'scripts/ast_grep_guard.py'
               - 'scripts/check_no_new_getattr.py'
+              - 'scripts/check_no_spec_mock.py'
               - 'scripts/check_no_new_controller_sqlalchemy.py'
               - 'scripts/check_no_new_session_mock.py'
               - 'scripts/lint_controller_sqlalchemy.py'
               - 'scripts/ast_grep_rules/no_new_getattr.yml'
+              - 'scripts/ast_grep_rules/no_spec_mock.yml'
               - 'scripts/ast_grep_rules/no_new_controller_sqlalchemy.yml'
               - 'scripts/ast_grep_rules/no_new_session_mock.yml'
               - '.github/workflows/style.yml'

.github/workflows/style.yml

@@ -93,6 +93,10 @@ jobs:
         if: inputs.run-python-style
         run: uv run --project api python scripts/check_no_new_session_mock.py --base-rev "${{ inputs.base-rev }}"
 
+      - name: Run No Spec Mock Guard
+        if: inputs.run-python-style
+        run: uv run --project api python scripts/check_no_spec_mock.py
+
       - name: Run Type Checks
         if: inputs.run-python-style
         env:
@@ -247,7 +251,7 @@ jobs:
           persist-credentials: false
 
       - name: Super-linter
-        uses: super-linter/super-linter/slim@4ce20838b8ab83717e78138c5b3a1407148e0918 # v8.7.0
+        uses: super-linter/super-linter/slim@2da136927bd4a73596db63044b504547c62cb854 # v9.0.0
         env:
           BASH_SEVERITY: warning
           DEFAULT_BRANCH: origin/main

.github/workflows/translate-i18n-claude.yml

@@ -163,7 +163,7 @@ jobs:
 
       - name: Run Claude Code for Translation Sync
         if: steps.context.outputs.CHANGED_FILES != ''
-        uses: anthropics/claude-code-action@756cc22e19660d20e8cc9496b4f242475a7f7790 # v1.0.235
+        uses: anthropics/claude-code-action@ed670b4cf9de2a5a570d130d2f6197b9e543cd64 # v1.0.240
         with:
           anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}
           github_token: ${{ secrets.GITHUB_TOKEN }}

Makefile

@@ -80,6 +80,7 @@ lint:
 	@uv run --project api --dev ruff format ./api
 	@uv run --project api --dev ruff check --fix ./api
 	@$(MAKE) api-contract-lint
+	@uv run --project api python scripts/check_no_spec_mock.py
 	@uv run --directory api --dev lint-imports
 	@uv run --project api --dev dotenv-linter ./api/.env.example ./web/.env.example
 	@echo "✅ Linting complete"

api/.importlinter

@@ -832,6 +832,36 @@ forbidden_modules =
     sqlalchemy
     werkzeug
 
+[importlinter:contract:api-based-extension-application-service-boundary]
+name = API-based extension application service is framework and persistence neutral
+type = forbidden
+source_modules =
+    services.api_based_extension_application_service
+forbidden_modules =
+    configs
+    controllers
+    core
+    extensions
+    flask
+    models
+    repositories
+    sqlalchemy
+    werkzeug
+
+[importlinter:contract:console-extension-controller-boundary]
+name = Console extension controller reaches persistence only through the composition root
+type = forbidden
+source_modules =
+    controllers.console.extension
+forbidden_modules =
+    extensions.ext_database
+    models
+    repositories
+    sqlalchemy
+# The composition root (extensions.ext_application_services) wires repositories,
+# so only direct imports are checked here.
+allow_indirect_imports = True
+
 [importlinter:contract:tool-file-download-service-boundary]
 name = Tool file download application service is framework and persistence neutral
 type = forbidden

api/commands/rbac.py

@@ -463,6 +463,36 @@ def _workspace_member_account_id_batches(tenant_id: str, batch_size: int) -> Ite
         last_join_id = str(rows[-1].id)
 
 
+def _workspace_non_admin_account_id_batches(
+    tenant_id: str,
+    operator_account_id: str,
+    batch_size: int,
+) -> Iterator[list[str]]:
+    """Yield workspace members except accounts bound to RBAC's builtin Owner/Admin roles.
+
+    Resource administrators receive access implicitly through RBAC and their
+    resource access-policy bindings are locked. Resolve their current RBAC
+    bindings in batches instead of relying on the legacy workspace role column.
+    """
+    for batch in _workspace_member_account_id_batches(tenant_id, batch_size):
+        member_roles = RBACService.MemberRoles.batch_get(
+            tenant_id=tenant_id,
+            account_id=operator_account_id,
+            member_account_ids=batch,
+        )
+        protected_account_ids = {
+            member.account_id
+            for member in member_roles
+            if any(
+                role.is_builtin and role.category == "global_system_default" and role.role_tag in {"owner", "admin"}
+                for role in member.roles
+            )
+        }
+        writable_batch = [account_id for account_id in batch if account_id not in protected_account_ids]
+        if writable_batch:
+            yield writable_batch
+
+
 def _replace_resource_whitelist(
     resource_type: str,
     *,
@@ -725,7 +755,7 @@ def migrate_resource_whitelist_scopes_to_automatic_include(
         migrated_count += 1
 
         if scope is RBACResourceWhitelistScope.ALL:
-            for batch in _workspace_member_account_id_batches(workspace_id, member_batch_size):
+            for batch in _workspace_non_admin_account_id_batches(workspace_id, operator_account_id, member_batch_size):
                 _replace_resource_default_access_policies(
                     current_resource_type,
                     tenant_id=workspace_id,
@@ -877,7 +907,7 @@ def migrate_only_me_resource_whitelist_scopes_to_automatic_include(
         )
         migrated_count += 1
 
-        for batch in _workspace_member_account_id_batches(workspace_id, member_batch_size):
+        for batch in _workspace_non_admin_account_id_batches(workspace_id, operator_account_id, member_batch_size):
             _replace_resource_default_access_policies(
                 current_resource_type,
                 tenant_id=workspace_id,
@@ -1187,7 +1217,11 @@ def _report_backing_app_specific_whitelist(options: _AgentAccessBootstrapOptions
 
 
 def _write_agent_access_rows(options: _AgentAccessBootstrapOptions) -> None:
-    for batch in _workspace_member_account_id_batches(options.tenant_id, options.member_batch_size):
+    for batch in _workspace_non_admin_account_id_batches(
+        options.tenant_id,
+        options.operator_account_id,
+        options.member_batch_size,
+    ):
         RBACService.AgentAccess.replace_user_access_policies(
             tenant_id=options.tenant_id,
             account_id=options.operator_account_id,

api/controllers/common/errors.py

@@ -92,6 +92,23 @@ class NotFoundError(BaseHTTPException):
     code = 404
 
 
+class UnauthorizedError(BaseHTTPException):
+    error_code = "unauthorized"
+    code = HTTPStatus.UNAUTHORIZED
+    description = "Authentication is required."
+
+
+class InternalServerError(BaseHTTPException):
+    """Expose a safe response while retaining the original exception in server logs."""
+
+    error_code = "internal_server_error"
+    code = HTTPStatus.INTERNAL_SERVER_ERROR
+    description = (
+        "The server encountered an internal error and was unable to complete your request. "
+        "Either the server is overloaded or there is an error in the application."
+    )
+
+
 class InvalidArgumentError(BaseHTTPException):
     error_code = "invalid_param"
     code = 400

api/controllers/console/explore/audio.py

@@ -7,10 +7,10 @@ from http import HTTPStatus
 
 from flask import Response, request
 from flask_restx import Resource
-from werkzeug.exceptions import HTTPException, InternalServerError
 
 from controllers.common.audio_response import audio_binary_response
 from controllers.common.controller_schemas import TextToAudioPayload
+from controllers.common.errors import InternalServerError
 from controllers.common.fields import AudioBinaryResponse, AudioTranscriptResponse
 from controllers.common.schema import register_response_schema_models, register_schema_model
 from controllers.console import console_ns
@@ -27,14 +27,17 @@ from controllers.console.app.error import (
     SpeechToTextDisabledError,
     UnsupportedAudioTypeError,
 )
+from controllers.console.explore.error import AgentVersionNotFoundHTTPError
 from controllers.console.explore.installed_app_admission import get_installed_app
 from controllers.console.flask_admission import console_account_admission
 from controllers.console.wraps import model_validate
 from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
 from extensions.ext_application_services import application_services
 from graphon.model_runtime.errors.invoke import InvokeError
+from libs.exception import BaseHTTPException
 from libs.helper import dump_response
 from machinery.context import RequestContext
+from services.agent.errors import AgentVersionNotFoundError
 from services.app_definition_query_service import AppDefinitionUnavailableError
 from services.audio_types import AudioAppRef, AudioUpload
 from services.errors.app_model_config import AppModelConfigBrokenError
@@ -84,7 +87,9 @@ def _audio_errors[**P, R](view: Callable[P, R]) -> Callable[P, R]:
             raise ProviderModelCurrentlyNotSupportError() from error
         except InvokeError as error:
             raise CompletionRequestError(error.description) from error
-        except (HTTPException, ValueError):
+        except AgentVersionNotFoundError as error:
+            raise AgentVersionNotFoundHTTPError() from error
+        except (BaseHTTPException, ValueError):
             raise
         except Exception as error:
             logger.exception("Installed-app audio operation failed")
@@ -101,10 +106,15 @@ class ChatAudioApi(Resource):
     @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AudioTranscriptResponse.__name__])
     @console_account_admission()
     @get_installed_app
-    @_audio_errors
     def post(self, request_context: RequestContext, installed_app: InstalledAppRef) -> dict[str, object]:
+        # Let Flask handle multipart parsing errors before mapping audio operation failures.
         file = request.files.get("file")
         audio = AudioUpload(stream=file.stream, mime_type=file.mimetype) if file is not None else None
+        return self._transcribe(installed_app, audio)
+
+    @_audio_errors
+    def _transcribe(self, installed_app: InstalledAppRef, audio: AudioUpload | None) -> dict[str, object]:
+        """Transcribe an upload; None lets the audio service report a missing file."""
         transcript = application_services().app_audio.transcript_asr(
             app=AudioAppRef(
                 app_id=installed_app.app_id,

api/controllers/console/explore/completion.py

@@ -6,9 +6,9 @@ from uuid import UUID
 from flask import Response
 from flask_restx import Resource
 from pydantic import BaseModel, Field, field_validator
-from werkzeug.exceptions import InternalServerError, NotFound, Unauthorized
 
 import services
+from controllers.common.errors import InternalServerError, NotFoundError, UnauthorizedError
 from controllers.common.fields import SimpleResultResponse
 from controllers.common.schema import register_response_schema_models, register_schema_models
 from controllers.console.app.error import (
@@ -112,11 +112,11 @@ class CompletionApi(Resource):
         except InstalledAppNotCompletionError:
             raise NotCompletionAppError() from None
         except InstalledAppNotFoundError:
-            raise NotFound("Installed app not found") from None
+            raise NotFoundError("Installed app not found") from None
         except AccountNotFoundError:
-            raise Unauthorized("Account no longer exists.") from None
+            raise UnauthorizedError("Account no longer exists.") from None
         except services.errors.conversation.ConversationNotExistsError:
-            raise NotFound("Conversation Not Exists.")
+            raise NotFoundError("Conversation Not Exists.")
         except services.errors.conversation.ConversationCompletedError:
             raise ConversationCompletedError()
         except services.errors.app_model_config.AppModelConfigBrokenError:
@@ -192,11 +192,11 @@ class ChatApi(Resource):
         except InstalledAppNotChatError:
             raise NotChatAppError() from None
         except InstalledAppNotFoundError:
-            raise NotFound("Installed app not found") from None
+            raise NotFoundError("Installed app not found") from None
         except AccountNotFoundError:
-            raise Unauthorized("Account no longer exists.") from None
+            raise UnauthorizedError("Account no longer exists.") from None
         except services.errors.conversation.ConversationNotExistsError:
-            raise NotFound("Conversation Not Exists.")
+            raise NotFoundError("Conversation Not Exists.")
         except services.errors.conversation.ConversationCompletedError:
             raise ConversationCompletedError()
         except services.errors.app_model_config.AppModelConfigBrokenError:

api/controllers/console/explore/error.py

@@ -4,6 +4,12 @@ from core.logging.context import get_request_id
 from libs.exception import BaseHTTPException
 
 
+class AgentVersionNotFoundHTTPError(BaseHTTPException):
+    error_code = "agent_version_not_found_error"
+    description = "Agent config version not found."
+    code = HTTPStatus.NOT_FOUND
+
+
 class InstalledAppHTTPError(BaseHTTPException):
     """A Console error with safe details and a request ID for troubleshooting."""
 

api/controllers/console/explore/message.py

@@ -10,9 +10,9 @@ from uuid import UUID
 from flask import Response
 from flask_restx import Resource
 from pydantic import BaseModel
-from werkzeug.exceptions import HTTPException, InternalServerError, Unauthorized
 
 from controllers.common.controller_schemas import MessageFeedbackPayload, MessageListQuery
+from controllers.common.errors import InternalServerError, UnauthorizedError
 from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
 from controllers.console import console_ns
 from controllers.console.app.error import (
@@ -42,6 +42,7 @@ from fields.conversation_fields import ResultResponse
 from fields.message_fields import ExploreMessageInfiniteScrollPagination, SuggestedQuestionsResponse
 from graphon.model_runtime.errors.invoke import InvokeError
 from libs import helper
+from libs.exception import BaseHTTPException
 from machinery.context import RequestContext
 from services.account_errors import AccountNotFoundError
 from services.app_definition_query_service import AppDefinitionUnavailableError
@@ -84,7 +85,7 @@ def _message_errors[**P, R](view: Callable[P, R]) -> Callable[P, R]:
         except AppDefinitionUnavailableError as error:
             raise AppUnavailableError() from error
         except (AccountNotFoundError, SuggestedQuestionsActorNotFoundError) as error:
-            raise Unauthorized("Account no longer exists.") from error
+            raise UnauthorizedError("Account no longer exists.") from error
         except MessageNotChatAppError as error:
             raise NotChatAppError() from error
         except InstalledAppNotCompletionError as error:
@@ -109,7 +110,7 @@ def _message_errors[**P, R](view: Callable[P, R]) -> Callable[P, R]:
             raise ProviderModelCurrentlyNotSupportError() from error
         except InvokeError as error:
             raise CompletionRequestError(error.description) from error
-        except (HTTPException, ValueError):
+        except (BaseHTTPException, ValueError):
             raise
         except Exception as error:
             logger.exception("Installed-app message operation failed")

api/controllers/console/explore/saved_message.py

@@ -2,9 +2,9 @@ from http import HTTPStatus
 from uuid import UUID
 
 from flask_restx import Resource
-from werkzeug.exceptions import NotFound
 
 from controllers.common.controller_schemas import SavedMessageCreatePayload, SavedMessageListQuery
+from controllers.common.errors import NotFoundError
 from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
 from controllers.console import console_ns
 from controllers.console.explore.error import NotCompletionAppError
@@ -72,7 +72,7 @@ class SavedMessageListApi(Resource):
                 message_id=str(req_data.message_id),
             )
         except MessageNotExistsError:
-            raise NotFound("Message Not Exists.")
+            raise NotFoundError("Message Not Exists.")
 
         return ResultResponse(result="success").model_dump(mode="json")
 

api/controllers/console/explore/trial.py

@@ -10,11 +10,11 @@ from uuid import UUID
 from flask import Response, request
 from flask_restx import Resource
 from pydantic import AliasChoices, BaseModel, Field, field_validator
-from werkzeug.exceptions import HTTPException, InternalServerError, NotFound, Unauthorized
 
 import services
 from configs import dify_config
 from controllers.common.audio_response import audio_binary_response
+from controllers.common.errors import InternalServerError, NotFoundError, UnauthorizedError
 from controllers.common.fields import (
     AudioBinaryResponse,
     AudioTranscriptResponse,
@@ -46,17 +46,18 @@ from controllers.console.app.error import (
 )
 from controllers.console.app.preview_admission import get_preview_app
 from controllers.console.explore.error import (
-    AppPreviewOwnerUnavailableError as AppPreviewOwnerUnavailableHttpError,
-)
-from controllers.console.explore.error import (
-    AppPreviewSiteUnavailableError as AppPreviewSiteUnavailableHttpError,
-)
-from controllers.console.explore.error import (
+    AgentVersionNotFoundHTTPError,
     AppSuggestedQuestionsAfterAnswerDisabledError,
     NotChatAppError,
     NotCompletionAppError,
     NotWorkflowAppError,
 )
+from controllers.console.explore.error import (
+    AppPreviewOwnerUnavailableError as AppPreviewOwnerUnavailableHttpError,
+)
+from controllers.console.explore.error import (
+    AppPreviewSiteUnavailableError as AppPreviewSiteUnavailableHttpError,
+)
 from controllers.console.explore.trial_app_admission import get_trial_app, get_trial_app_for_stop
 from controllers.console.files import FILE_UPLOAD_PARAMS, upload_file_from_request_context
 from controllers.console.flask_admission import console_account_admission
@@ -82,6 +83,7 @@ from fields.message_fields import SuggestedQuestionsResponse
 from graphon.model_runtime.errors.invoke import InvokeError
 from graphon.variables import SecretVariable, VariableBase
 from libs import helper
+from libs.exception import BaseHTTPException
 from libs.helper import dump_response, to_timestamp, uuid_value
 from libs.stream import close_stream
 from libs.url_utils import normalize_api_base_url
@@ -89,6 +91,7 @@ from machinery.context import RequestContext
 from models.enums import CreatorUserRole
 from models.model import AppMode
 from services.account_errors import AccountNotFoundError
+from services.agent.errors import AgentVersionNotFoundError
 from services.app_definition_query_service import AppDefinitionUnavailableError
 from services.app_preview_query_service import (
     AppPreviewOwnerUnavailableError,
@@ -531,7 +534,7 @@ class TrialAppWorkflowRunApi(Resource):
         except AppDefinitionUnavailableError as error:
             raise AppUnavailableError() from error
         except AccountNotFoundError as error:
-            raise Unauthorized("Account no longer exists.") from error
+            raise UnauthorizedError("Account no longer exists.") from error
         except services.errors.app_model_config.AppModelConfigBrokenError as error:
             raise AppUnavailableError() from error
         except ProviderTokenNotInitError as ex:
@@ -600,9 +603,9 @@ class TrialChatApi(Resource):
         except AppDefinitionUnavailableError as error:
             raise AppUnavailableError() from error
         except AccountNotFoundError as error:
-            raise Unauthorized("Account no longer exists.") from error
+            raise UnauthorizedError("Account no longer exists.") from error
         except services.errors.conversation.ConversationNotExistsError:
-            raise NotFound("Conversation Not Exists.")
+            raise NotFoundError("Conversation Not Exists.")
         except services.errors.conversation.ConversationCompletedError:
             raise ConversationCompletedError()
         except services.errors.app_model_config.AppModelConfigBrokenError:
@@ -671,11 +674,11 @@ class TrialMessageSuggestedQuestionApi(Resource):
         except AppDefinitionUnavailableError as error:
             raise AppUnavailableError() from error
         except SuggestedQuestionsActorNotFoundError as error:
-            raise Unauthorized("Account no longer exists.") from error
+            raise UnauthorizedError("Account no longer exists.") from error
         except MessageNotExistsError:
-            raise NotFound("Message not found")
+            raise NotFoundError("Message not found")
         except ConversationNotExistsError:
-            raise NotFound("Conversation not found")
+            raise NotFoundError("Conversation not found")
         except SuggestedQuestionsAfterAnswerDisabledError:
             raise AppSuggestedQuestionsAfterAnswerDisabledError()
         except ProviderTokenNotInitError as ex:
@@ -718,7 +721,9 @@ def _trial_audio_errors[**P, R](view: Callable[P, R]) -> Callable[P, R]:
             raise ProviderModelCurrentlyNotSupportError() from error
         except InvokeError as error:
             raise CompletionRequestError(error.description) from error
-        except (HTTPException, ValueError):
+        except AgentVersionNotFoundError as error:
+            raise AgentVersionNotFoundHTTPError() from error
+        except (BaseHTTPException, ValueError):
             raise
         except Exception as error:
             logger.exception("Trial audio operation failed")
@@ -735,10 +740,17 @@ class TrialChatAudioApi(Resource):
     @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[AudioTranscriptResponse.__name__])
     @console_account_admission()
     @get_trial_app
-    @_trial_audio_errors
     def post(self, request_context: RequestContext, trial_app: TrialAppRef) -> dict[str, object]:
+        # Let Flask handle multipart parsing errors before mapping audio operation failures.
         file = request.files.get("file")
         audio = AudioUpload(stream=file.stream, mime_type=file.mimetype) if file is not None else None
+        return self._transcript_asr(request_context=request_context, trial_app=trial_app, audio=audio)
+
+    @_trial_audio_errors
+    def _transcript_asr(
+        self, *, request_context: RequestContext, trial_app: TrialAppRef, audio: AudioUpload | None
+    ) -> dict[str, object]:
+        """Transcribe an upload; None lets the audio service report a missing file."""
         transcript = application_services().app_audio.transcript_asr(
             app=AudioAppRef(app_id=trial_app.app_id, tenant_id=trial_app.tenant_id, app_mode=trial_app.app_mode),
             audio=audio,
@@ -807,9 +819,9 @@ class TrialCompletionApi(Resource):
         except AppDefinitionUnavailableError as error:
             raise AppUnavailableError() from error
         except AccountNotFoundError as error:
-            raise Unauthorized("Account no longer exists.") from error
+            raise UnauthorizedError("Account no longer exists.") from error
         except services.errors.conversation.ConversationNotExistsError:
-            raise NotFound("Conversation Not Exists.")
+            raise NotFoundError("Conversation Not Exists.")
         except services.errors.conversation.ConversationCompletedError:
             raise ConversationCompletedError()
         except services.errors.app_model_config.AppModelConfigBrokenError:
@@ -890,7 +902,7 @@ class AppApi(Resource):
         except AppDefinitionUnavailableError as error:
             raise AppUnavailableError() from error
         except AccountNotFoundError as error:
-            raise Unauthorized("Account no longer exists.") from error
+            raise UnauthorizedError("Account no longer exists.") from error
         except AppPreviewSiteUnavailableError as error:
             raise AppPreviewSiteUnavailableHttpError(str(error)) from error
 

api/controllers/console/explore/workflow.py

@@ -3,9 +3,9 @@ from http import HTTPStatus
 
 from flask import Response
 from flask_restx import Resource
-from werkzeug.exceptions import InternalServerError, Unauthorized
 
 from controllers.common.controller_schemas import WorkflowRunPayload
+from controllers.common.errors import InternalServerError, UnauthorizedError
 from controllers.common.fields import SimpleResultResponse
 from controllers.common.schema import register_response_schema_models, register_schema_model
 from controllers.console.app.error import (
@@ -70,7 +70,7 @@ class InstalledAppWorkflowRunApi(Resource):
         except (AppDefinitionUnavailableError, InstalledAppNotWorkflowError):
             raise NotWorkflowAppError() from None
         except AccountNotFoundError:
-            raise Unauthorized("Account no longer exists.") from None
+            raise UnauthorizedError("Account no longer exists.") from None
         except ProviderTokenNotInitError as ex:
             raise ProviderNotInitializeError(ex.description)
         except QuotaExceededError:

api/controllers/console/extension.py

@@ -6,17 +6,23 @@ from flask_restx import Resource
 from pydantic import BaseModel, Field, RootModel, field_validator
 
 from constants import HIDDEN_VALUE
-from extensions.ext_database import db
+from controllers.console.flask_admission import console_account_admission
+from extensions.ext_application_services import application_services
 from fields.base import ResponseModel
 from libs.helper import dump_response, to_timestamp
 from libs.login import login_required
-from models.api_based_extension import APIBasedExtension
-from services.api_based_extension_service import APIBasedExtensionService
+from machinery.context import RequestContext
+from services.api_based_extension_application_service import (
+    APIBasedExtensionError,
+    APIBasedExtensionInput,
+    APIBasedExtensionRecord,
+    APIBasedExtensionUpdate,
+)
 from services.code_based_extension_service import CodeBasedExtensionService
 
 from ..common.schema import query_params_from_model, register_response_schema_models, register_schema_models
 from . import console_ns
-from .wraps import account_initialization_required, model_validate, setup_required, with_current_tenant_id
+from .wraps import account_initialization_required, model_validate, setup_required
 
 
 class CodeBasedExtensionQuery(BaseModel):
@@ -77,6 +83,10 @@ register_response_schema_models(
 )
 
 
+def _extension_response(extension: APIBasedExtensionRecord) -> dict[str, Any]:
+    return dump_response(APIBasedExtensionResponse, extension)
+
+
 @console_ns.route("/code-based-extension")
 class CodeBasedExtensionAPI(Resource):
     @console_ns.doc("get_code_based_extension")
@@ -103,41 +113,32 @@ class APIBasedExtensionAPI(Resource):
     @console_ns.doc("get_api_based_extensions")
     @console_ns.doc(description="Get all API-based extensions for current tenant")
     @console_ns.response(200, "Success", console_ns.models[APIBasedExtensionListResponse.__name__])
-    @setup_required
-    @login_required
-    @account_initialization_required
-    @with_current_tenant_id
-    def get(self, current_tenant_id: str):
-        return dump_response(
-            APIBasedExtensionListResponse,
-            APIBasedExtensionService.get_all_by_tenant_id(current_tenant_id, session=db.session()),
-        )
+    @console_account_admission()
+    def get(self, request_context: RequestContext):
+        extensions = application_services().api_based_extensions.list_extensions(request_context)
+        return dump_response(APIBasedExtensionListResponse, extensions)
 
     @console_ns.doc("create_api_based_extension")
     @console_ns.doc(description="Create a new API-based extension")
     @console_ns.expect(console_ns.models[APIBasedExtensionPayload.__name__])
     @console_ns.response(201, "Extension created successfully", console_ns.models[APIBasedExtensionResponse.__name__])
-    @setup_required
-    @login_required
-    @account_initialization_required
-    @with_current_tenant_id
+    @console_account_admission()
     @model_validate(APIBasedExtensionPayload)
-    def post(self, req_data: APIBasedExtensionPayload, current_tenant_id: str):
-        extension_data = APIBasedExtension(
-            tenant_id=current_tenant_id,
-            name=req_data.name,
-            api_endpoint=req_data.api_endpoint,
-            api_key=req_data.api_key,
-        )
+    def post(self, req_data: APIBasedExtensionPayload, request_context: RequestContext):
+        try:
+            extension = application_services().api_based_extensions.create_extension(
+                request_context,
+                APIBasedExtensionInput(
+                    name=req_data.name,
+                    api_endpoint=req_data.api_endpoint,
+                    api_key=req_data.api_key,
+                ),
+            )
+        except APIBasedExtensionError as error:
+            # Legacy clients receive the 400 `invalid_param` envelope for every extension failure.
+            raise ValueError(str(error)) from None
 
-        extension = APIBasedExtensionService.save(extension_data, session=db.session())
-        return APIBasedExtensionResponse(
-            id=extension.id,
-            name=extension.name,
-            api_endpoint=extension.api_endpoint,
-            api_key=req_data.api_key,
-            created_at=to_timestamp(extension.created_at),
-        ).model_dump(mode="json"), 201
+        return _extension_response(extension), 201
 
 
 @console_ns.route("/api-based-extension/<uuid:id>")
@@ -146,70 +147,48 @@ class APIBasedExtensionDetailAPI(Resource):
     @console_ns.doc(description="Get API-based extension by ID")
     @console_ns.doc(params={"id": "Extension ID"})
     @console_ns.response(200, "Success", console_ns.models[APIBasedExtensionResponse.__name__])
-    @setup_required
-    @login_required
-    @account_initialization_required
-    @with_current_tenant_id
-    def get(self, current_tenant_id: str, id: UUID):
-        api_based_extension_id = str(id)
+    @console_account_admission()
+    def get(self, request_context: RequestContext, id: UUID):
+        try:
+            extension = application_services().api_based_extensions.get_extension(request_context, str(id))
+        except APIBasedExtensionError as error:
+            raise ValueError(str(error)) from None
 
-        return dump_response(
-            APIBasedExtensionResponse,
-            APIBasedExtensionService.get_with_tenant_id(
-                current_tenant_id, api_based_extension_id, session=db.session()
-            ),
-        )
+        return _extension_response(extension)
 
     @console_ns.doc("update_api_based_extension")
     @console_ns.doc(description="Update API-based extension")
     @console_ns.doc(params={"id": "Extension ID"})
     @console_ns.expect(console_ns.models[APIBasedExtensionPayload.__name__])
     @console_ns.response(200, "Extension updated successfully", console_ns.models[APIBasedExtensionResponse.__name__])
-    @setup_required
-    @login_required
-    @account_initialization_required
-    @with_current_tenant_id
+    @console_account_admission()
     @model_validate(APIBasedExtensionPayload)
-    def post(self, req_data: APIBasedExtensionPayload, current_tenant_id: str, id: UUID):
-        api_based_extension_id = str(id)
+    def post(self, req_data: APIBasedExtensionPayload, request_context: RequestContext, id: UUID):
+        try:
+            extension = application_services().api_based_extensions.update_extension(
+                request_context,
+                str(id),
+                APIBasedExtensionUpdate(
+                    name=req_data.name,
+                    api_endpoint=req_data.api_endpoint,
+                    # The console echoes the masked key back; keep the stored key when it is sent unchanged.
+                    api_key=None if req_data.api_key == HIDDEN_VALUE else req_data.api_key,
+                ),
+            )
+        except APIBasedExtensionError as error:
+            raise ValueError(str(error)) from None
 
-        extension_data_from_db = APIBasedExtensionService.get_with_tenant_id(
-            current_tenant_id, api_based_extension_id, session=db.session()
-        )
-
-        api_key_for_response = extension_data_from_db.api_key
-
-        extension_data_from_db.name = req_data.name
-        extension_data_from_db.api_endpoint = req_data.api_endpoint
-
-        if req_data.api_key != HIDDEN_VALUE:
-            extension_data_from_db.api_key = req_data.api_key
-            api_key_for_response = req_data.api_key
-
-        APIBasedExtensionService.save(extension_data_from_db, session=db.session())
-        return APIBasedExtensionResponse(
-            id=extension_data_from_db.id,
-            name=extension_data_from_db.name,
-            api_endpoint=extension_data_from_db.api_endpoint,
-            api_key=api_key_for_response,
-            created_at=to_timestamp(extension_data_from_db.created_at),
-        ).model_dump(mode="json")
+        return _extension_response(extension)
 
     @console_ns.doc("delete_api_based_extension")
     @console_ns.doc(description="Delete API-based extension")
     @console_ns.doc(params={"id": "Extension ID"})
     @console_ns.response(204, "Extension deleted successfully")
-    @setup_required
-    @login_required
-    @account_initialization_required
-    @with_current_tenant_id
-    def delete(self, current_tenant_id: str, id: UUID):
-        api_based_extension_id = str(id)
-
-        extension_data_from_db = APIBasedExtensionService.get_with_tenant_id(
-            current_tenant_id, api_based_extension_id, session=db.session()
-        )
-
-        APIBasedExtensionService.delete(extension_data_from_db, session=db.session())
+    @console_account_admission()
+    def delete(self, request_context: RequestContext, id: UUID):
+        try:
+            application_services().api_based_extensions.delete_extension(request_context, str(id))
+        except APIBasedExtensionError as error:
+            raise ValueError(str(error)) from None
 
         return "", 204

api/controllers/service_api/app/audio.py

@@ -2,12 +2,14 @@ import logging
 
 from flask import request
 from flask_restx import Resource
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import InternalServerError
 
 import services
 from controllers.common.controller_schemas import TextToAudioPayload
 from controllers.common.fields import AudioBinaryResponse, AudioTranscriptResponse
 from controllers.common.schema import register_response_schema_models, register_schema_model
+from controllers.common.session import with_session
 from controllers.console.wraps import model_validate
 from controllers.service_api import service_api_ns
 from controllers.service_api.app.error import (
@@ -26,7 +28,6 @@ from controllers.service_api.schema import binary_response, expect_with_user, mu
 from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token
 from core.base.tts.audio_mime import SUPPORTED_TTS_AUDIO_MIME_TYPES
 from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
-from extensions.ext_database import db
 from graphon.model_runtime.errors.invoke import InvokeError
 from libs.helper import dump_response
 from models.model import App, EndUser
@@ -98,7 +99,8 @@ class AudioApi(Resource):
         service_api_ns.models[AudioTranscriptResponse.__name__],
     )
     @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.FORM))
-    def post(self, app_model: App, end_user: EndUser):
+    @with_session(write=False)
+    def post(self, session: Session, app_model: App, end_user: EndUser):
         """Convert audio to text using speech-to-text.
 
         Accepts an audio file upload and returns the transcribed text.
@@ -109,7 +111,7 @@ class AudioApi(Resource):
             response = AudioService.transcript_asr(
                 app_model=app_model,
                 file=file,
-                session=db.session(),
+                session=session,
                 end_user=end_user.id,
             )
 
@@ -182,8 +184,9 @@ class TextApi(Resource):
     # TTS returns provider audio bytes, so the success response is intentionally schema-less.
     @service_api_ns.response(200, "Text successfully converted to audio")
     @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.JSON))
+    @with_session(write=False)
     @model_validate(TextToAudioPayload)
-    def post(self, payload: TextToAudioPayload, app_model: App, end_user: EndUser):
+    def post(self, payload: TextToAudioPayload, session: Session, app_model: App, end_user: EndUser):
         """Convert text to audio using text-to-speech.
 
         Converts the provided text to audio using the specified voice.
@@ -202,7 +205,7 @@ class TextApi(Resource):
                 )
             return AudioService.transcript_tts(
                 app_model=app_model,
-                session=db.session(),
+                session=session,
                 text=text,
                 voice=voice,
                 end_user=end_user.external_user_id,

api/controllers/web/audio.py

@@ -2,10 +2,12 @@ import logging
 
 from flask import request
 from pydantic import field_validator
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import InternalServerError
 
 import services
 from controllers.common.controller_schemas import TextToAudioPayload as TextToAudioPayloadBase
+from controllers.common.session import with_session
 from controllers.console.wraps import model_validate
 from controllers.web import web_ns
 from controllers.web.error import (
@@ -22,7 +24,6 @@ from controllers.web.error import (
 )
 from controllers.web.wraps import WebApiResource
 from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
-from extensions.ext_database import db
 from fields.base import ResponseModel
 from graphon.model_runtime.errors.invoke import InvokeError
 from libs.helper import dump_response, uuid_value
@@ -75,7 +76,8 @@ class AudioApi(WebApiResource):
         }
     )
     @web_ns.response(200, "Success", web_ns.models[AudioToTextResponse.__name__])
-    def post(self, app_model: App, end_user: EndUser):
+    @with_session(write=False)
+    def post(self, session: Session, app_model: App, end_user: EndUser):
         """Convert audio to text"""
         file = request.files.get("file")
 
@@ -83,7 +85,7 @@ class AudioApi(WebApiResource):
             response = AudioService.transcript_asr(
                 app_model=app_model,
                 file=file,
-                session=db.session(),
+                session=session,
                 end_user=end_user.external_user_id,
             )
 
@@ -132,8 +134,9 @@ class TextApi(WebApiResource):
     )
     # response-contract:ignore provider audio bytes; TODO: model binary audio response if shape is standardized.
     @web_ns.response(200, "Success")
+    @with_session(write=False)
     @model_validate(TextToAudioPayload)
-    def post(self, payload: TextToAudioPayload, app_model: App, end_user: EndUser):
+    def post(self, payload: TextToAudioPayload, session: Session, app_model: App, end_user: EndUser):
         """Convert text to audio"""
         try:
             message_id = payload.message_id
@@ -149,7 +152,7 @@ class TextApi(WebApiResource):
                 )
             return AudioService.transcript_tts(
                 app_model=app_model,
-                session=db.session(),
+                session=session,
                 text=text,
                 voice=voice,
                 end_user=end_user.external_user_id,

api/core/rag/datasource/vdb/vector_factory.py

@@ -29,7 +29,16 @@ logger = logging.getLogger(__name__)
 
 class AbstractVectorFactory(ABC):
     @abstractmethod
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> BaseVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> BaseVector:
+        """Initialize the backend using the caller's session for metadata reads.
+
+        ``None`` is used by callers that initialize a resolved vector backend
+        outside their database transaction; backend-specific legacy reads may
+        still use their own session in that case. A supplied session belongs to
+        the caller and must not be committed, rolled back, or closed here.
+        """
         raise NotImplementedError
 
     @staticmethod
@@ -142,7 +151,7 @@ class Vector:
         self._session = session
         if vector_type is not None:
             self._vector_processor = self.get_vector_factory(vector_type)().init_vector(
-                dataset, self._attributes, self._embeddings
+                dataset, self._attributes, self._embeddings, session=session
             )
         else:
             if session is None:
@@ -172,7 +181,7 @@ class Vector:
     def _init_vector(self, *, session: Session) -> BaseVector:
         vector_type = self.resolve_vector_type(self._dataset, session=session)
         vector_factory_cls = self.get_vector_factory(vector_type)
-        return vector_factory_cls().init_vector(self._dataset, self._attributes, self._embeddings)
+        return vector_factory_cls().init_vector(self._dataset, self._attributes, self._embeddings, session=session)
 
     @staticmethod
     def get_vector_factory(vector_type: str) -> type[AbstractVectorFactory]:

api/extensions/ext_application_services.py

@@ -49,6 +49,7 @@ from models.model import EndUser
 from repositories.account.repository import SQLAlchemyAccountRepository
 from repositories.account_activation_repository import SQLAlchemyAccountActivationRepository
 from repositories.account_integration_repository import SQLAlchemyAccountIntegrationRepository
+from repositories.api_based_extension_repository import APIBasedExtensionRepository
 from repositories.app.mcp_server_repository import AppMCPServerRepository
 from repositories.app.site_command_repository import AppSiteCommandRepository
 from repositories.app.tracing_config_repository import SQLAlchemyAppTracingConfigRepository
@@ -92,6 +93,8 @@ from services.account.login_adapters import RedisConsoleAuthSecurityGateway
 from services.account.service import AccountSetupProvisioner
 from services.account_password_hasher import DefaultAccountPasswordHasher
 from services.agent.roster_package_exporter import RosterAgentPackageExporter
+from services.api_based_extension_adapters import APIBasedExtensionPingProbe, WorkspaceTokenCipher
+from services.api_based_extension_application_service import APIBasedExtensionApplicationService
 from services.app.advanced_prompt_template_service import AdvancedPromptTemplateService
 from services.app.api_key_service import AppApiKeyService
 from services.app.mcp_server_service import AppMCPServerService
@@ -240,6 +243,7 @@ class AppScopedEndUserServices:
 class ApplicationServices:
     agent_apps: AgentAppServices
     advanced_prompt_templates: AdvancedPromptTemplateService
+    api_based_extensions: APIBasedExtensionApplicationService
     credential_queries: CredentialQuery
     accounts: AccountServices
     app_api_keys: AppApiKeyService
@@ -677,6 +681,11 @@ def build_application_services(
             now=lambda: datetime.now(UTC),
             access_token_expire_minutes=dify_config.ACCESS_TOKEN_EXPIRE_MINUTES,
         ),
+        api_based_extensions=APIBasedExtensionApplicationService(
+            extensions=APIBasedExtensionRepository(session_factory=database_client),
+            secrets=WorkspaceTokenCipher(),
+            probe=APIBasedExtensionPingProbe(),
+        ),
         tags=tags,
         workflow_statistics=WorkflowStatisticQueryService(
             workflow_runs=DifyAPIRepositoryFactory.create_api_workflow_run_repository(

api/libs/broadcast_channel/redis/_subscription.py

@@ -196,6 +196,8 @@ class RedisSubscriptionBase(Subscription):
         except queue.Empty:
             return None
 
+        if item == SIG_CLOSE:
+            raise SubscriptionClosedError(f"The Redis {self._get_subscription_type()} subscription is closed")
         return item
 
     @override

api/libs/schedule_utils.py

@@ -1,7 +1,8 @@
 from datetime import UTC, datetime
+from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 
-import pytz  # type: ignore[import-untyped]
 from croniter import croniter
+from pytz.exceptions import UnknownTimeZoneError
 
 
 def calculate_next_run_at(
@@ -38,7 +39,12 @@ def calculate_next_run_at(
             f"(@daily, @weekly, etc.). Got {len(parts)} fields: '{cron_expression}'"
         )
 
-    tz = pytz.timezone(timezone)
+    try:
+        tz = ZoneInfo(timezone)
+    except ZoneInfoNotFoundError as exc:
+        # Preserve the public error contract: callers and tests expect the same
+        # exception type pytz raised for unknown timezone names.
+        raise UnknownTimeZoneError(timezone) from exc
 
     if base_time is None:
         base_time = datetime.now(UTC)

api/providers/vdb/vdb-alibabacloud-mysql/src/dify_vdb_alibabacloud_mysql/alibabacloud_mysql_vector.py

@@ -8,6 +8,7 @@ from typing import Any, Literal, cast, override
 import mysql.connector
 from mysql.connector import Error as MySQLError
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import parse_metadata_json
@@ -365,7 +366,9 @@ class AlibabaCloudMySQLVectorFactory(AbstractVectorFactory):
         return cast(Literal["cosine", "euclidean"], distance_function)
 
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> AlibabaCloudMySQLVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> AlibabaCloudMySQLVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-alibabacloud-mysql/tests/unit_tests/test_alibabacloud_mysql_factory.py

@@ -39,7 +39,7 @@ def test_factory_init_vector_uses_existing_index_struct_class_prefix(monkeypatch
     monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_HNSW_M", 6)
 
     with patch.object(alibaba_module, "AlibabaCloudMySQLVector", return_value="vector") as vector_cls:
-        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result == "vector"
     assert vector_cls.call_args.kwargs["collection_name"] == "existing_collection"
@@ -61,7 +61,7 @@ def test_factory_init_vector_generates_collection_name_when_index_struct_is_miss
     monkeypatch.setattr(alibaba_module.dify_config, "ALIBABACLOUD_MYSQL_HNSW_M", 12)
 
     with patch.object(alibaba_module, "AlibabaCloudMySQLVector", return_value="vector") as vector_cls:
-        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result == "vector"
     vector_cls.assert_called_once()

api/providers/vdb/vdb-analyticdb/src/dify_vdb_analyticdb/analyticdb_vector.py

@@ -1,6 +1,8 @@
 import json
 from typing import Any, override
 
+from sqlalchemy.orm import Session
+
 from configs import dify_config
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
@@ -74,7 +76,9 @@ class AnalyticdbVector(BaseVector):
 
 class AnalyticdbVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> AnalyticdbVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> AnalyticdbVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-analyticdb/tests/unit_tests/test_analyticdb_vector.py

@@ -100,7 +100,7 @@ def test_factory_builds_openapi_config_when_host_is_missing(monkeypatch: pytest.
     monkeypatch.setattr(analyticdb_module.dify_config, "ANALYTICDB_NAMESPACE_PASSWORD", "ns-password")
 
     with patch.object(analyticdb_module, "AnalyticdbVector", return_value="vector") as vector_cls:
-        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result == "vector"
     args = vector_cls.call_args.args
@@ -123,7 +123,7 @@ def test_factory_builds_sql_config_when_host_is_present(monkeypatch: pytest.Monk
     monkeypatch.setattr(analyticdb_module.dify_config, "ANALYTICDB_NAMESPACE", "dify")
 
     with patch.object(analyticdb_module, "AnalyticdbVector", return_value="vector") as vector_cls:
-        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result == "vector"
     args = vector_cls.call_args.args

api/providers/vdb/vdb-baidu/src/dify_vdb_baidu/baidu_vector.py

@@ -26,6 +26,7 @@ from pymochow.model.schema import (
     VectorIndex,
 )  # type: ignore
 from pymochow.model.table import AnnSearch, BM25SearchRequest, HNSWSearchParams, Partition, Row  # type: ignore
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import Field as VDBField
@@ -378,7 +379,9 @@ class BaiduVector(BaseVector):
 
 class BaiduVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> BaiduVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> BaiduVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-baidu/tests/unit_tests/test_baidu_vector.py

@@ -274,7 +274,7 @@ def test_factory_initializes_collection_name_and_index_struct(baidu_module, monk
     monkeypatch.setattr(baidu_module.dify_config, "BAIDU_VECTOR_DB_REBUILD_INDEX_TIMEOUT_IN_SECONDS", 300)
 
     with patch.object(baidu_module, "BaiduVector", return_value="vector") as vector_cls:
-        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result == "vector"
     assert vector_cls.call_args.kwargs["collection_name"] == "auto_collection"
@@ -547,7 +547,7 @@ def test_factory_uses_existing_collection_prefix_when_index_struct_exists(
     monkeypatch.setattr(baidu_module.dify_config, "BAIDU_VECTOR_DB_REBUILD_INDEX_TIMEOUT_IN_SECONDS", 300)
 
     with patch.object(baidu_module, "BaiduVector", return_value="vector") as vector_cls:
-        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result == "vector"
     assert vector_cls.call_args.kwargs["collection_name"] == "existing_collection"

api/providers/vdb/vdb-chroma/src/dify_vdb_chroma/chroma_vector.py

@@ -4,6 +4,7 @@ from typing import Any, TypedDict, override
 import chromadb
 from chromadb import QueryResult, Settings
 from pydantic import BaseModel
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.vector_base import BaseVector, VectorIndexStructDict
@@ -158,7 +159,9 @@ class ChromaVector(BaseVector):
 
 class ChromaVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> BaseVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> BaseVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-chroma/tests/unit_tests/test_chroma_vector.py

@@ -190,8 +190,8 @@ def test_factory_init_vector_uses_existing_or_generated_collection(chroma_module
     monkeypatch.setattr(chroma_module.dify_config, "CHROMA_AUTH_CREDENTIALS", None)
 
     with patch.object(chroma_module, "ChromaVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-clickzetta/src/dify_vdb_clickzetta/clickzetta_vector.py

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, override
 
 import clickzetta  # type: ignore
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 
 if TYPE_CHECKING:
     from clickzetta.connector.v0.connection import Connection  # type: ignore
@@ -1069,7 +1070,9 @@ class ClickzettaVectorFactory(AbstractVectorFactory):
     """Factory for creating Clickzetta vector instances."""
 
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> BaseVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> BaseVector:
         """Initialize a Clickzetta vector instance."""
         # Get configuration from environment variables or dataset config
         config = ClickzettaConfig(

api/providers/vdb/vdb-clickzetta/tests/unit_tests/test_clickzetta_vector.py

@@ -238,7 +238,7 @@ def test_factory_initializes_clickzetta_vector(clickzetta_module, monkeypatch: p
     monkeypatch.setattr(clickzetta_module.dify_config, "CLICKZETTA_VECTOR_DISTANCE_FUNCTION", "cosine_distance")
 
     with patch.object(clickzetta_module, "ClickzettaVector", return_value="vector") as vector_cls:
-        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result == "vector"
     assert vector_cls.call_args.kwargs["collection_name"] == "collection"

api/providers/vdb/vdb-couchbase/src/dify_vdb_couchbase/couchbase_vector.py

@@ -15,6 +15,7 @@ from couchbase.options import ClusterOptions, SearchOptions  # type: ignore
 from couchbase.vector_search import VectorQuery, VectorSearch  # type: ignore
 from flask import current_app
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
@@ -366,7 +367,9 @@ class CouchbaseVector(BaseVector):
 
 class CouchbaseVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> CouchbaseVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> CouchbaseVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-couchbase/tests/unit_tests/test_couchbase_vector.py

@@ -357,8 +357,8 @@ def test_delete_collection_and_factory(couchbase_module, monkeypatch: pytest.Mon
     )
 
     with patch.object(couchbase_module, "CouchbaseVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-elasticsearch/src/dify_vdb_elasticsearch/elasticsearch_ja_vector.py

@@ -3,6 +3,7 @@ import logging
 from typing import Any, override
 
 from flask import current_app
+from sqlalchemy.orm import Session
 
 from core.rag.datasource.vdb.field import Field
 from core.rag.datasource.vdb.vector_type import VectorType
@@ -87,7 +88,9 @@ class ElasticSearchJaVector(ElasticSearchVector):
 
 class ElasticSearchJaVectorFactory(ElasticSearchVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> ElasticSearchJaVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> ElasticSearchJaVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-elasticsearch/src/dify_vdb_elasticsearch/elasticsearch_vector.py

@@ -9,6 +9,7 @@ from elasticsearch import Elasticsearch
 from flask import current_app
 from packaging.version import parse as parse_version
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 
 from core.rag.datasource.vdb.field import Field
 from core.rag.datasource.vdb.vector_base import BaseVector
@@ -307,7 +308,9 @@ class ElasticSearchVector(BaseVector):
 
 class ElasticSearchVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> ElasticSearchVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> ElasticSearchVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-elasticsearch/tests/unit_tests/test_elasticsearch_ja_vector.py

@@ -112,8 +112,8 @@ def test_ja_factory_uses_existing_or_generated_collection(elasticsearch_ja_modul
     )
 
     with patch.object(elasticsearch_ja_module, "ElasticSearchJaVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-elasticsearch/tests/unit_tests/test_elasticsearch_vector.py

@@ -358,7 +358,7 @@ def test_elasticsearch_factory_branches(elasticsearch_module, monkeypatch: pytes
     )
 
     with patch.object(elasticsearch_module, "ElasticSearchVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
     assert result_1 == "vector"
     cfg = vector_cls.call_args.kwargs["config"]
     assert cfg.use_cloud is False
@@ -377,7 +377,7 @@ def test_elasticsearch_factory_branches(elasticsearch_module, monkeypatch: pytes
         ),
     )
     with patch.object(elasticsearch_module, "ElasticSearchVector", return_value="vector") as vector_cls:
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
     assert result_2 == "vector"
     cfg = vector_cls.call_args.kwargs["config"]
     assert cfg.use_cloud is True
@@ -399,7 +399,7 @@ def test_elasticsearch_factory_branches(elasticsearch_module, monkeypatch: pytes
         ),
     )
     with patch.object(elasticsearch_module, "ElasticSearchVector", return_value="vector") as vector_cls:
-        factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
     cfg = vector_cls.call_args.kwargs["config"]
     assert cfg.use_cloud is False
     assert cfg.host == "fallback-host"

api/providers/vdb/vdb-hologres/src/dify_vdb_hologres/hologres_vector.py

@@ -7,6 +7,7 @@ import holo_search_sdk as holo  # type: ignore
 from holo_search_sdk.types import BaseQuantizationType, DistanceType, TokenizerType
 from psycopg import sql as psql
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import parse_metadata_json
@@ -344,7 +345,9 @@ class HologresVectorFactory(AbstractVectorFactory):
     """Factory class for creating HologresVector instances."""
 
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> HologresVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> HologresVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-hologres/tests/unit_tests/test_hologres_vector.py

@@ -353,8 +353,8 @@ def test_hologres_factory_uses_existing_or_generated_collection(hologres_module,
     monkeypatch.setattr(hologres_module.dify_config, "HOLOGRES_EF_CONSTRUCTION", 400)
 
     with patch.object(hologres_module, "HologresVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-huawei-cloud/src/dify_vdb_huawei_cloud/huawei_cloud_vector.py

@@ -5,6 +5,7 @@ from typing import Any, override
 
 from elasticsearch import Elasticsearch
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 from typing_extensions import TypedDict
 
 from configs import dify_config
@@ -217,7 +218,9 @@ class HuaweiCloudVector(BaseVector):
 
 class HuaweiCloudVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> HuaweiCloudVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> HuaweiCloudVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-huawei-cloud/tests/unit_tests/test_huawei_cloud_vector.py

@@ -233,8 +233,8 @@ def test_huawei_factory_branches(huawei_module, monkeypatch: pytest.MonkeyPatch)
     monkeypatch.setattr(huawei_module.dify_config, "HUAWEI_CLOUD_PASSWORD", "pass")
 
     with patch.object(huawei_module, "HuaweiCloudVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-iris/src/dify_vdb_iris/iris_vector.py

@@ -13,6 +13,8 @@ import uuid
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, override
 
+from sqlalchemy.orm import Session
+
 from configs import dify_config
 from configs.middleware.vdb.iris_config import IrisVectorConfig
 from core.rag.datasource.vdb.field import parse_metadata_json
@@ -491,7 +493,9 @@ class IrisVectorFactory(AbstractVectorFactory):
     """Factory for creating IrisVector instances."""
 
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> IrisVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> IrisVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-iris/tests/unit_tests/test_iris_vector.py

@@ -402,8 +402,8 @@ def test_iris_vector_delete_create_collection_and_factory(iris_module, monkeypat
     monkeypatch.setattr(iris_module.dify_config, "IRIS_TEXT_INDEX_LANGUAGE", "en")
 
     with patch.object(iris_module, "IrisVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-lindorm/src/dify_vdb_lindorm/lindorm_vector.py

@@ -6,6 +6,7 @@ from typing import Any, override
 from opensearchpy import OpenSearch, helpers
 from opensearchpy.helpers import BulkIndexError
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 from tenacity import retry, stop_after_attempt, wait_exponential
 from typing_extensions import TypedDict
 
@@ -388,7 +389,9 @@ class LindormVectorStore(BaseVector):
 
 class LindormVectorStoreFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> LindormVectorStore:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> LindormVectorStore:
         lindorm_config = LindormVectorStoreConfig(
             hosts=dify_config.LINDORM_URL,
             username=dify_config.LINDORM_USERNAME,

api/providers/vdb/vdb-lindorm/tests/unit_tests/test_lindorm_vector.py

@@ -361,7 +361,7 @@ def test_lindorm_factory_branches(lindorm_module, monkeypatch: pytest.MonkeyPatc
 
     monkeypatch.setattr(lindorm_module.dify_config, "LINDORM_USING_UGC", None)
     with pytest.raises(ValueError, match="LINDORM_USING_UGC is not set"):
-        factory.init_vector(dataset, attributes=[], embeddings=embeddings)
+        factory.init_vector(dataset, attributes=[], embeddings=embeddings, session=None)
 
     monkeypatch.setattr(lindorm_module.dify_config, "LINDORM_USING_UGC", False)
 
@@ -369,7 +369,7 @@ def test_lindorm_factory_branches(lindorm_module, monkeypatch: pytest.MonkeyPatc
         id="dataset-1", index_struct=json.dumps({"vector_store": {"class_prefix": "EXISTING"}, "using_ugc": False})
     )
     with patch.object(lindorm_module, "LindormVectorStore", return_value="vector") as store_cls:
-        result = factory.init_vector(dataset_existing_plain, attributes=[], embeddings=embeddings)
+        result = factory.init_vector(dataset_existing_plain, attributes=[], embeddings=embeddings, session=None)
     assert result == "vector"
     assert store_cls.call_args.args[0] == "existing"
 
@@ -386,7 +386,7 @@ def test_lindorm_factory_branches(lindorm_module, monkeypatch: pytest.MonkeyPatc
         ),
     )
     with patch.object(lindorm_module, "LindormVectorStore", return_value="vector") as store_cls:
-        factory.init_vector(dataset_existing_ugc, attributes=[], embeddings=embeddings)
+        factory.init_vector(dataset_existing_ugc, attributes=[], embeddings=embeddings, session=None)
     assert store_cls.call_args.args[0] == "ugc_index_1536_hnsw_l2"
     assert store_cls.call_args.kwargs["routing_value"] == "ROUTING"
 
@@ -394,7 +394,7 @@ def test_lindorm_factory_branches(lindorm_module, monkeypatch: pytest.MonkeyPatc
 
     monkeypatch.setattr(lindorm_module.dify_config, "LINDORM_USING_UGC", True)
     with patch.object(lindorm_module, "LindormVectorStore", return_value="vector") as store_cls:
-        factory.init_vector(dataset_new, attributes=[], embeddings=embeddings)
+        factory.init_vector(dataset_new, attributes=[], embeddings=embeddings, session=None)
     assert store_cls.call_args.args[0] == "ugc_index_3_hnsw_l2"
     assert store_cls.call_args.kwargs["routing_value"] == "auto_collection"
     assert dataset_new.index_struct is not None
@@ -402,6 +402,6 @@ def test_lindorm_factory_branches(lindorm_module, monkeypatch: pytest.MonkeyPatc
     dataset_new_plain = Dataset(id="dataset-3")
     monkeypatch.setattr(lindorm_module.dify_config, "LINDORM_USING_UGC", False)
     with patch.object(lindorm_module, "LindormVectorStore", return_value="vector") as store_cls:
-        factory.init_vector(dataset_new_plain, attributes=[], embeddings=embeddings)
+        factory.init_vector(dataset_new_plain, attributes=[], embeddings=embeddings, session=None)
     assert store_cls.call_args.args[0] == "auto_collection"
     assert store_cls.call_args.kwargs["routing_value"] is None

api/providers/vdb/vdb-matrixone/src/dify_vdb_matrixone/matrixone_vector.py

@@ -7,6 +7,7 @@ from typing import Any, Concatenate, override
 
 from mo_vector.client import MoVectorClient  # type: ignore
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import parse_metadata_json
@@ -228,7 +229,9 @@ class MatrixoneVector(BaseVector):
 
 class MatrixoneVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> MatrixoneVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> MatrixoneVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-matrixone/tests/unit_tests/test_matrixone_vector.py

@@ -246,8 +246,8 @@ def test_matrixone_factory_uses_existing_or_generated_collection(matrixone_modul
     monkeypatch.setattr(matrixone_module.dify_config, "MATRIXONE_METRIC", "l2")
 
     with patch.object(matrixone_module, "MatrixoneVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-milvus/src/dify_vdb_milvus/milvus_vector.py

@@ -6,6 +6,7 @@ from packaging import version
 from pydantic import BaseModel, model_validator
 from pymilvus import MilvusClient, MilvusException  # type: ignore
 from pymilvus.milvus_client import IndexParams  # type: ignore
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import Field
@@ -422,7 +423,9 @@ class MilvusVectorFactory(AbstractVectorFactory):
     """
 
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> MilvusVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> MilvusVector:
         """
         Initialize a MilvusVector instance for the given dataset.
         """

api/providers/vdb/vdb-milvus/tests/unit_tests/test_milvus.py

@@ -436,8 +436,8 @@ def test_factory_initializes_milvus_vector(milvus_module, monkeypatch: pytest.Mo
     monkeypatch.setattr(milvus_module.dify_config, "MILVUS_ANALYZER_PARAMS", '{"tokenizer":"standard"}')
 
     with patch.object(milvus_module, "MilvusVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-myscale/src/dify_vdb_myscale/myscale_vector.py

@@ -7,6 +7,7 @@ from typing import Any, override
 
 from clickhouse_connect import get_client  # type: ignore[import-untyped]
 from pydantic import BaseModel
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.vector_base import BaseVector
@@ -196,7 +197,9 @@ class MyScaleVector(BaseVector):
 
 class MyScaleVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> MyScaleVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> MyScaleVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-myscale/tests/unit_tests/test_myscale_vector.py

@@ -108,8 +108,8 @@ def test_factory_initializes_lower_case_collection_name(myscale_module, monkeypa
     monkeypatch.setattr(myscale_module.dify_config, "MYSCALE_FTS_PARAMS", "")
 
     with patch.object(myscale_module, "MyScaleVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-oceanbase/src/dify_vdb_oceanbase/oceanbase_vector.py

@@ -8,6 +8,7 @@ from pyobvector import VECTOR, ObVecClient, cosine_distance, inner_product, l2_d
 from sqlalchemy import JSON, Column, String
 from sqlalchemy.dialects.mysql import LONGTEXT
 from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import parse_metadata_json
@@ -564,6 +565,8 @@ class OceanBaseVectorFactory(AbstractVectorFactory):
         dataset: Dataset,
         attributes: list,
         embeddings: Embeddings,
+        *,
+        session: Session | None,
     ) -> BaseVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]

api/providers/vdb/vdb-oceanbase/tests/unit_tests/test_oceanbase_vector.py

@@ -591,8 +591,8 @@ def test_oceanbase_factory_uses_existing_or_generated_collection(oceanbase_modul
     monkeypatch.setattr(oceanbase_module.dify_config, "OCEANBASE_HNSW_REFRESH_THRESHOLD", 1000)
 
     with patch.object(oceanbase_module, "OceanBaseVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-opengauss/src/dify_vdb_opengauss/opengauss.py

@@ -6,6 +6,7 @@ from typing import Any, override
 import psycopg2.extras
 import psycopg2.pool
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.vector_base import BaseVector
@@ -250,7 +251,9 @@ class OpenGauss(BaseVector):
 
 class OpenGaussFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> OpenGauss:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> OpenGauss:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-opengauss/tests/unit_tests/test_opengauss.py

@@ -390,8 +390,8 @@ def test_opengauss_factory_uses_existing_or_generated_collection(opengauss_modul
     monkeypatch.setattr(opengauss_module.dify_config, "OPENGAUSS_ENABLE_PQ", False)
 
     with patch.object(opengauss_module, "OpenGauss", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-opensearch/src/dify_vdb_opensearch/opensearch_vector.py

@@ -6,6 +6,7 @@ from uuid import uuid4
 from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection, helpers
 from opensearchpy.helpers import BulkIndexError
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 from typing_extensions import TypedDict
 
 from configs import dify_config
@@ -310,7 +311,9 @@ class OpenSearchVector(BaseVector):
 
 class OpenSearchVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> OpenSearchVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> OpenSearchVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-opensearch/tests/unit_tests/test_opensearch_vector.py

@@ -352,8 +352,8 @@ def test_opensearch_factory_initializes_expected_collection_name(opensearch_modu
     monkeypatch.setattr(opensearch_module.dify_config, "OPENSEARCH_AWS_SERVICE", None)
 
     with patch.object(opensearch_module, "OpenSearchVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-oracle/src/dify_vdb_oracle/oraclevector.py

@@ -10,6 +10,7 @@ import numpy
 import oracledb
 from oracledb.connection import Connection
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.vector_base import BaseVector
@@ -383,7 +384,9 @@ class OracleVector(BaseVector):
 
 class OracleVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> OracleVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> OracleVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-oracle/tests/unit_tests/test_oraclevector.py

@@ -367,8 +367,8 @@ def test_oracle_factory_init_vector_uses_existing_or_generated_collection(
     monkeypatch.setattr(oracle_module.dify_config, "ORACLE_IS_AUTONOMOUS", False)
 
     with patch.object(oracle_module, "OracleVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-pgvecto-rs/src/dify_vdb_pgvecto_rs/pgvecto_rs.py

@@ -217,7 +217,9 @@ class PGVectoRS(BaseVector):
 
 class PGVectoRSFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> PGVectoRS:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> PGVectoRS:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-pgvecto-rs/tests/unit_tests/test_pgvecto_rs.py

@@ -334,8 +334,8 @@ def test_factory_uses_existing_or_generated_collection(pgvecto_module, monkeypat
     embeddings.embed_query.return_value = [0.1, 0.2, 0.3]
 
     with patch.object(module, "PGVectoRS", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=embeddings)
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=embeddings)
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=embeddings, session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=embeddings, session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-pgvector/src/dify_vdb_pgvector/pgvector.py

@@ -9,6 +9,7 @@ import psycopg2.errors
 import psycopg2.extras
 import psycopg2.pool
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.vector_base import BaseVector
@@ -280,7 +281,9 @@ class PGVector(BaseVector):
 
 class PGVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> PGVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> PGVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-pgvector/tests/unit_tests/test_pgvector.py

@@ -484,8 +484,8 @@ def test_pgvector_factory_initializes_expected_collection_name(monkeypatch: pyte
     monkeypatch.setattr(pgvector_module.dify_config, "PGVECTOR_PG_BIGM", False)
 
     with patch.object(pgvector_module, "PGVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-qdrant/src/dify_vdb_qdrant/qdrant_vector.py

@@ -19,6 +19,7 @@ from qdrant_client.http.models import (
 )
 from qdrant_client.local.qdrant_local import QdrantLocal
 from sqlalchemy import select
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.field import Field
@@ -497,10 +498,15 @@ class QdrantVector(BaseVector):
 
 class QdrantVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> QdrantVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> QdrantVector:
         if dataset.collection_binding_id:
             stmt = select(DatasetCollectionBinding).where(DatasetCollectionBinding.id == dataset.collection_binding_id)
-            dataset_collection_binding = db.session.scalars(stmt).one_or_none()
+            # Annotation bindings may only be flushed in the caller's transaction.
+            # Keep the legacy lookup for workflows initialized with session=None.
+            binding_session = session if session is not None else db.session
+            dataset_collection_binding = binding_session.scalars(stmt).one_or_none()
             if dataset_collection_binding:
                 collection_name = dataset_collection_binding.collection_name
             else:

api/providers/vdb/vdb-qdrant/tests/unit_tests/test_qdrant_vector.py

@@ -4,13 +4,16 @@ import os
 import sys
 import types
 from collections import UserDict
+from pathlib import Path
 from types import SimpleNamespace
 from typing import override
 from unittest.mock import MagicMock, patch
 
 import pytest
+from sqlalchemy import create_engine, select
 from sqlalchemy.orm import Session
 
+from core.rag.datasource.vdb.vector_factory import Vector
 from core.rag.models.document import Document
 from models.dataset import Dataset, DatasetCollectionBinding
 from models.enums import CollectionBindingType
@@ -295,12 +298,15 @@ def test_search_and_helper_methods(qdrant_module):
 
 
 @pytest.mark.parametrize("sqlite3_session", [(DatasetCollectionBinding,)], indirect=True)
+@pytest.mark.parametrize("use_caller_session", [True, False])
 def test_qdrant_factory_paths(
     qdrant_module,
     monkeypatch: pytest.MonkeyPatch,
     sqlite3_session: Session,
+    use_caller_session: bool,
 ):
     factory = qdrant_module.QdrantVectorFactory()
+    session = sqlite3_session if use_caller_session else None
     dataset = Dataset(id="dataset-1", tenant_id="tenant-1", collection_binding_id=None)
     monkeypatch.setattr(qdrant_module.Dataset, "gen_collection_name_by_id", lambda _id: "AUTO_COLLECTION")
     monkeypatch.setattr(qdrant_module, "current_app", SimpleNamespace(config=SimpleNamespace(root_path="/root")))
@@ -312,7 +318,7 @@ def test_qdrant_factory_paths(
     monkeypatch.setattr(qdrant_module.dify_config, "QDRANT_REPLICATION_FACTOR", 1)
 
     with patch.object(qdrant_module, "QdrantVector", return_value="vector") as vector_cls:
-        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        result = factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=session)
     assert result == "vector"
     assert vector_cls.call_args.kwargs["collection_name"] == "AUTO_COLLECTION"
     assert dataset.index_struct is not None
@@ -331,10 +337,53 @@ def test_qdrant_factory_paths(
     monkeypatch.setattr(qdrant_module.db, "session", sqlite3_session)
 
     with patch.object(qdrant_module, "QdrantVector", return_value="vector") as vector_cls:
-        factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=session)
     assert vector_cls.call_args.kwargs["collection_name"] == "BOUND_COLLECTION"
 
     sqlite3_session.delete(binding)
     sqlite3_session.commit()
     with pytest.raises(ValueError, match="Dataset Collection Bindings does not exist"):
-        factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=session)
+
+
+@pytest.mark.parametrize("vector_type", [None, "qdrant"])
+def test_vector_reads_uncommitted_annotation_binding(
+    qdrant_module, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, vector_type: str | None
+):
+    # Separate connections must not see each other's uncommitted rows. An
+    # in-memory SQLite engine would share a connection and hide this regression.
+    engine = create_engine(f"sqlite:///{tmp_path / 'bindings.db'}")
+    DatasetCollectionBinding.__table__.create(engine)
+    monkeypatch.setattr(Vector, "get_vector_factory", staticmethod(lambda _: qdrant_module.QdrantVectorFactory))
+    monkeypatch.setattr(qdrant_module, "current_app", SimpleNamespace(config=SimpleNamespace(root_path="/root")))
+    monkeypatch.setattr(qdrant_module.dify_config, "VECTOR_STORE", "qdrant")
+    monkeypatch.setattr(qdrant_module.dify_config, "VECTOR_STORE_WHITELIST_ENABLE", False)
+    monkeypatch.setattr(qdrant_module.dify_config, "QDRANT_URL", "http://localhost:6333")
+
+    try:
+        with Session(engine) as session, Session(engine) as other_session:
+            monkeypatch.setattr(qdrant_module.db, "session", other_session)
+            binding = DatasetCollectionBinding(
+                provider_name="provider",
+                model_name="model",
+                type=CollectionBindingType.ANNOTATION,
+                collection_name="ANNOTATION_COLLECTION",
+            )
+            session.add(binding)
+            session.flush()
+            stmt = select(DatasetCollectionBinding).where(DatasetCollectionBinding.id == binding.id)
+            assert other_session.scalar(stmt) is None
+            dataset = Dataset(id="app-1", tenant_id="tenant-1", collection_binding_id=binding.id)
+
+            vector = Vector(
+                dataset, attributes=["doc_id", "annotation_id", "app_id"], session=session, vector_type=vector_type
+            )
+
+            assert vector._vector_processor.collection_name == "ANNOTATION_COLLECTION"
+            assert dataset.index_struct_dict is not None
+            assert dataset.index_struct_dict["vector_store"]["class_prefix"] == "ANNOTATION_COLLECTION"
+            assert other_session.scalar(stmt) is None
+            session.rollback()
+            assert session.scalar(stmt) is None
+    finally:
+        engine.dispose()

api/providers/vdb/vdb-relyt/src/dify_vdb_relyt/relyt_vector.py

@@ -302,7 +302,9 @@ class RelytVector(BaseVector):
 
 class RelytVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> RelytVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> RelytVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-relyt/tests/unit_tests/test_relyt_vector.py

@@ -312,8 +312,8 @@ def test_relyt_factory_existing_and_generated_collection(relyt_module, monkeypat
     monkeypatch.setattr(relyt_module.dify_config, "RELYT_DATABASE", "relyt")
 
     with patch.object(relyt_module, "RelytVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-tablestore/src/dify_vdb_tablestore/tablestore_vector.py

@@ -6,6 +6,7 @@ from typing import Any, override
 
 import tablestore  # type: ignore
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 from tablestore import BatchGetRowRequest, TableInBatchGetRowItem
 
 from configs import dify_config
@@ -404,7 +405,9 @@ class TableStoreVector(BaseVector):
 
 class TableStoreVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TableStoreVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> TableStoreVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-tablestore/tests/unit_tests/test_tablestore_vector.py

@@ -305,8 +305,8 @@ def test_tablestore_factory_uses_existing_or_generated_collection(tablestore_mod
     monkeypatch.setattr(tablestore_module.dify_config, "TABLESTORE_NORMALIZE_FULLTEXT_BM25_SCORE", True)
 
     with patch.object(tablestore_module, "TableStoreVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-tencent/src/dify_vdb_tencent/tencent_vector.py

@@ -4,6 +4,7 @@ import math
 from typing import Any, TypedDict, override
 
 from pydantic import BaseModel
+from sqlalchemy.orm import Session
 from tcvdb_text.encoder import BM25Encoder  # type: ignore
 from tcvectordb import RPCVectorDBClient, VectorDBException  # type: ignore
 from tcvectordb.model import document, enum  # type: ignore
@@ -337,7 +338,9 @@ class TencentVector(BaseVector):
 
 class TencentVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TencentVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> TencentVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-tencent/tests/unit_tests/test_tencent_vector.py

@@ -316,8 +316,8 @@ def test_tencent_factory_existing_and_generated_collection(tencent_module, monke
     monkeypatch.setattr(tencent_module.dify_config, "TENCENT_VECTOR_DB_ENABLE_HYBRID_SEARCH", True)
 
     with patch.object(tencent_module, "TencentVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_on_qdrant_vector.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, override
 
 import httpx
 import qdrant_client
+from sqlalchemy.orm import Session
 
 logger = logging.getLogger(__name__)
 from flask import current_app
@@ -437,7 +438,9 @@ class TidbOnQdrantVector(BaseVector):
 
 class TidbOnQdrantVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TidbOnQdrantVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> TidbOnQdrantVector:
         logger.info("init_vector: tenant_id=%s, dataset_id=%s", dataset.tenant_id, dataset.id)
         stmt = select(TidbAuthBinding).where(TidbAuthBinding.tenant_id == dataset.tenant_id)
         tidb_auth_binding = db.session.scalars(stmt).one_or_none()

api/providers/vdb/vdb-tidb-vector/src/dify_vdb_tidb_vector/tidb_vector.py

@@ -356,7 +356,9 @@ class TiDBVector(BaseVector):
 
 class TiDBVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TiDBVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> TiDBVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-tidb-vector/tests/unit_tests/test_tidb_vector.py

@@ -548,8 +548,8 @@ def test_tidb_factory_uses_existing_or_generated_collection(tidb_module, monkeyp
     monkeypatch.setattr(tidb_module.dify_config, "APPLICATION_NAME", "dify-app")
 
     with patch.object(tidb_module, "TiDBVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-upstash/src/dify_vdb_upstash/upstash_vector.py

@@ -3,6 +3,7 @@ from typing import Any, override
 from uuid import uuid4
 
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 from upstash_vector import Index, Vector
 
 from configs import dify_config
@@ -136,7 +137,9 @@ class UpstashVector(BaseVector):
 
 class UpstashVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> UpstashVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> UpstashVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-upstash/tests/unit_tests/test_upstash_vector.py

@@ -176,8 +176,8 @@ def test_upstash_factory_uses_existing_or_generated_collection(upstash_module, m
     monkeypatch.setattr(upstash_module.dify_config, "UPSTASH_VECTOR_TOKEN", "token-123")
 
     with patch.object(upstash_module, "UpstashVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-vastbase/src/dify_vdb_vastbase/vastbase_vector.py

@@ -9,6 +9,7 @@ import psycopg2.errors
 import psycopg2.extras
 import psycopg2.pool
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 
 from configs import dify_config
 from core.rag.datasource.vdb.vector_base import BaseVector
@@ -288,7 +289,9 @@ class VastbaseVector(BaseVector):
 
 class VastbaseVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> VastbaseVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> VastbaseVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix

api/providers/vdb/vdb-vastbase/tests/unit_tests/test_vastbase_vector.py

@@ -382,8 +382,8 @@ def test_vastbase_factory_uses_existing_or_generated_collection(vastbase_module,
     monkeypatch.setattr(vastbase_module.dify_config, "VASTBASE_MAX_CONNECTION", 5)
 
     with patch.object(vastbase_module, "VastbaseVector", return_value="vector") as vector_cls:
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"

api/providers/vdb/vdb-vikingdb/src/dify_vdb_vikingdb/vikingdb_vector.py

@@ -2,6 +2,7 @@ import json
 from typing import Any, override
 
 from pydantic import BaseModel
+from sqlalchemy.orm import Session
 from volcengine.viking_db import (  # type: ignore
     Data,
     DistanceType,
@@ -220,7 +221,9 @@ class VikingDBVector(BaseVector):
 
 class VikingDBVectorFactory(AbstractVectorFactory):
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> VikingDBVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> VikingDBVector:
         if dataset.index_struct_dict:
             class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"]
             collection_name = class_prefix.lower()

api/providers/vdb/vdb-vikingdb/tests/unit_tests/test_vikingdb_vector.py

@@ -272,8 +272,8 @@ def test_vikingdb_factory_validates_config_and_builds_vector(vikingdb_module, mo
         monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_CONNECTION_TIMEOUT", 10)
         monkeypatch.setattr(vikingdb_module.dify_config, "VIKINGDB_SOCKET_TIMEOUT", 20)
 
-        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock())
-        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock())
+        result_1 = factory.init_vector(dataset_with_index, attributes=[], embeddings=MagicMock(), session=None)
+        result_2 = factory.init_vector(dataset_without_index, attributes=[], embeddings=MagicMock(), session=None)
 
     assert result_1 == "vector"
     assert result_2 == "vector"
@@ -306,4 +306,4 @@ def test_vikingdb_factory_raises_when_required_config_missing(
     monkeypatch.setattr(vikingdb_module.dify_config, field, None)
 
     with pytest.raises(ValueError, match=message):
-        factory.init_vector(dataset, attributes=[], embeddings=MagicMock())
+        factory.init_vector(dataset, attributes=[], embeddings=MagicMock(), session=None)

api/providers/vdb/vdb-weaviate/src/dify_vdb_weaviate/weaviate_vector.py

@@ -18,6 +18,7 @@ from urllib.parse import urlparse
 import weaviate
 import weaviate.classes.config as wc
 from pydantic import BaseModel, model_validator
+from sqlalchemy.orm import Session
 from weaviate.classes.data import DataObject
 from weaviate.classes.init import Auth
 from weaviate.classes.query import Filter, MetadataQuery
@@ -606,7 +607,9 @@ class WeaviateVectorFactory(AbstractVectorFactory):
     """Factory class for creating WeaviateVector instances."""
 
     @override
-    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> WeaviateVector:
+    def init_vector(
+        self, dataset: Dataset, attributes: list, embeddings: Embeddings, *, session: Session | None
+    ) -> WeaviateVector:
         """
         Initializes a WeaviateVector instance for the given dataset.
 

api/providers/vdb/vdb-weaviate/tests/unit_tests/test_weaviate_vector.py

@@ -893,7 +893,7 @@ class TestWeaviateVectorFactory(unittest.TestCase):
             patch("dify_vdb_weaviate.weaviate_vector.WeaviateVector", return_value="vector") as mock_vector,
         ):
             factory = weaviate_vector_module.WeaviateVectorFactory()
-            result = factory.init_vector(dataset, attributes, MagicMock())
+            result = factory.init_vector(dataset, attributes, MagicMock(), session=None)
 
         assert result == "vector"
         config = mock_vector.call_args.kwargs["config"]
@@ -922,7 +922,7 @@ class TestWeaviateVectorFactory(unittest.TestCase):
             patch("dify_vdb_weaviate.weaviate_vector.WeaviateVector", return_value="vector") as mock_vector,
         ):
             factory = weaviate_vector_module.WeaviateVectorFactory()
-            result = factory.init_vector(dataset, attributes, MagicMock())
+            result = factory.init_vector(dataset, attributes, MagicMock(), session=None)
 
         assert result == "vector"
         assert mock_vector.call_args.kwargs["collection_name"] == "GeneratedCollection_Node"

api/pyproject.toml

@@ -6,7 +6,7 @@ requires-python = "~=3.12.0"
 dependencies = [
     # Legacy: mature and widely deployed
     "bleach>=6.4.0,<7.0.0",
-    "boto3>=1.43.102,<2.0.0",
+    "boto3>=1.43.107,<2.0.0",
     "celery>=5.6.3,<6.0.0",
     "croniter>=6.2.2,<7.0.0",
     "dify-agent",
@@ -15,7 +15,7 @@ dependencies = [
     "gevent>=26.4.0,<26.5.0",
     "gevent-websocket==0.10.1",
     "gmpy2>=2.3.0,<3.0.0",
-    "google-api-python-client>=2.200.0,<3.0.0",
+    "google-api-python-client>=2.201.0,<3.0.0",
     "gunicorn>=26.2.0,<27.0.0",
     "psycogreen>=1.0.2,<2.0.0",
     "psycopg2-binary>=2.9.12,<3.0.0",
@@ -25,13 +25,13 @@ dependencies = [
     "sseclient-py>=1.8.0,<2.0.0",
     # Stable: production-proven, cap below the next major
     "aliyun-log-python-sdk==0.9.44",
-    "azure-identity>=1.25.3,<2.0.0",
+    "azure-identity>=1.26.0,<2.0.0",
     "flask-compress>=1.25,<2.0.0",
     "flask-login==0.6.3",
     "flask-migrate>=4.1.0,<5.0.0",
     "flask-orjson>=2.0.0,<3.0.0",
     "flask-restx>=1.3.2,<2.0.0",
-    "google-cloud-aiplatform>=2.2.0,<3.0.0",
+    "google-cloud-aiplatform>=2.3.0,<3.0.0",
     "httpx[socks]==0.28.1",
     "opentelemetry-distro==0.65b0",
     "opentelemetry-instrumentation-celery==0.65b0",
@@ -39,7 +39,7 @@ dependencies = [
     "opentelemetry-instrumentation-httpx==0.65b0",
     "opentelemetry-instrumentation-redis==0.65b0",
     "opentelemetry-instrumentation-sqlalchemy==0.65b0",
-    "opentelemetry-propagator-b3>=1.44.0,<2.0.0",
+    "opentelemetry-propagator-b3>=1.45.0,<2.0.0",
     "readabilipy==0.3.0",
     "resend>=2.27.0,<3.0.0",
     "zstandard==0.25.0",
@@ -201,11 +201,11 @@ kms = [
 # Required for storage clients
 ############################################################
 storage = [
-    "azure-storage-blob>=12.30.3,<13.0.0",
+    "azure-storage-blob>=12.31.0,<13.0.0",
     "bce-python-sdk==0.9.79",
     "cos-python-sdk-v5>=1.9.44,<2.0.0",
     "esdk-obs-python>=3.26.6,<4.0.0",
-    "google-cloud-storage>=3.15.0,<4.0.0",
+    "google-cloud-storage>=3.16.0,<4.0.0",
     "opendal==0.46.0",
     "oss2>=2.19.1,<3.0.0",
     "supabase>=2.31.0,<3.0.0",

api/repositories/api_based_extension_repository.py

@@ -0,0 +1,104 @@
+"""SQLAlchemy persistence adapter for Console API-based extensions."""
+
+from typing import override
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session, sessionmaker
+
+from models.api_based_extension import APIBasedExtension
+from services.api_based_extension_application_service import (
+    APIBasedExtensionInput,
+    APIBasedExtensionNotFoundError,
+    APIBasedExtensionRecord,
+    APIBasedExtensionStore,
+)
+
+
+class APIBasedExtensionRepository(APIBasedExtensionStore):
+    def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
+        self._session_factory = session_factory
+
+    @override
+    def list_extensions(self, workspace_id: str) -> tuple[APIBasedExtensionRecord, ...]:
+        with self._session_factory() as session:
+            return tuple(
+                self._to_record(extension)
+                for extension in session.scalars(
+                    select(APIBasedExtension)
+                    .where(APIBasedExtension.tenant_id == workspace_id)
+                    .order_by(APIBasedExtension.created_at.desc(), APIBasedExtension.id.desc())
+                ).all()
+            )
+
+    @override
+    def find_extension(self, workspace_id: str, extension_id: str) -> APIBasedExtensionRecord | None:
+        with self._session_factory() as session:
+            extension = self._find(session, workspace_id, extension_id)
+            return self._to_record(extension) if extension is not None else None
+
+    @override
+    def name_exists(self, workspace_id: str, name: str, *, exclude_id: str | None = None) -> bool:
+        stmt = select(APIBasedExtension.id).where(
+            APIBasedExtension.tenant_id == workspace_id,
+            APIBasedExtension.name == name,
+        )
+        if exclude_id is not None:
+            stmt = stmt.where(APIBasedExtension.id != exclude_id)
+        with self._session_factory() as session:
+            return session.scalar(stmt.limit(1)) is not None
+
+    @override
+    def create_extension(self, workspace_id: str, extension: APIBasedExtensionInput) -> APIBasedExtensionRecord:
+        with self._session_factory.begin() as session:
+            model = APIBasedExtension(
+                tenant_id=workspace_id,
+                name=extension.name,
+                api_endpoint=extension.api_endpoint,
+                api_key=extension.api_key,
+            )
+            session.add(model)
+            session.flush()
+            session.refresh(model)
+            return self._to_record(model)
+
+    @override
+    def update_extension(
+        self, workspace_id: str, extension_id: str, extension: APIBasedExtensionInput
+    ) -> APIBasedExtensionRecord:
+        with self._session_factory.begin() as session:
+            model = self._require(session, workspace_id, extension_id)
+            model.name = extension.name
+            model.api_endpoint = extension.api_endpoint
+            model.api_key = extension.api_key
+            session.flush()
+            return self._to_record(model)
+
+    @override
+    def delete_extension(self, workspace_id: str, extension_id: str) -> None:
+        with self._session_factory.begin() as session:
+            session.delete(self._require(session, workspace_id, extension_id))
+
+    @staticmethod
+    def _find(session: Session, workspace_id: str, extension_id: str) -> APIBasedExtension | None:
+        return session.scalar(
+            select(APIBasedExtension)
+            .where(APIBasedExtension.tenant_id == workspace_id, APIBasedExtension.id == extension_id)
+            .limit(1)
+        )
+
+    @classmethod
+    def _require(cls, session: Session, workspace_id: str, extension_id: str) -> APIBasedExtension:
+        extension = cls._find(session, workspace_id, extension_id)
+        if extension is None:
+            raise APIBasedExtensionNotFoundError
+        return extension
+
+    @staticmethod
+    def _to_record(extension: APIBasedExtension) -> APIBasedExtensionRecord:
+        return APIBasedExtensionRecord(
+            id=extension.id,
+            name=extension.name,
+            api_endpoint=extension.api_endpoint,
+            api_key=extension.api_key,
+            created_at=extension.created_at,
+        )

api/repositories/api_workflow_run_repository.py

@@ -401,6 +401,20 @@ class APIWorkflowRunRepository(Protocol):
         """
         ...
 
+    def delete_runs_with_related_in_session(
+        self,
+        session: Session,
+        runs: Sequence[WorkflowRun],
+        delete_node_executions: Callable[[Session, Sequence[WorkflowRun]], tuple[int, int]] | None = None,
+        delete_trigger_logs: Callable[[Session, Sequence[str]], int] | None = None,
+    ) -> RunsWithRelatedCountsDict:
+        """
+        Delete workflow runs and related records using the caller-owned transaction.
+
+        The implementation must not commit or roll back the supplied session.
+        """
+        ...
+
     def get_pause_records_by_run_id(
         self,
         session: Session,

api/repositories/sqlalchemy_api_workflow_run_repository.py

@@ -653,57 +653,39 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
         runs: Sequence[WorkflowRun],
         delete_node_executions: Callable[[Session, Sequence[WorkflowRun]], tuple[int, int]] | None = None,
         delete_trigger_logs: Callable[[Session, Sequence[str]], int] | None = None,
+    ) -> RunsWithRelatedCountsDict:
+        with self._session_maker.begin() as session:
+            return self.delete_runs_with_related_in_session(
+                session,
+                runs,
+                delete_node_executions=delete_node_executions,
+                delete_trigger_logs=delete_trigger_logs,
+            )
+
+    @override
+    def delete_runs_with_related_in_session(
+        self,
+        session: Session,
+        runs: Sequence[WorkflowRun],
+        delete_node_executions: Callable[[Session, Sequence[WorkflowRun]], tuple[int, int]] | None = None,
+        delete_trigger_logs: Callable[[Session, Sequence[str]], int] | None = None,
     ) -> RunsWithRelatedCountsDict:
         if not runs:
-            return {
-                "runs": 0,
-                "node_executions": 0,
-                "offloads": 0,
-                "app_logs": 0,
-                "trigger_logs": 0,
-                "pauses": 0,
-                "pause_reasons": 0,
-            }
+            return self._empty_runs_with_related_counts()
 
-        with self._session_maker() as session:
-            run_ids = [run.id for run in runs]
-            if delete_node_executions:
-                node_executions_deleted, offloads_deleted = delete_node_executions(session, runs)
-            else:
-                node_executions_deleted, offloads_deleted = 0, 0
+        runs = list(runs)
 
-            app_logs_result = session.execute(delete(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id.in_(run_ids)))
-            app_logs_deleted = cast(CursorResult, app_logs_result).rowcount or 0
+        def delete_node_executions_by_ids(active_session: Session, _run_ids: Sequence[str]) -> tuple[int, int]:
+            if delete_node_executions is None:
+                return 0, 0
+            return delete_node_executions(active_session, runs)
 
-            pause_stmt = select(WorkflowPause.id).where(WorkflowPause.workflow_run_id.in_(run_ids))
-            pause_ids = session.scalars(pause_stmt).all()
-            pause_reasons_deleted = 0
-            pauses_deleted = 0
-
-            if pause_ids:
-                pause_reasons_result = session.execute(
-                    delete(WorkflowPauseReason).where(WorkflowPauseReason.pause_id.in_(pause_ids))
-                )
-                pause_reasons_deleted = cast(CursorResult, pause_reasons_result).rowcount or 0
-                pauses_result = session.execute(delete(WorkflowPause).where(WorkflowPause.id.in_(pause_ids)))
-                pauses_deleted = cast(CursorResult, pauses_result).rowcount or 0
-
-            trigger_logs_deleted = delete_trigger_logs(session, run_ids) if delete_trigger_logs else 0
-
-            runs_result = session.execute(delete(WorkflowRun).where(WorkflowRun.id.in_(run_ids)))
-            runs_deleted = cast(CursorResult, runs_result).rowcount or 0
-
-            session.commit()
-
-            return {
-                "runs": runs_deleted,
-                "node_executions": node_executions_deleted,
-                "offloads": offloads_deleted,
-                "app_logs": app_logs_deleted,
-                "trigger_logs": trigger_logs_deleted,
-                "pauses": pauses_deleted,
-                "pause_reasons": pause_reasons_deleted,
-            }
+        return self._delete_runs_with_related_by_ids_in_session(
+            session,
+            [run.id for run in runs],
+            delete_node_executions=delete_node_executions_by_ids if delete_node_executions else None,
+            delete_trigger_logs=delete_trigger_logs,
+        )
 
     @override
     def delete_runs_with_related_by_ids(
@@ -711,49 +693,61 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
         run_ids: Sequence[str],
         delete_node_executions: Callable[[Session, Sequence[str]], tuple[int, int]] | None = None,
         delete_trigger_logs: Callable[[Session, Sequence[str]], int] | None = None,
+    ) -> RunsWithRelatedCountsDict:
+        with self._session_maker.begin() as session:
+            return self._delete_runs_with_related_by_ids_in_session(
+                session,
+                run_ids,
+                delete_node_executions=delete_node_executions,
+                delete_trigger_logs=delete_trigger_logs,
+            )
+
+    def _delete_runs_with_related_by_ids_in_session(
+        self,
+        session: Session,
+        run_ids: Sequence[str],
+        delete_node_executions: Callable[[Session, Sequence[str]], tuple[int, int]] | None = None,
+        delete_trigger_logs: Callable[[Session, Sequence[str]], int] | None = None,
     ) -> RunsWithRelatedCountsDict:
         if not run_ids:
             return self._empty_runs_with_related_counts()
 
         run_ids = list(run_ids)
-        with self._session_maker() as session:
-            if delete_node_executions:
-                node_executions_deleted, offloads_deleted = delete_node_executions(session, run_ids)
-            else:
-                node_executions_deleted, offloads_deleted = 0, 0
+        if delete_node_executions:
+            node_executions_deleted, offloads_deleted = delete_node_executions(session, run_ids)
+        else:
+            node_executions_deleted, offloads_deleted = 0, 0
 
-            app_logs_result = session.execute(delete(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id.in_(run_ids)))
-            app_logs_deleted = cast(CursorResult, app_logs_result).rowcount or 0
+        app_logs_result = session.execute(delete(WorkflowAppLog).where(WorkflowAppLog.workflow_run_id.in_(run_ids)))
+        app_logs_deleted = cast(CursorResult, app_logs_result).rowcount or 0
 
-            pause_stmt = select(WorkflowPause.id).where(WorkflowPause.workflow_run_id.in_(run_ids))
-            pause_ids = session.scalars(pause_stmt).all()
-            pause_reasons_deleted = 0
-            pauses_deleted = 0
+        pause_stmt = select(WorkflowPause.id).where(WorkflowPause.workflow_run_id.in_(run_ids))
+        pause_ids = session.scalars(pause_stmt).all()
+        pause_reasons_deleted = 0
+        pauses_deleted = 0
 
-            if pause_ids:
-                pause_reasons_result = session.execute(
-                    delete(WorkflowPauseReason).where(WorkflowPauseReason.pause_id.in_(pause_ids))
-                )
-                pause_reasons_deleted = cast(CursorResult, pause_reasons_result).rowcount or 0
-                pauses_result = session.execute(delete(WorkflowPause).where(WorkflowPause.id.in_(pause_ids)))
-                pauses_deleted = cast(CursorResult, pauses_result).rowcount or 0
+        if pause_ids:
+            pause_reasons_result = session.execute(
+                delete(WorkflowPauseReason).where(WorkflowPauseReason.pause_id.in_(pause_ids))
+            )
+            pause_reasons_deleted = cast(CursorResult, pause_reasons_result).rowcount or 0
+            pauses_result = session.execute(delete(WorkflowPause).where(WorkflowPause.id.in_(pause_ids)))
+            pauses_deleted = cast(CursorResult, pauses_result).rowcount or 0
 
-            trigger_logs_deleted = delete_trigger_logs(session, run_ids) if delete_trigger_logs else 0
+        trigger_logs_deleted = delete_trigger_logs(session, run_ids) if delete_trigger_logs else 0
 
-            runs_result = session.execute(delete(WorkflowRun).where(WorkflowRun.id.in_(run_ids)))
-            runs_deleted = cast(CursorResult, runs_result).rowcount or 0
+        runs_result = session.execute(delete(WorkflowRun).where(WorkflowRun.id.in_(run_ids)))
+        runs_deleted = cast(CursorResult, runs_result).rowcount or 0
 
-            session.commit()
-
-            return {
-                "runs": runs_deleted,
-                "node_executions": node_executions_deleted,
-                "offloads": offloads_deleted,
-                "app_logs": app_logs_deleted,
-                "trigger_logs": trigger_logs_deleted,
-                "pauses": pauses_deleted,
-                "pause_reasons": pause_reasons_deleted,
-            }
+        return {
+            "runs": runs_deleted,
+            "node_executions": node_executions_deleted,
+            "offloads": offloads_deleted,
+            "app_logs": app_logs_deleted,
+            "trigger_logs": trigger_logs_deleted,
+            "pauses": pauses_deleted,
+            "pause_reasons": pause_reasons_deleted,
+        }
 
     @override
     def get_app_logs_by_run_id(

api/schedule/clean_workflow_runlogs_precise.py

@@ -155,7 +155,8 @@ def _delete_batch(
                 trigger_repo = SQLAlchemyWorkflowTriggerLogRepository(active_session)
                 return trigger_repo.delete_by_run_ids(run_ids)
 
-            workflow_run_repo.delete_runs_with_related(
+            workflow_run_repo.delete_runs_with_related_in_session(
+                session,
                 workflow_runs,
                 delete_node_executions=_delete_node_executions,
                 delete_trigger_logs=_delete_trigger_logs,

api/services/api_based_extension_adapters.py

@@ -0,0 +1,40 @@
+"""Runtime adapters for the API-based extension application service."""
+
+from typing import override
+
+from core.extension.api_based_extension_requestor import APIBasedExtensionRequestor
+from core.helper.encrypter import decrypt_token, encrypt_token
+from models.api_based_extension import APIBasedExtensionPoint
+from services.api_based_extension_application_service import (
+    APIBasedExtensionConnectionError,
+    APIBasedExtensionEndpointProbe,
+    APIBasedExtensionSecretCipher,
+)
+
+
+class WorkspaceTokenCipher(APIBasedExtensionSecretCipher):
+    """Encrypts API keys with the workspace key provider used by the rest of the platform."""
+
+    @override
+    def encrypt(self, workspace_id: str, value: str) -> str:
+        return encrypt_token(workspace_id, value)
+
+    @override
+    def decrypt(self, workspace_id: str, value: str) -> str:
+        return decrypt_token(workspace_id, value)
+
+
+class APIBasedExtensionPingProbe(APIBasedExtensionEndpointProbe):
+    """Sends the ``ping`` extension point through the SSRF-guarded extension requestor."""
+
+    @override
+    def ping(self, api_endpoint: str, api_key: str) -> None:
+        try:
+            response = APIBasedExtensionRequestor(api_endpoint, api_key).request(
+                point=APIBasedExtensionPoint.PING, params={}
+            )
+        except Exception as error:
+            raise APIBasedExtensionConnectionError(str(error)) from error
+        # The requestor only promises JSON; an endpoint may answer with a list or null instead of an object.
+        if not isinstance(response, dict) or response.get("result") != "pong":
+            raise APIBasedExtensionConnectionError(str(response))

api/services/api_based_extension_application_service.py

@@ -0,0 +1,178 @@
+"""Application boundary for Console API-based extension management."""
+
+from collections.abc import Sequence
+from datetime import datetime
+from typing import NamedTuple, Protocol
+
+from machinery.context import RequestContext
+
+_MIN_API_KEY_LENGTH = 5
+
+
+class APIBasedExtensionRecord(NamedTuple):
+    """One API-based extension of a workspace.
+
+    The store persists and returns ``api_key`` as ciphertext; the application service
+    hands callers the plaintext key.
+    """
+
+    id: str
+    name: str
+    api_endpoint: str
+    api_key: str
+    created_at: datetime
+
+
+class APIBasedExtensionInput(NamedTuple):
+    """Field values for creating an extension; ``api_key`` is plaintext."""
+
+    name: str
+    api_endpoint: str
+    api_key: str
+
+
+class APIBasedExtensionUpdate(NamedTuple):
+    """Field edits for an extension; a ``None`` ``api_key`` keeps the stored key."""
+
+    name: str
+    api_endpoint: str
+    api_key: str | None
+
+
+class APIBasedExtensionStore(Protocol):
+    """Persistence for API-based extensions, always scoped by workspace."""
+
+    def list_extensions(self, workspace_id: str) -> Sequence[APIBasedExtensionRecord]: ...
+
+    def find_extension(self, workspace_id: str, extension_id: str) -> APIBasedExtensionRecord | None: ...
+
+    def name_exists(self, workspace_id: str, name: str, *, exclude_id: str | None = None) -> bool: ...
+
+    def create_extension(self, workspace_id: str, extension: APIBasedExtensionInput) -> APIBasedExtensionRecord:
+        """Insert a new extension; ``extension.api_key`` is already ciphertext."""
+        ...
+
+    def update_extension(
+        self, workspace_id: str, extension_id: str, extension: APIBasedExtensionInput
+    ) -> APIBasedExtensionRecord:
+        """Overwrite all editable fields; raise ``APIBasedExtensionNotFoundError`` if absent from this scope."""
+        ...
+
+    def delete_extension(self, workspace_id: str, extension_id: str) -> None:
+        """Delete the extension; raise ``APIBasedExtensionNotFoundError`` if absent from this scope."""
+        ...
+
+
+class APIBasedExtensionSecretCipher(Protocol):
+    """Workspace-scoped encryption for stored API keys."""
+
+    def encrypt(self, workspace_id: str, value: str) -> str: ...
+
+    def decrypt(self, workspace_id: str, value: str) -> str: ...
+
+
+class APIBasedExtensionEndpointProbe(Protocol):
+    """Verifies that an extension endpoint answers the ``ping`` extension point."""
+
+    def ping(self, api_endpoint: str, api_key: str) -> None:
+        """Return normally when the endpoint answers ``pong``; raise ``APIBasedExtensionConnectionError`` otherwise."""
+        ...
+
+
+class APIBasedExtensionError(Exception):
+    """Base class for framework-neutral API-based extension failures."""
+
+
+class APIBasedExtensionNotFoundError(APIBasedExtensionError):
+    def __init__(self) -> None:
+        super().__init__("API based extension is not found")
+
+
+class APIBasedExtensionNameConflictError(APIBasedExtensionError):
+    def __init__(self) -> None:
+        super().__init__("name must be unique, it is already existed")
+
+
+class APIBasedExtensionInvalidError(APIBasedExtensionError):
+    """A field value does not satisfy the extension contract."""
+
+
+class APIBasedExtensionConnectionError(APIBasedExtensionError):
+    def __init__(self, reason: str) -> None:
+        super().__init__(f"connection error: {reason}")
+
+
+class APIBasedExtensionApplicationService:
+    """Owns the extension rules: non-empty fields, a workspace-unique name, a key of at least five
+    characters, an endpoint that answers ``ping``, and keys stored only as ciphertext."""
+
+    def __init__(
+        self,
+        *,
+        extensions: APIBasedExtensionStore,
+        secrets: APIBasedExtensionSecretCipher,
+        probe: APIBasedExtensionEndpointProbe,
+    ) -> None:
+        self._extensions = extensions
+        self._secrets = secrets
+        self._probe = probe
+
+    def list_extensions(self, context: RequestContext) -> tuple[APIBasedExtensionRecord, ...]:
+        workspace_id = context.active_workspace_id
+        return tuple(
+            self._with_plaintext_key(workspace_id, record) for record in self._extensions.list_extensions(workspace_id)
+        )
+
+    def get_extension(self, context: RequestContext, extension_id: str) -> APIBasedExtensionRecord:
+        workspace_id = context.active_workspace_id
+        return self._with_plaintext_key(workspace_id, self._require_extension(workspace_id, extension_id))
+
+    def create_extension(self, context: RequestContext, extension: APIBasedExtensionInput) -> APIBasedExtensionRecord:
+        workspace_id = context.active_workspace_id
+        self._validate(workspace_id, extension, exclude_id=None)
+        stored = self._extensions.create_extension(
+            workspace_id,
+            extension._replace(api_key=self._secrets.encrypt(workspace_id, extension.api_key)),
+        )
+        return stored._replace(api_key=extension.api_key)
+
+    def update_extension(
+        self, context: RequestContext, extension_id: str, update: APIBasedExtensionUpdate
+    ) -> APIBasedExtensionRecord:
+        workspace_id = context.active_workspace_id
+        current = self._require_extension(workspace_id, extension_id)
+        api_key = update.api_key if update.api_key is not None else self._secrets.decrypt(workspace_id, current.api_key)
+        extension = APIBasedExtensionInput(name=update.name, api_endpoint=update.api_endpoint, api_key=api_key)
+        self._validate(workspace_id, extension, exclude_id=extension_id)
+        stored = self._extensions.update_extension(
+            workspace_id,
+            extension_id,
+            extension._replace(api_key=self._secrets.encrypt(workspace_id, api_key)),
+        )
+        return stored._replace(api_key=api_key)
+
+    def delete_extension(self, context: RequestContext, extension_id: str) -> None:
+        # The store raises APIBasedExtensionNotFoundError itself, so no separate existence read is needed.
+        self._extensions.delete_extension(context.active_workspace_id, extension_id)
+
+    def _require_extension(self, workspace_id: str, extension_id: str) -> APIBasedExtensionRecord:
+        record = self._extensions.find_extension(workspace_id, extension_id)
+        if record is None:
+            raise APIBasedExtensionNotFoundError
+        return record
+
+    def _with_plaintext_key(self, workspace_id: str, record: APIBasedExtensionRecord) -> APIBasedExtensionRecord:
+        return record._replace(api_key=self._secrets.decrypt(workspace_id, record.api_key))
+
+    def _validate(self, workspace_id: str, extension: APIBasedExtensionInput, *, exclude_id: str | None) -> None:
+        if not extension.name:
+            raise APIBasedExtensionInvalidError("name must not be empty")
+        if self._extensions.name_exists(workspace_id, extension.name, exclude_id=exclude_id):
+            raise APIBasedExtensionNameConflictError
+        if not extension.api_endpoint:
+            raise APIBasedExtensionInvalidError("api_endpoint must not be empty")
+        if not extension.api_key:
+            raise APIBasedExtensionInvalidError("api_key must not be empty")
+        if len(extension.api_key) < _MIN_API_KEY_LENGTH:
+            raise APIBasedExtensionInvalidError(f"api_key must be at least {_MIN_API_KEY_LENGTH} characters")
+        self._probe.ping(extension.api_endpoint, extension.api_key)

api/services/api_based_extension_service.py

@@ -1,111 +0,0 @@
-from sqlalchemy import select
-from sqlalchemy.orm import Session
-
-from core.extension.api_based_extension_requestor import APIBasedExtensionRequestor
-from core.helper.encrypter import decrypt_token, encrypt_token
-from models.api_based_extension import APIBasedExtension, APIBasedExtensionPoint
-
-
-class APIBasedExtensionService:
-    @staticmethod
-    def get_all_by_tenant_id(tenant_id: str, *, session: Session) -> list[APIBasedExtension]:
-        extension_list = list(
-            session.scalars(
-                select(APIBasedExtension)
-                .where(APIBasedExtension.tenant_id == tenant_id)
-                .order_by(APIBasedExtension.created_at.desc())
-            ).all()
-        )
-
-        for extension in extension_list:
-            extension.api_key = decrypt_token(extension.tenant_id, extension.api_key)
-
-        return extension_list
-
-    @classmethod
-    def save(cls, extension_data: APIBasedExtension, *, session: Session) -> APIBasedExtension:
-        cls._validation(session, extension_data)
-
-        extension_data.api_key = encrypt_token(extension_data.tenant_id, extension_data.api_key)
-
-        session.add(extension_data)
-        session.commit()
-        return extension_data
-
-    @staticmethod
-    def delete(extension_data: APIBasedExtension, *, session: Session):
-        session.delete(extension_data)
-        session.commit()
-
-    @staticmethod
-    def get_with_tenant_id(tenant_id: str, api_based_extension_id: str, *, session: Session) -> APIBasedExtension:
-        extension = session.scalar(
-            select(APIBasedExtension)
-            .where(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id)
-            .limit(1)
-        )
-
-        if not extension:
-            raise ValueError("API based extension is not found")
-
-        extension.api_key = decrypt_token(extension.tenant_id, extension.api_key)
-
-        return extension
-
-    @classmethod
-    def _validation(cls, session: Session, extension_data: APIBasedExtension):
-        # name
-        if not extension_data.name:
-            raise ValueError("name must not be empty")
-
-        if not extension_data.id:
-            # case one: check new data, name must be unique
-            is_name_existed = session.scalar(
-                select(APIBasedExtension)
-                .where(
-                    APIBasedExtension.tenant_id == extension_data.tenant_id,
-                    APIBasedExtension.name == extension_data.name,
-                )
-                .limit(1)
-            )
-
-            if is_name_existed:
-                raise ValueError("name must be unique, it is already existed")
-        else:
-            # case two: check existing data, name must be unique
-            is_name_existed = session.scalar(
-                select(APIBasedExtension)
-                .where(
-                    APIBasedExtension.tenant_id == extension_data.tenant_id,
-                    APIBasedExtension.name == extension_data.name,
-                    APIBasedExtension.id != extension_data.id,
-                )
-                .limit(1)
-            )
-
-            if is_name_existed:
-                raise ValueError("name must be unique, it is already existed")
-
-        # api_endpoint
-        if not extension_data.api_endpoint:
-            raise ValueError("api_endpoint must not be empty")
-
-        # api_key
-        if not extension_data.api_key:
-            raise ValueError("api_key must not be empty")
-
-        if len(extension_data.api_key) < 5:
-            raise ValueError("api_key must be at least 5 characters")
-
-        # check endpoint
-        cls._ping_connection(extension_data)
-
-    @staticmethod
-    def _ping_connection(extension_data: APIBasedExtension):
-        try:
-            client = APIBasedExtensionRequestor(extension_data.api_endpoint, extension_data.api_key)
-            resp = client.request(point=APIBasedExtensionPoint.PING, params={})
-            if resp.get("result") != "pong":
-                raise ValueError(resp)
-        except Exception as e:
-            raise ValueError(f"connection error: {e}")

api/services/workflow_event_snapshot_service.py

@@ -43,6 +43,7 @@ from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus
 from graphon.runtime import GraphRuntimeState
 from graphon.runtime.graph_runtime_state_protocol import ReadOnlyVariablePool
 from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter
+from libs.broadcast_channel.exc import SubscriptionClosedError
 from libs.datetime_utils import to_utc_timestamp
 from models.human_input import HumanInputForm
 from models.model import AppMode, Message
@@ -669,6 +670,8 @@ def _start_buffering(subscription) -> BufferState:
                     except queue.Full:
                         continue
                     logger.warning("Dropped buffered workflow event, total_dropped=%s", dropped_count)
+        except SubscriptionClosedError:
+            pass
         except Exception:
             logger.exception("Failed while buffering workflow events")
         finally:

api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py

@@ -9,6 +9,7 @@ from flask import Flask
 from sqlalchemy.orm import Session, sessionmaker
 
 from controllers.console.datasets.data_source import DataSourceNotionListApi, DataSourceNotionListQuery
+from extensions.application_services.data_sources import build_data_source_credentials
 from machinery.context import RequestContext
 from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
 from models.dataset import Dataset, Document
@@ -16,7 +17,6 @@ from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
 from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository
 from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository
 from repositories.workspace.workspace_repository import WorkspaceRepository
-from services.data_source.credential_gateway import ActorDatasourceCredentialResolver
 from services.data_source.notion_import_adapters import PluginNotionSourceGateway
 from services.data_source.notion_import_application_service import NotionImportApplicationService
 from services.knowledge.dataset_access import DatasetAccessService
@@ -89,8 +89,7 @@ def test_notion_page_is_marked_bound_from_persisted_document(
 
     sessions = sessionmaker(bind=db_session_with_containers.get_bind(), expire_on_commit=False)
     datasets = SQLAlchemyDatasetRepository(session_factory=sessions)
-    credentials = MagicMock(spec=ActorDatasourceCredentialResolver)
-    credentials.resolve.return_value = {"token": "token"}
+    credentials = build_data_source_credentials(database_client=sessions).actor
     notion_imports = NotionImportApplicationService(
         dataset_access=DatasetAccessService(
             datasets=datasets,
@@ -106,6 +105,7 @@ def test_notion_page_is_marked_bound_from_persisted_document(
     )
 
     with (
+        patch.object(credentials, "resolve", return_value={"token": "token"}) as resolve_credentials,
         flask_app_with_containers.test_request_context(f"/?credential_id=c1&dataset_id={dataset_id}"),
         patch(
             "controllers.console.datasets.data_source.application_services",
@@ -120,7 +120,7 @@ def test_notion_page_is_marked_bound_from_persisted_document(
 
     assert status == 200
     assert response["notion_info"][0]["pages"][0]["is_bound"] is True
-    credentials.resolve.assert_called_once_with(
+    resolve_credentials.assert_called_once_with(
         workspace_id=tenant_id,
         actor_id=account.id,
         credential_id="c1",

api/tests/test_containers_integration_tests/controllers/console/test_api_based_extension.py

@@ -10,9 +10,9 @@ from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from constants import HIDDEN_VALUE
+from core.helper.encrypter import encrypt_token
 from libs.rsa import generate_key_pair
 from models.api_based_extension import APIBasedExtension
-from services.api_based_extension_service import APIBasedExtensionService
 from tests.test_containers_integration_tests.controllers.console.helpers import (
     authenticate_console_client,
     create_console_account_and_tenant,
@@ -41,7 +41,7 @@ def api_extension_client(
 
 @pytest.fixture(autouse=True)
 def mock_api_based_extension_ping():
-    with patch("services.api_based_extension_service.APIBasedExtensionRequestor") as requestor:
+    with patch("services.api_based_extension_adapters.APIBasedExtensionRequestor") as requestor:
         requestor.return_value.request.return_value = {"result": "pong"}
         yield requestor
 
@@ -96,15 +96,16 @@ def test_list_scopes_api_based_extensions_to_authenticated_tenant(
     )
     assert account_create_response.status_code == 201
 
-    APIBasedExtensionService.save(
+    # Seed another workspace's extension directly; the list endpoint must not expose it.
+    db_session_with_containers.add(
         APIBasedExtension(
             tenant_id=foreign_tenant_id,
             name="Foreign API",
             api_endpoint="https://foreign.example.com/hook",
-            api_key="foreign-secret-12345",
-        ),
-        session=db_session_with_containers,
+            api_key=encrypt_token(foreign_tenant_id, "foreign-secret-12345"),
+        )
     )
+    db_session_with_containers.commit()
 
     response = test_client_with_containers.get(
         "/console/api/api-based-extension",

api/tests/test_containers_integration_tests/core/workflow/test_human_input_resume_node_execution.py

@@ -1,7 +1,5 @@
 import time
 import uuid
-from datetime import timedelta
-from unittest.mock import MagicMock
 
 import pytest
 from sqlalchemy import delete, select
@@ -10,7 +8,7 @@ from sqlalchemy.orm import Session
 from core.app.app_config.entities import WorkflowUIBasedAppConfig
 from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity
 from core.app.workflow.layers import PersistenceWorkflowInfo, WorkflowPersistenceLayer
-from core.repositories.human_input_repository import HumanInputFormEntity, HumanInputFormRepository
+from core.repositories.human_input_repository import HumanInputFormRepository, HumanInputFormRepositoryImpl
 from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository
 from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository
 from core.workflow.nodes.human_input.callback import (
@@ -33,40 +31,12 @@ from libs.datetime_utils import naive_utc_now
 from models import Account
 from models.account import AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus
 from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom
+from models.human_input import HumanInputDelivery, HumanInputForm, HumanInputFormRecipient
 from models.model import App, AppMode, IconType
 from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowRun
 from tests.workflow_test_utils import build_test_graph_init_params
 
 
-def _mock_form_repository_without_submission() -> HumanInputFormRepository:
-    repo = MagicMock(spec=HumanInputFormRepository)
-    form_entity = MagicMock(spec=HumanInputFormEntity)
-    form_entity.id = "test-form-id"
-    form_entity.submission_token = "test-form-token"
-    form_entity.recipients = []
-    form_entity.rendered_content = "rendered"
-    form_entity.submitted = False
-    repo.create_form.return_value = form_entity
-    repo.get_form.return_value = None
-    return repo
-
-
-def _mock_form_repository_with_submission(action_id: str) -> HumanInputFormRepository:
-    repo = MagicMock(spec=HumanInputFormRepository)
-    form_entity = MagicMock(spec=HumanInputFormEntity)
-    form_entity.id = "test-form-id"
-    form_entity.submission_token = "test-form-token"
-    form_entity.recipients = []
-    form_entity.rendered_content = "rendered"
-    form_entity.submitted = True
-    form_entity.selected_action_id = action_id
-    form_entity.submitted_data = {}
-    form_entity.status = HumanInputFormStatus.WAITING
-    form_entity.expiration_time = naive_utc_now() + timedelta(hours=1)
-    repo.get_form.return_value = form_entity
-    return repo
-
-
 def _build_runtime_state(workflow_execution_id: str, app_id: str, workflow_id: str, user_id: str) -> GraphRuntimeState:
     variable_pool = VariablePool.from_bootstrap(
         system_variables=build_system_variables(
@@ -248,6 +218,16 @@ class TestHumanInputResumeNodeExecutionIntegration:
 
         yield
 
+        form_ids = select(HumanInputForm.id).where(
+            HumanInputForm.tenant_id == self.tenant.id, HumanInputForm.app_id == self.app.id
+        )
+        self.session.execute(delete(HumanInputFormRecipient).where(HumanInputFormRecipient.form_id.in_(form_ids)))
+        self.session.execute(delete(HumanInputDelivery).where(HumanInputDelivery.form_id.in_(form_ids)))
+        self.session.execute(
+            delete(HumanInputForm).where(
+                HumanInputForm.tenant_id == self.tenant.id, HumanInputForm.app_id == self.app.id
+            )
+        )
         self.session.execute(delete(WorkflowNodeExecutionModel))
         self.session.execute(delete(WorkflowRun))
         self.session.execute(delete(Workflow).where(Workflow.id == self.workflow.id))
@@ -310,27 +290,42 @@ class TestHumanInputResumeNodeExecutionIntegration:
             workflow_id=self.workflow.id,
             user_id=self.account.id,
         )
-        pause_repo = _mock_form_repository_without_submission()
+        form_repository = HumanInputFormRepositoryImpl(
+            tenant_id=self.tenant.id, app_id=self.app.id, workflow_execution_id=execution_id
+        )
         paused_graph = _build_graph(
             runtime_state,
             self.tenant.id,
             self.app.id,
             self.workflow.id,
             self.account.id,
-            pause_repo,
+            form_repository,
         )
         self._run_graph(paused_graph, runtime_state, execution_id)
 
         snapshot = runtime_state.dumps()
         resumed_state = GraphRuntimeState.from_snapshot(snapshot)
-        resume_repo = _mock_form_repository_with_submission(action_id="continue")
+        form = self.session.scalar(
+            select(HumanInputForm).where(
+                HumanInputForm.tenant_id == self.tenant.id,
+                HumanInputForm.workflow_run_id == execution_id,
+                HumanInputForm.node_id == "human",
+            )
+        )
+        assert form is not None
+        assert form.status == HumanInputFormStatus.WAITING
+        form.selected_action_id = "continue"
+        form.submitted_data = "{}"
+        form.submitted_at = naive_utc_now()
+        form.status = HumanInputFormStatus.SUBMITTED
+        self.session.commit()
         resumed_graph = _build_graph(
             resumed_state,
             self.tenant.id,
             self.app.id,
             self.workflow.id,
             self.account.id,
-            resume_repo,
+            form_repository,
         )
         self._run_graph(resumed_graph, resumed_state, execution_id)
 

api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_sharded_channel.py

@@ -237,6 +237,7 @@ class TestShardedRedisBroadcastChannelIntegration:
         def consumer_thread() -> set[bytes]:
             received_msgs: set[bytes] = set()
             with subscription:
+                subscription.receive(timeout=0.1)
                 consumer_ready.set()
                 while True:
                     try:
@@ -260,7 +261,10 @@ class TestShardedRedisBroadcastChannelIntegration:
             for future in as_completed(producer_futures, timeout=30.0):
                 sent_msgs.update(future.result())
 
-            consumer_received_msgs = consumer_future.result(timeout=60.0)
+            try:
+                consumer_received_msgs = consumer_future.result(timeout=60.0)
+            finally:
+                subscription.close()
 
         assert sent_msgs == consumer_received_msgs
 

api/tests/test_containers_integration_tests/pyrefly.toml

@@ -52,7 +52,6 @@ project-excludes = [
   "services/test_advanced_prompt_template_service.py",
   "services/test_agent_service.py",
   "services/test_annotation_service.py",
-  "services/test_api_based_extension_service.py",
   "services/test_app_dsl_service.py",
   "services/test_app_generate_service.py",
   "services/test_app_service.py",

api/tests/test_containers_integration_tests/repositories/test_recommended_app_catalog_repository.py

@@ -1,6 +1,9 @@
-from unittest.mock import MagicMock, patch
+from collections.abc import Iterator
+from unittest.mock import patch
 from uuid import uuid4
 
+import pytest
+from redis import Redis
 from sqlalchemy.orm import Session, object_session, sessionmaker
 
 from extensions.ext_redis import RedisClientWrapper
@@ -58,17 +61,23 @@ def _add_catalog_app(
     return app
 
 
-def _repository(session: Session) -> DatabaseRecommendedAppCatalogRepository:
-    redis = MagicMock(spec=RedisClientWrapper)
-    redis.get.return_value = None
+@pytest.fixture
+def catalog_redis() -> Iterator[RedisClientWrapper]:
+    with Redis() as client, patch.object(client, "execute_command", return_value=None):
+        redis = RedisClientWrapper()
+        redis.initialize(client)
+        yield redis
+
+
+def _repository(session: Session, redis: RedisClientWrapper) -> DatabaseRecommendedAppCatalogRepository:
     return DatabaseRecommendedAppCatalogRepository(
-        sessionmaker(bind=session.get_bind(), expire_on_commit=False),
-        redis=redis,
+        sessionmaker(bind=session.get_bind(), expire_on_commit=False), redis=redis
     )
 
 
 def test_list_maps_postgres_models_with_owned_session(
     db_session_with_containers: Session,
+    catalog_redis: RedisClientWrapper,
 ) -> None:
     app = _add_catalog_app(
         db_session_with_containers,
@@ -77,7 +86,7 @@ def test_list_maps_postgres_models_with_owned_session(
     private_app = _add_catalog_app(db_session_with_containers, is_public=False)
     no_site_app = _add_catalog_app(db_session_with_containers, with_site=False)
 
-    page = _repository(db_session_with_containers).list_recommended("fr-FR")
+    page = _repository(db_session_with_containers, catalog_redis).list_recommended("fr-FR")
 
     record = next(item for item in page.recommended_apps if item.app_id == app.id)
     assert record.app is not None
@@ -92,9 +101,10 @@ def test_list_maps_postgres_models_with_owned_session(
 
 def test_membership_does_not_export_dsl_with_owned_session(
     db_session_with_containers: Session,
+    catalog_redis: RedisClientWrapper,
 ) -> None:
     app = _add_catalog_app(db_session_with_containers, with_site=False)
-    repository = _repository(db_session_with_containers)
+    repository = _repository(db_session_with_containers, catalog_redis)
 
     def export_dsl(*, app_model: App, session: Session) -> str:
         assert object_session(app_model) is session

api/tests/test_containers_integration_tests/repositories/test_sqlalchemy_api_workflow_run_repository.py

@@ -86,6 +86,17 @@ def _create_workflow_run(
     return workflow_run
 
 
+def _unpersisted_pause(*, workflow_id: str, workflow_run_id: str) -> WorkflowPauseEntity:
+    return _PrivateWorkflowPauseEntity(
+        pause_model=WorkflowPause(
+            workflow_id=workflow_id,
+            workflow_run_id=workflow_run_id,
+            state_object_key="unpersisted-state",
+        ),
+        reason_models=[],
+    )
+
+
 def _cleanup_scope_data(session: Session, scope: _TestScope) -> None:
     """Remove test-created DB rows and storage objects for a test scope."""
 
@@ -736,8 +747,7 @@ class TestResumeWorkflowPause:
             test_scope,
             status=WorkflowExecutionStatus.RUNNING,
         )
-        pause_entity = Mock(spec=WorkflowPauseEntity)
-        pause_entity.id = str(uuid4())
+        pause_entity = _unpersisted_pause(workflow_id=test_scope.workflow_id, workflow_run_id=workflow_run.id)
 
         with pytest.raises(_WorkflowRunError, match="WorkflowRun is not in PAUSED status"):
             repository.resume_workflow_pause(
@@ -769,8 +779,9 @@ class TestResumeWorkflowPause:
         assert pause_model is not None
         test_scope.state_keys.add(pause_model.state_object_key)
 
-        mismatched_pause_entity = Mock(spec=WorkflowPauseEntity)
-        mismatched_pause_entity.id = str(uuid4())
+        mismatched_pause_entity = _unpersisted_pause(
+            workflow_id=test_scope.workflow_id, workflow_run_id=workflow_run.id
+        )
 
         with pytest.raises(_WorkflowRunError, match="different id in WorkflowPause and WorkflowPauseEntity"):
             repository.resume_workflow_pause(
@@ -819,8 +830,7 @@ class TestDeleteWorkflowPause:
     ) -> None:
         """Raise _WorkflowRunError when deleting a non-existent pause."""
 
-        pause_entity = Mock(spec=WorkflowPauseEntity)
-        pause_entity.id = str(uuid4())
+        pause_entity = _unpersisted_pause(workflow_id=str(uuid4()), workflow_run_id=str(uuid4()))
 
         with pytest.raises(_WorkflowRunError, match="WorkflowPause not found"):
             repository.delete_workflow_pause(pause_entity=pause_entity)

api/tests/test_containers_integration_tests/services/test_api_based_extension_service.py

@@ -1,678 +0,0 @@
-from unittest.mock import patch
-
-import pytest
-from faker import Faker
-from sqlalchemy.orm import Session
-
-from models.api_based_extension import APIBasedExtension
-from services.api_based_extension_service import APIBasedExtensionService
-from tests.test_containers_integration_tests.helpers import accounts as account_fixtures
-from tests.test_containers_integration_tests.helpers import generate_valid_password
-
-
-class TestAPIBasedExtensionService:
-    """Integration tests for APIBasedExtensionService using testcontainers."""
-
-    @pytest.fixture
-    def mock_external_service_dependencies(self):
-        """Mock setup for external service dependencies."""
-        with (
-            patch("services.account.login_adapters.SystemFeatureService") as mock_account_feature_service,
-            patch("services.api_based_extension_service.APIBasedExtensionRequestor") as mock_requestor,
-        ):
-            # Mock successful ping response
-            mock_requestor_instance = mock_requestor.return_value
-            mock_requestor_instance.request.return_value = {"result": "pong"}
-
-            yield {
-                "account_feature_service": mock_account_feature_service,
-                "requestor": mock_requestor,
-                "requestor_instance": mock_requestor_instance,
-            }
-
-    def _create_test_account_and_tenant(self, db_session_with_containers: Session, mock_external_service_dependencies):
-        """
-        Helper method to create a test account and tenant for testing.
-
-        Args:
-            db_session_with_containers: Database session from testcontainers infrastructure
-            mock_external_service_dependencies: Mock dependencies
-
-        Returns:
-            tuple: (account, tenant) - Created account and tenant instances
-        """
-        fake = Faker()
-
-        # Setup mocks for account creation
-        mock_external_service_dependencies["account_feature_service"].is_registration_allowed.return_value = True
-
-        # Create account and tenant
-        account = account_fixtures.create_account(
-            email=fake.email(),
-            name=fake.name(),
-            interface_language="en-US",
-            password=generate_valid_password(fake),
-            session=db_session_with_containers,
-        )
-        account_fixtures.create_owner_workspace(account, name=fake.company(), session=db_session_with_containers)
-        tenant = account.current_tenant
-
-        return account, tenant
-
-    def test_save_extension_success(self, db_session_with_containers: Session, mock_external_service_dependencies):
-        """
-        Test successful saving of API-based extension.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-        # Setup extension data
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        # Save extension
-        saved_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-        # Verify extension was saved correctly
-        assert saved_extension.id is not None
-        assert saved_extension.tenant_id == tenant.id
-        assert saved_extension.name == extension_data.name
-        assert saved_extension.api_endpoint == extension_data.api_endpoint
-        assert saved_extension.api_key == extension_data.api_key  # Should be decrypted when retrieved
-        assert saved_extension.created_at is not None
-
-        # Verify extension was saved to database
-
-        db_session_with_containers.refresh(saved_extension)
-        assert saved_extension.id is not None
-
-        # Verify ping connection was called
-        mock_external_service_dependencies["requestor_instance"].request.assert_called_once()
-
-    def test_save_extension_validation_errors(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test validation errors when saving extension with invalid data.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-        # Test empty name
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name="",
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        with pytest.raises(ValueError, match="name must not be empty"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-        # Test empty api_endpoint
-        extension_data.name = fake.company()
-        extension_data.api_endpoint = ""
-
-        with pytest.raises(ValueError, match="api_endpoint must not be empty"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-        # Test empty api_key
-        extension_data.api_endpoint = f"https://{fake.domain_name()}/api"
-        extension_data.api_key = ""
-
-        with pytest.raises(ValueError, match="api_key must not be empty"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-    def test_get_all_by_tenant_id_success(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test successful retrieval of all extensions by tenant ID.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-
-        # Create multiple extensions
-        extensions = []
-        assert tenant is not None
-        for i in range(3):
-            extension_data = APIBasedExtension(
-                tenant_id=tenant.id,
-                name=f"Extension {i}: {fake.company()}",
-                api_endpoint=f"https://{fake.domain_name()}/api",
-                api_key=fake.password(length=20),
-            )
-
-            saved_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-            extensions.append(saved_extension)
-
-        # Get all extensions for tenant
-        extension_list = APIBasedExtensionService.get_all_by_tenant_id(tenant.id, session=db_session_with_containers)
-
-        # Verify results
-        assert len(extension_list) == 3
-
-        # Verify all extensions belong to the correct tenant and are ordered by created_at desc
-        for i, extension in enumerate(extension_list):
-            assert extension.tenant_id == tenant.id
-            assert extension.api_key is not None  # Should be decrypted
-            if i > 0:
-                # Verify descending order (newer first)
-                assert extension.created_at <= extension_list[i - 1].created_at
-
-    def test_get_with_tenant_id_success(self, db_session_with_containers: Session, mock_external_service_dependencies):
-        """
-        Test successful retrieval of extension by tenant ID and extension ID.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-        # Create an extension
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        created_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-        # Get extension by ID
-        retrieved_extension = APIBasedExtensionService.get_with_tenant_id(
-            tenant.id, created_extension.id, session=db_session_with_containers
-        )
-
-        # Verify extension was retrieved correctly
-        assert retrieved_extension is not None
-        assert retrieved_extension.id == created_extension.id
-        assert retrieved_extension.tenant_id == tenant.id
-        assert retrieved_extension.name == extension_data.name
-        assert retrieved_extension.api_endpoint == extension_data.api_endpoint
-        assert retrieved_extension.api_key == extension_data.api_key  # Should be decrypted
-        assert retrieved_extension.created_at is not None
-
-    def test_get_with_tenant_id_not_found(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test retrieval of extension when extension is not found.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        non_existent_extension_id = fake.uuid4()
-
-        # Try to get non-existent extension
-        with pytest.raises(ValueError, match="API based extension is not found"):
-            APIBasedExtensionService.get_with_tenant_id(
-                tenant.id, non_existent_extension_id, session=db_session_with_containers
-            )
-
-    def test_delete_extension_success(self, db_session_with_containers: Session, mock_external_service_dependencies):
-        """
-        Test successful deletion of extension.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-        # Create an extension first
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        created_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-        extension_id = created_extension.id
-
-        # Delete the extension
-        APIBasedExtensionService.delete(created_extension, session=db_session_with_containers)
-
-        # Verify extension was deleted
-
-        deleted_extension = (
-            db_session_with_containers.query(APIBasedExtension).where(APIBasedExtension.id == extension_id).first()
-        )
-        assert deleted_extension is None
-
-    def test_save_extension_duplicate_name(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test validation error when saving extension with duplicate name.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-        # Create first extension
-        extension_data1 = APIBasedExtension(
-            tenant_id=tenant.id,
-            name="Test Extension",
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        APIBasedExtensionService.save(extension_data1, session=db_session_with_containers)
-        # Try to create second extension with same name
-        extension_data2 = APIBasedExtension(
-            tenant_id=tenant.id,
-            name="Test Extension",  # Same name
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        with pytest.raises(ValueError, match="name must be unique, it is already existed"):
-            APIBasedExtensionService.save(extension_data2, session=db_session_with_containers)
-
-    def test_save_extension_update_existing(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test successful update of existing extension.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-        # Create initial extension
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        created_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-        # Save original values for later comparison
-        original_name = created_extension.name
-        original_endpoint = created_extension.api_endpoint
-
-        # Update the extension with guaranteed different values
-        new_name = fake.company()
-        # Ensure new endpoint is different from original
-        new_endpoint = f"https://{fake.domain_name()}/api"
-        # If by chance they're the same, generate a new one
-        while new_endpoint == original_endpoint:
-            new_endpoint = f"https://{fake.domain_name()}/api"
-        new_api_key = fake.password(length=20)
-
-        created_extension.name = new_name
-        created_extension.api_endpoint = new_endpoint
-        created_extension.api_key = new_api_key
-
-        updated_extension = APIBasedExtensionService.save(created_extension, session=db_session_with_containers)
-
-        # Verify extension was updated correctly
-        assert updated_extension.id == created_extension.id
-        assert updated_extension.tenant_id == tenant.id
-        assert updated_extension.name == new_name
-        assert updated_extension.api_endpoint == new_endpoint
-
-        # Verify original values were changed
-        assert updated_extension.name != original_name
-        assert updated_extension.api_endpoint != original_endpoint
-
-        # Verify ping connection was called for both create and update
-        assert mock_external_service_dependencies["requestor_instance"].request.call_count == 2
-
-        # Verify the update by retrieving the extension again
-        retrieved_extension = APIBasedExtensionService.get_with_tenant_id(
-            tenant.id, created_extension.id, session=db_session_with_containers
-        )
-        assert retrieved_extension.name == new_name
-        assert retrieved_extension.api_endpoint == new_endpoint
-        assert retrieved_extension.api_key == new_api_key  # Should be decrypted when retrieved
-
-    def test_save_extension_connection_error(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test connection error when saving extension with invalid endpoint.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-
-        # Mock connection error
-        mock_external_service_dependencies["requestor_instance"].request.side_effect = ValueError(
-            "connection error: request timeout"
-        )
-        assert tenant is not None
-        # Setup extension data
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint="https://invalid-endpoint.com/api",
-            api_key=fake.password(length=20),
-        )
-
-        # Try to save extension with connection error
-        with pytest.raises(ValueError, match="connection error: request timeout"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-    def test_save_extension_invalid_api_key_length(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test validation error when saving extension with API key that is too short.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-        # Setup extension data with short API key
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key="1234",  # Less than 5 characters
-        )
-
-        # Try to save extension with short API key
-        with pytest.raises(ValueError, match="api_key must be at least 5 characters"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-    def test_save_extension_empty_fields(self, db_session_with_containers: Session, mock_external_service_dependencies):
-        """
-        Test validation errors when saving extension with empty required fields.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-        # Test with None values
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=None,  # type: ignore # why str become None here???
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        with pytest.raises(ValueError, match="name must not be empty"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-        # Test with None api_endpoint
-        extension_data.name = fake.company()
-        extension_data.api_endpoint = None
-
-        with pytest.raises(ValueError, match="api_endpoint must not be empty"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-        # Test with None api_key
-        extension_data.api_endpoint = f"https://{fake.domain_name()}/api"
-        extension_data.api_key = None
-
-        with pytest.raises(ValueError, match="api_key must not be empty"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-    def test_get_all_by_tenant_id_empty_list(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test retrieval of extensions when no extensions exist for tenant.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-
-        # Get all extensions for tenant (none exist)
-        extension_list = APIBasedExtensionService.get_all_by_tenant_id(tenant.id, session=db_session_with_containers)
-
-        # Verify empty list is returned
-        assert len(extension_list) == 0
-        assert extension_list == []
-
-    def test_save_extension_invalid_ping_response(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test validation error when ping response is invalid.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-
-        # Mock invalid ping response
-        mock_external_service_dependencies["requestor_instance"].request.return_value = {"result": "invalid"}
-        assert tenant is not None
-        # Setup extension data
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        # Try to save extension with invalid ping response
-        with pytest.raises(ValueError, match="{'result': 'invalid'}"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-    def test_save_extension_missing_ping_result(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test validation error when ping response is missing result field.
-        """
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-
-        # Mock ping response without result field
-        mock_external_service_dependencies["requestor_instance"].request.return_value = {"status": "ok"}
-        assert tenant is not None
-        # Setup extension data
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        # Try to save extension with missing ping result
-        with pytest.raises(ValueError, match="{'status': 'ok'}"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-    def test_get_with_tenant_id_wrong_tenant(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """
-        Test retrieval of extension when tenant ID doesn't match.
-        """
-        fake = Faker()
-        account1, tenant1 = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-
-        # Create second account and tenant
-        account2, tenant2 = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant1 is not None
-        # Create extension in first tenant
-        extension_data = APIBasedExtension(
-            tenant_id=tenant1.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        created_extension = APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-        # Try to get extension with wrong tenant ID
-        with pytest.raises(ValueError, match="API based extension is not found"):
-            APIBasedExtensionService.get_with_tenant_id(
-                tenant2.id, created_extension.id, session=db_session_with_containers
-            )
-
-    def test_save_extension_api_key_exactly_four_chars_rejected(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """API key with exactly 4 characters should be rejected (boundary)."""
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key="1234",
-        )
-
-        with pytest.raises(ValueError, match="api_key must be at least 5 characters"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-    def test_save_extension_api_key_exactly_five_chars_accepted(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """API key with exactly 5 characters should be accepted (boundary)."""
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key="12345",
-        )
-
-        saved = APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-        assert saved.id is not None
-
-    def test_save_extension_requestor_constructor_error(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """Exception raised by requestor constructor is wrapped in ValueError."""
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-
-        mock_external_service_dependencies["requestor"].side_effect = RuntimeError("bad config")
-
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        with pytest.raises(ValueError, match="connection error: bad config"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-    def test_save_extension_network_exception(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """Network exceptions during ping are wrapped in ValueError."""
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-
-        mock_external_service_dependencies["requestor_instance"].request.side_effect = ConnectionError(
-            "network failure"
-        )
-
-        extension_data = APIBasedExtension(
-            tenant_id=tenant.id,
-            name=fake.company(),
-            api_endpoint=f"https://{fake.domain_name()}/api",
-            api_key=fake.password(length=20),
-        )
-
-        with pytest.raises(ValueError, match="connection error: network failure"):
-            APIBasedExtensionService.save(extension_data, session=db_session_with_containers)
-
-    def test_save_extension_update_duplicate_name_rejected(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """Updating an existing extension to use another extension's name should fail."""
-        fake = Faker()
-        account, tenant = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant is not None
-
-        ext1 = APIBasedExtensionService.save(
-            APIBasedExtension(
-                tenant_id=tenant.id,
-                name="Extension Alpha",
-                api_endpoint=f"https://{fake.domain_name()}/api",
-                api_key=fake.password(length=20),
-            ),
-            session=db_session_with_containers,
-        )
-        ext2 = APIBasedExtensionService.save(
-            APIBasedExtension(
-                tenant_id=tenant.id,
-                name="Extension Beta",
-                api_endpoint=f"https://{fake.domain_name()}/api",
-                api_key=fake.password(length=20),
-            ),
-            session=db_session_with_containers,
-        )
-
-        # Try to rename ext2 to ext1's name
-        ext2.name = "Extension Alpha"
-        with pytest.raises(ValueError, match="name must be unique, it is already existed"):
-            APIBasedExtensionService.save(ext2, session=db_session_with_containers)
-
-    def test_get_all_returns_empty_for_different_tenant(
-        self, db_session_with_containers: Session, mock_external_service_dependencies
-    ):
-        """Extensions from one tenant should not be visible to another."""
-        fake = Faker()
-        _, tenant1 = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        _, tenant2 = self._create_test_account_and_tenant(
-            db_session_with_containers, mock_external_service_dependencies
-        )
-        assert tenant1 is not None
-
-        APIBasedExtensionService.save(
-            APIBasedExtension(
-                tenant_id=tenant1.id,
-                name=fake.company(),
-                api_endpoint=f"https://{fake.domain_name()}/api",
-                api_key=fake.password(length=20),
-            ),
-            session=db_session_with_containers,
-        )
-
-        assert tenant2 is not None
-        result = APIBasedExtensionService.get_all_by_tenant_id(tenant2.id, session=db_session_with_containers)
-        assert result == []

api/tests/test_containers_integration_tests/services/test_hit_testing_service.py

@@ -160,14 +160,14 @@ class TestHitTestingService:
         self, mock_format: MagicMock, db_session_with_containers: Session
     ) -> None:
         query = "test query"
-        mock_doc = MagicMock(spec=Document)
+        document = Document(page_content="Test document content")
 
         mock_record = MagicMock()
         mock_record.model_dump.return_value = {"content": "formatted content"}
         mock_format.return_value = [mock_record]
 
         response = _RetrieveResponse.model_validate(
-            HitTestingService.compact_retrieve_response(query, [mock_doc], session=db_session_with_containers)
+            HitTestingService.compact_retrieve_response(query, [document], session=db_session_with_containers)
         )
 
         assert response.query.content == query
@@ -175,7 +175,7 @@ class TestHitTestingService:
         assert response.records[0].content == "formatted content"
         mock_format.assert_called_once()
         assert mock_format.call_args.args[0] is not db_session_with_containers
-        assert mock_format.call_args.args[1] == [mock_doc]
+        assert mock_format.call_args.args[1] == [document]
 
     def test_compact_external_retrieve_response_should_return_records_for_external_provider(
         self, db_session_with_containers: Session